


"""simple docstring"""







import argparse
import json
import os

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase_  :  Optional[int]     		=    logging.get_logger(__name__)
def _A		(		lowercase						):







    """simple docstring"""
    a											=ConvNextVaConfig()

    if "atto" in checkpoint_url:
        a											=[2, 2, 6, 2]
        a											=[40, 80, 1_60, 3_20]
    if "femto" in checkpoint_url:
        a											=[2, 2, 6, 2]
        a											=[48, 96, 1_92, 3_84]
    if "pico" in checkpoint_url:
        a											=[2, 2, 6, 2]
        a											=[64, 1_28, 2_56, 5_12]
    if "nano" in checkpoint_url:
        a											=[2, 2, 8, 2]
        a											=[80, 1_60, 3_20, 6_40]
    if "tiny" in checkpoint_url:
        a											=[3, 3, 9, 3]
        a											=[96, 1_92, 3_84, 7_68]
    if "base" in checkpoint_url:
        a											=[3, 3, 27, 3]
        a											=[1_28, 2_56, 5_12, 10_24]
    if "large" in checkpoint_url:
        a											=[3, 3, 27, 3]
        a											=[1_92, 3_84, 7_68, 15_36]
    if "huge" in checkpoint_url:
        a											=[3, 3, 27, 3]
        a											=[3_52, 7_04, 14_08, 28_16]

    a											=10_00
    a											='''imagenet-1k-id2label.json'''
    a											=(1, 10_00)

    a											='''huggingface/label-files'''
    a											=num_labels
    a											=json.load(open(hf_hub_download(lowerCAmelCase_			,       lowerCAmelCase_			,       repo_type='''dataset'''						)			,       '''r'''						)						)
    a											={int(lowerCAmelCase_						): v for k, v in idalabel.items()}

    a											=idalabel
    a											={v: k for k, v in idalabel.items()}
    a											=hidden_sizes
    a											=depths

    return config, expected_shape
def _A		(		lowercase						):







    """simple docstring"""



    if "downsample_layers.0.0" in name:
        a											=name.replace('''downsample_layers.0.0'''			,       '''embeddings.patch_embeddings'''						)
    if "downsample_layers.0.1" in name:
        a											=name.replace('''downsample_layers.0.1'''			,       '''embeddings.norm'''						)  # we rename to layernorm later on
    if "downsample_layers.1.0" in name:
        a											=name.replace('''downsample_layers.1.0'''			,       '''stages.1.downsampling_layer.0'''						)
    if "downsample_layers.1.1" in name:
        a											=name.replace('''downsample_layers.1.1'''			,       '''stages.1.downsampling_layer.1'''						)
    if "downsample_layers.2.0" in name:
        a											=name.replace('''downsample_layers.2.0'''			,       '''stages.2.downsampling_layer.0'''						)
    if "downsample_layers.2.1" in name:
        a											=name.replace('''downsample_layers.2.1'''			,       '''stages.2.downsampling_layer.1'''						)
    if "downsample_layers.3.0" in name:
        a											=name.replace('''downsample_layers.3.0'''			,       '''stages.3.downsampling_layer.0'''						)
    if "downsample_layers.3.1" in name:
        a											=name.replace('''downsample_layers.3.1'''			,       '''stages.3.downsampling_layer.1'''						)
    if "stages" in name and "downsampling_layer" not in name:
        # stages.0.0. for instance should be renamed to stages.0.layers.0.
        a											=name[: len('''stages.0'''						)] + '''.layers''' + name[len('''stages.0'''						) :]
    if "gamma" in name:
        a											=name.replace('''gamma'''			,       '''weight'''						)
    if "beta" in name:
        a											=name.replace('''beta'''			,       '''bias'''						)
    if "stages" in name:
        a											=name.replace('''stages'''			,       '''encoder.stages'''						)
    if "norm" in name:
        a											=name.replace('''norm'''			,       '''layernorm'''						)
    if "head" in name:
        a											=name.replace('''head'''			,       '''classifier'''						)

    return name
def _A		(		):







    """simple docstring"""
    a											='''http://images.cocodataset.org/val2017/000000039769.jpg'''
    a											=Image.open(requests.get(lowerCAmelCase_			,       stream=lowerCAmelCase_						).raw						)
    return im
def _A		(		lowercase						):







    """simple docstring"""



    if "224" in checkpoint_url:
        a											=2_24
        a											=2_24 / 2_56
    elif "384" in checkpoint_url:
        a											=3_84
        a											=None
    else:
        a											=5_12
        a											=None

    return ConvNextImageProcessor(
        size=lowerCAmelCase_			,       crop_pct=lowerCAmelCase_			,       image_mean=[0.485, 0.456, 0.406]			,       image_std=[0.229, 0.224, 0.225]			,       resample=PILImageResampling.BICUBIC			,       )
@torch.no_grad()
def _A		(		lowercase			,       lowercase			,       lowercase			,       lowercase						):







    """simple docstring"""
    print('''Downloading original model from checkpoint...'''						)
    # define ConvNeXTV2 configuration based on URL
    a											=get_convnextva_config(lowerCAmelCase_						)
    # load original state_dict from URL
    a											=torch.hub.load_state_dict_from_url(lowerCAmelCase_						)['''model''']

    print('''Converting model parameters...'''						)
    # rename keys
    for key in state_dict.copy().keys():
        a											=state_dict.pop(lowerCAmelCase_						)
        a											=val
    # add prefix to all keys expect classifier head
    for key in state_dict.copy().keys():
        a											=state_dict.pop(lowerCAmelCase_						)
        if not key.startswith('''classifier'''						):
            a											='''convnextv2.''' + key
        a											=val

    # load HuggingFace model
    a											=ConvNextVaForImageClassification(lowerCAmelCase_						)
    model.load_state_dict(lowerCAmelCase_						)
    model.eval()

    # Check outputs on an image, prepared by ConvNextImageProcessor
    a											=convert_preprocessor(lowerCAmelCase_						)
    a											=preprocessor(images=prepare_img()			,       return_tensors='''pt'''						)
    a											=model(**lowerCAmelCase_						).logits

    # note: the logits below were obtained without center cropping
    if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
        a											=torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
        a											=torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
        a											=torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
        a											=torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
        a											=torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
        a											=torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
        a											=torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
        a											=torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
        a											=torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
        a											=torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
        a											=torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
        a											=torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
        a											=torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
        a											=torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
        a											=torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
        a											=torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
        a											=torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]						)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
        a											=torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]						)
    else:
        raise ValueError(f'''Unknown URL: {checkpoint_url}'''						)

    assert torch.allclose(logits[0, :5]			,       lowerCAmelCase_			,       atol=1E-3						)
    assert logits.shape == expected_shape
    print('''Model outputs match the original results!'''						)

    if save_model:
        print('''Saving model to local...'''						)
        # Create folder to save model
        if not os.path.isdir(lowerCAmelCase_						):
            os.mkdir(lowerCAmelCase_						)

        model.save_pretrained(lowerCAmelCase_						)
        preprocessor.save_pretrained(lowerCAmelCase_						)

    a											='''convnextv2'''
    if "atto" in checkpoint_url:
        model_name += "-atto"
    if "femto" in checkpoint_url:
        model_name += "-femto"
    if "pico" in checkpoint_url:
        model_name += "-pico"
    if "nano" in checkpoint_url:
        model_name += "-nano"
    elif "tiny" in checkpoint_url:
        model_name += "-tiny"
    elif "base" in checkpoint_url:
        model_name += "-base"
    elif "large" in checkpoint_url:
        model_name += "-large"
    elif "huge" in checkpoint_url:
        model_name += "-huge"
    if "22k" in checkpoint_url and "1k" not in checkpoint_url:
        model_name += "-22k"
    elif "22k" in checkpoint_url and "1k" in checkpoint_url:
        model_name += "-22k-1k"
    elif "1k" in checkpoint_url:
        model_name += "-1k"
    if "224" in checkpoint_url:
        model_name += "-224"
    elif "384" in checkpoint_url:
        model_name += "-384"
    elif "512" in checkpoint_url:
        model_name += "-512"

    if push_to_hub:
        print(f'''Pushing {model_name} to the hub...'''						)
        model.push_to_hub(lowerCAmelCase_						)
        preprocessor.push_to_hub(lowerCAmelCase_						)


if __name__ == "__main__":
       lowerCamelCase_  :  Optional[Any]     		=    argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           """--checkpoint_url""",
           default="""https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt""",
           type=str,
           help="""URL of the original ConvNeXTV2 checkpoint you\'d like to convert.""",
       )
       parser.add_argument(
           """--pytorch_dump_folder_path""",
           default="""model""",
           type=str,
           help="""Path to the output PyTorch model directory.""",
       )
       parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
       parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image preprocessor to the hub""")

       lowerCamelCase_  :  str     		=    parser.parse_args()
       convert_convnextva_checkpoint(
           args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
       )






'''simple docstring'''

import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import NatBackbone, NatForImageClassification, NatModel
   from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor





class lowerCamelCase    :







  def __init__(		self  ,  lowercase__  ,  lowercase__=13  ,  lowercase__=64  ,  lowercase__=4  ,  lowercase__=3  ,  lowercase__=16  ,  lowercase__=[1, 2, 1]  ,  lowercase__=[2, 4, 8]  ,  lowercase__=3  ,  lowercase__=2.0  ,  lowercase__=True  ,  lowercase__=0.0  ,  lowercase__=0.0  ,  lowercase__=0.1  ,  lowercase__="gelu"  ,  lowercase__=True  ,  lowercase__=0.02  ,  lowercase__=1E-5  ,  lowercase__=True  ,  lowercase__=None  ,  lowercase__=True  ,  lowercase__=10  ,  lowercase__=["stage1", "stage2"]  ,  lowercase__=[1, 2]  ,  )		->    List[str]:






        """simple docstring"""





        _snake_case		:						Optional[int]         =  parent
        _snake_case		:						List[str]         =  batch_size
        _snake_case		:						int         =  image_size
        _snake_case		:						Union[str, Any]         =  patch_size
        _snake_case		:						List[Any]         =  num_channels
        _snake_case		:						Tuple         =  embed_dim
        _snake_case		:						Tuple         =  depths
        _snake_case		:						Optional[Any]         =  num_heads
        _snake_case		:						Union[str, Any]         =  kernel_size
        _snake_case		:						Optional[Any]         =  mlp_ratio
        _snake_case		:						Optional[int]         =  qkv_bias
        _snake_case		:						List[Any]         =  hidden_dropout_prob
        _snake_case		:						Optional[int]         =  attention_probs_dropout_prob
        _snake_case		:						List[str]         =  drop_path_rate
        _snake_case		:						Optional[Any]         =  hidden_act
        _snake_case		:						Any         =  patch_norm
        _snake_case		:						Any         =  layer_norm_eps
        _snake_case		:						Union[str, Any]         =  initializer_range
        _snake_case		:						List[Any]         =  is_training
        _snake_case		:						List[Any]         =  scope
        _snake_case		:						Dict         =  use_labels
        _snake_case		:						Tuple         =  num_labels
        _snake_case		:						Union[str, Any]         =  out_features
        _snake_case		:						Dict         =  out_indices







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        _snake_case		:						List[str]         =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

        _snake_case		:						Union[str, Any]         =  None
        if self.use_labels:
              _snake_case		:						int         =  ids_tensor([self.batch_size]  ,  self.num_labels		)

        _snake_case		:						Union[str, Any]         =  self.get_config()

        return config, pixel_values, labels







  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        return NatConfig(
            num_labels=self.num_labels  ,  image_size=self.image_size  ,  patch_size=self.patch_size  ,  num_channels=self.num_channels  ,  embed_dim=self.embed_dim  ,  depths=self.depths  ,  num_heads=self.num_heads  ,  kernel_size=self.kernel_size  ,  mlp_ratio=self.mlp_ratio  ,  qkv_bias=self.qkv_bias  ,  hidden_dropout_prob=self.hidden_dropout_prob  ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,  drop_path_rate=self.drop_path_rate  ,  hidden_act=self.hidden_act  ,  patch_norm=self.patch_norm  ,  layer_norm_eps=self.layer_norm_eps  ,  initializer_range=self.initializer_range  ,  out_features=self.out_features  ,  out_indices=self.out_indices  ,  )







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    str:






        """simple docstring"""





        _snake_case		:						int         =  NatModel(config=lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						Optional[Any]         =  model(lowercase__		)

        _snake_case		:						Optional[int]         =  (config.image_size // config.patch_size) // (2 ** (len(config.depths		) - 1))
        _snake_case		:						str         =  int(config.embed_dim * 2 ** (len(config.depths		) - 1)		)

        self.parent.assertEqual(
            result.last_hidden_state.shape  ,  (self.batch_size, expected_height, expected_width, expected_dim)		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    Tuple:






        """simple docstring"""





        _snake_case		:						int         =  NatForImageClassification(lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						List[str]         =  model(lowercase__  ,  labels=lowercase__		)
        self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)		)

        # test greyscale images
        _snake_case		:						str         =  1
        _snake_case		:						Dict         =  NatForImageClassification(lowercase__		)
        model.to(lowercase__		)
        model.eval()

        _snake_case		:						int         =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
        _snake_case		:						Optional[Any]         =  model(lowercase__		)
        self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    Union[str, Any]:






        """simple docstring"""





        _snake_case		:						int         =  NatBackbone(config=lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						Any         =  model(lowercase__		)

        # verify hidden states
        self.parent.assertEqual(len(result.feature_maps		)  ,  len(config.out_features		)		)
        self.parent.assertListEqual(list(result.feature_maps[0].shape		)  ,  [self.batch_size, model.channels[0], 16, 16]		)

        # verify channels
        self.parent.assertEqual(len(model.channels		)  ,  len(config.out_features		)		)

        # verify backbone works with out_features=None
        _snake_case		:						Dict         =  None
        _snake_case		:						Any         =  NatBackbone(config=lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						str         =  model(lowercase__		)

        # verify feature maps
        self.parent.assertEqual(len(result.feature_maps		)  ,  1		)
        self.parent.assertListEqual(list(result.feature_maps[0].shape		)  ,  [self.batch_size, model.channels[-1], 4, 4]		)

        # verify channels
        self.parent.assertEqual(len(model.channels		)  ,  1		)







  def 					UpperCAmelCase_					(		self		)		->    Any:






        """simple docstring"""





        _snake_case		:						Union[str, Any]         =  self.prepare_config_and_inputs()
        _snake_case  ,				_snake_case  ,				_snake_case		:						Optional[Any]         =  config_and_inputs
        _snake_case		:						List[Any]         =  {'''pixel_values''': pixel_values}
        return config, inputs_dict





@require_natten
@require_torch
class lowerCamelCase    (a__    ,   a__    ,   unittest.TestCase       ):
  _lowercase	:						int							       =					(
      (
          NatModel,
          NatForImageClassification,
          NatBackbone,
      )
      if is_torch_available()
      else ()
  )
  _lowercase	:						Optional[Any]							       =					(
      {"""feature-extraction""": NatModel, """image-classification""": NatForImageClassification}
      if is_torch_available()
      else {}
  )
  _lowercase	:						List[str]							       =					False

  _lowercase	:						str							       =					False
  _lowercase	:						Dict							       =					False
  _lowercase	:						Tuple							       =					False
  _lowercase	:						List[str]							       =					False







  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        _snake_case		:						List[str]         =  NatModelTester(self		)
        _snake_case		:						List[str]         =  ConfigTester(self  ,  config_class=lowercase__  ,  embed_dim=37		)







  def 					UpperCAmelCase_					(		self		)		->    Union[str, Any]:






        """simple docstring"""





        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        return







  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        _snake_case		:						Tuple         =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						Optional[int]         =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    Any:






        """simple docstring"""





        _snake_case		:						List[Any]         =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*lowercase__		)







  @unittest.skip(reason='''Nat does not use inputs_embeds'''		)
  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        pass







  @unittest.skip(reason='''Nat does not use feedforward chunking'''		)
  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        pass







  def 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Tuple         =  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _snake_case		:						List[Any]         =  model_class(lowercase__		)
              self.assertIsInstance(model.get_input_embeddings()  ,  (nn.Module)		)
              _snake_case		:						str         =  model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(lowercase__  ,  nn.Linear		)		)







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Dict         =  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _snake_case		:						Union[str, Any]         =  model_class(lowercase__		)
              _snake_case		:						Tuple         =  inspect.signature(model.forward		)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              _snake_case		:						List[Any]         =  [*signature.parameters.keys()]

              _snake_case		:						str         =  ['''pixel_values''']
              self.assertListEqual(arg_names[:1]  ,  lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.'''		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    Any:






        """simple docstring"""





        _snake_case		:						str         =  model_class(lowercase__		)
        model.to(lowercase__		)
        model.eval()

        with torch.no_grad():
              _snake_case		:						Optional[Any]         =  model(**self._prepare_for_class(lowercase__  ,  lowercase__		)		)

        _snake_case		:						Any         =  outputs.hidden_states

        _snake_case		:						Optional[Any]         =  getattr(
            self.model_tester  ,  '''expected_num_hidden_layers'''  ,  len(self.model_tester.depths		) + 1		)
        self.assertEqual(len(lowercase__		)  ,  lowercase__		)

        # Nat has a different seq_length
        _snake_case		:						Tuple         =  (
            config.patch_size
            if isinstance(config.patch_size  ,  collections.abc.Iterable		)
            else (config.patch_size, config.patch_size)
        )

        _snake_case		:						int         =  image_size[0] // patch_size[0]
        _snake_case		:						Optional[Any]         =  image_size[1] // patch_size[1]

        self.assertListEqual(
            list(hidden_states[0].shape[-3:]		)  ,  [height, width, self.model_tester.embed_dim]  ,  )

        if model_class.__name__ != "NatBackbone":
              _snake_case		:						Dict         =  outputs.reshaped_hidden_states
              self.assertEqual(len(lowercase__		)  ,  lowercase__		)

              _snake_case  ,				_snake_case  ,				_snake_case  ,				_snake_case		:						Any         =  reshaped_hidden_states[0].shape
              _snake_case		:						str         =  (
                  reshaped_hidden_states[0].view(lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		).permute(0  ,  2  ,  3  ,  1		)
              )
              self.assertListEqual(
                  list(reshaped_hidden_states.shape[-3:]		)  ,  [height, width, self.model_tester.embed_dim]  ,  )







  def 					UpperCAmelCase_					(		self		)		->    Union[str, Any]:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						str         =  self.model_tester.prepare_config_and_inputs_for_common()

        _snake_case		:						Dict         =  (
            self.model_tester.image_size
            if isinstance(self.model_tester.image_size  ,  collections.abc.Iterable		)
            else (self.model_tester.image_size, self.model_tester.image_size)
        )

        for model_class in self.all_model_classes:
              _snake_case		:						Optional[Any]         =  True
              self.check_hidden_states_output(lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              _snake_case		:						str         =  True

              self.check_hidden_states_output(lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		)







  @slow
  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              _snake_case		:						Optional[int]         =  NatModel.from_pretrained(lowercase__		)
              self.assertIsNotNone(lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Tuple         =  self.model_tester.prepare_config_and_inputs_for_common()

        _snake_case		:						str         =  _config_zero_init(lowercase__		)
        for model_class in self.all_model_classes:
              _snake_case		:						int         =  model_class(config=lowercase__		)
              for name, param in model.named_parameters():
                    if "embeddings" not in name and param.requires_grad:
                          self.assertIn(
                              ((param.data.mean() * 1E9).round() / 1E9).item()  ,  [0.0, 1.0]  ,  msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''  ,  )





@require_natten
@require_vision
@require_torch
class lowerCamelCase    (unittest.TestCase       ):







  @cached_property
  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224'''		) if is_vision_available() else None







  @slow
  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						Tuple         =  NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''		).to(lowercase__		)
        _snake_case		:						Optional[int]         =  self.default_image_processor

        _snake_case		:						List[str]         =  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
        _snake_case		:						List[Any]         =  image_processor(images=lowercase__  ,  return_tensors='''pt'''		).to(lowercase__		)

        # forward pass
        with torch.no_grad():
              _snake_case		:						Tuple         =  model(**lowercase__		)

        # verify the logits
        _snake_case		:						List[str]         =  torch.Size((1, 1_000)		)
        self.assertEqual(outputs.logits.shape  ,  lowercase__		)
        _snake_case		:						List[str]         =  torch.tensor([0.3_805, -0.8_676, -0.3_912]		).to(lowercase__		)
        self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,  lowercase__  ,  atol=1E-4		)		)





@require_torch
@require_natten
class lowerCamelCase    (unittest.TestCase    ,   a__       ):
  _lowercase	:						str							       =					(NatBackbone,) if is_torch_available() else ()
  _lowercase	:						List[Any]							       =					NatConfig







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						Any         =  NatModelTester(self		)


