





"""simple docstring"""





import copy
import inspect
import unittest

from transformers import PretrainedConfig, SwiftFormerConfig
from transformers.testing_utils import (
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
  import torch
  from torch import nn

  from transformers import SwiftFormerForImageClassification, SwiftFormerModel
  from transformers.models.swiftformer.modeling_swiftformer import SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor




class 							snake_case							:

      '''simple docstring'''



      def __init__(							self       :      str						,				__lowercase       :      List[Any]						,				__lowercase       :      Optional[Any]=13						,				__lowercase       :      Optional[int]=3						,				__lowercase       :      Tuple=True						,				__lowercase       :      List[Any]=True						,				__lowercase       :      Any=0.1						,				__lowercase       :      int=0.1						,				__lowercase       :      Optional[int]=224						,				__lowercase       :      Optional[Any]=1_000						,				__lowercase       :      List[Any]=[3, 3, 6, 4]						,				__lowercase       :      int=[48, 56, 112, 220]						,				):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						parent
        __UpperCAmelCase    :						Any       =						batch_size
        __UpperCAmelCase    :						Optional[int]       =						num_channels
        __UpperCAmelCase    :						Tuple       =						is_training
        __UpperCAmelCase    :						int       =						use_labels
        __UpperCAmelCase    :						Optional[Any]       =						hidden_dropout_prob
        __UpperCAmelCase    :						List[str]       =						attention_probs_dropout_prob
        __UpperCAmelCase    :						Union[str, Any]       =						num_labels
        __UpperCAmelCase    :						Optional[Any]       =						image_size
        __UpperCAmelCase    :						Optional[int]       =						layer_depths
        __UpperCAmelCase    :						str       =						embed_dims

      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Tuple       =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

        __UpperCAmelCase    :						int       =						None
        if self.use_labels:
          __UpperCAmelCase    :						List[str]       =						ids_tensor([self.batch_size]						,				self.num_labels )

        __UpperCAmelCase    :						List[Any]       =						self.get_config()

        return config, pixel_values, labels

      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        return SwiftFormerConfig(
            depths=self.layer_depths						,				embed_dims=self.embed_dims						,				mlp_ratio=4						,				downsamples=[True, True, True, True]						,				hidden_act='''gelu'''						,				num_labels=self.num_labels						,				down_patch_size=3						,				down_stride=2						,				down_pad=1						,				drop_rate=0.0						,				drop_path_rate=0.0						,				use_layer_scale=_A						,				layer_scale_init_value=1e-5						,				)

      def       A_					(							self       :      Tuple						,				__lowercase       :      int						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						Union[str, Any]       =						SwiftFormerModel(config=_A )
        model.to(_A )
        model.eval()
        __UpperCAmelCase    :						str       =						model(_A )
        self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.embed_dims[-1], 7, 7) )

      def       A_					(							self       :      List[Any]						,				__lowercase       :      List[Any]						,				__lowercase       :      Tuple						,				__lowercase       :      str ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						self.num_labels
        __UpperCAmelCase    :						List[Any]       =						SwiftFormerForImageClassification(_A )
        model.to(_A )
        model.eval()
        __UpperCAmelCase    :						Tuple       =						model(_A						,				labels=_A )
        self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.num_labels) )

        __UpperCAmelCase    :						Optional[int]       =						SwiftFormerForImageClassification(_A )
        model.to(_A )
        model.eval()

        __UpperCAmelCase    :						Tuple       =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )
        __UpperCAmelCase    :						Tuple       =						model(_A )
        self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.num_labels) )







      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        (__UpperCAmelCase)    :						Any       =						self.prepare_config_and_inputs()
        __UpperCAmelCase    :						int       =						{'pixel_values': pixel_values}
        return config, inputs_dict




@require_torch
class 							snake_case							(			a__      ,  a__      ,  unittest.TestCase ):

      '''simple docstring'''



      _A     :       Optional[Any]         =       (SwiftFormerModel, SwiftFormerForImageClassification) if is_torch_available() else ()
      _A     :       Any         =       (
          {'feature-extraction': SwiftFormerModel, 'image-classification': SwiftFormerForImageClassification}
          if is_torch_available()
          else {}
      )

      _A     :       Optional[Any]         =       False
      _A     :       Any         =       False
      _A     :       int         =       False
      _A     :       List[Any]         =       False
      _A     :       str         =       False

      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Union[str, Any]       =						SwiftFormerModelTester(self )
        __UpperCAmelCase    :						Tuple       =						ConfigTester(
            self						,				config_class=_A						,				has_text_modality=_A						,				hidden_size=37						,				num_attention_heads=12						,				num_hidden_layers=12						,				)

      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        self.config_tester.run_common_tests()

      @unittest.skip(reason='''SwiftFormer does not use inputs_embeds''' )
      def       A_					(							self       :      Any ):


        '''simple docstring'''




        pass

      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						int       =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						List[str]       =						model_class(_A )
          __UpperCAmelCase    :						Union[str, Any]       =						model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(_A						,				nn.Linear ) )

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						str       =						self.model_tester.prepare_config_and_inputs_for_common()

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

          __UpperCAmelCase    :						List[Any]       =						['pixel_values']
          self.assertListEqual(arg_names[:1]						,				_A )

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_A )

      def       A_					(							self       :      Union[str, Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Union[str, Any]       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*_A )

      @slow
      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        for model_name in SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          __UpperCAmelCase    :						Union[str, Any]       =						SwiftFormerModel.from_pretrained(_A )
          self.assertIsNotNone(_A )

      @unittest.skip(reason='''SwiftFormer does not output attentions''' )
      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




        pass

      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        def check_hidden_states_output(__lowercase       :      Union[str, Any]						,				__lowercase       :      Tuple						,				__lowercase       :      List[Any] ):
          __UpperCAmelCase    :						Tuple       =						model_class(_A )
          model.to(_A )
          model.eval()

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

          __UpperCAmelCase    :						Dict       =						outputs.hidden_states

          __UpperCAmelCase    :						List[str]       =						8
          self.assertEqual(len(_A )						,				_A )  # TODO

          # SwiftFormer's feature maps are of shape (batch_size, embed_dims, height, width)
          # with the width and height being successively divided by 2, after every 2 blocks
          for i in range(len(_A ) ):
            self.assertEqual(
                hidden_states[i].shape						,				torch.Size(
                    [
                        self.model_tester.batch_size,
                        self.model_tester.embed_dims[i // 2],
                        (self.model_tester.image_size // 4) // 2 ** (i // 2),
                        (self.model_tester.image_size // 4) // 2 ** (i // 2),
                    ] )						,				)

        __UpperCAmelCase    :						List[Any]       =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						Dict       =						True
          check_hidden_states_output(_A						,				_A						,				_A )

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

          check_hidden_states_output(_A						,				_A						,				_A )

      def       A_					(							self       :      str ):


        '''simple docstring'''




        def _config_zero_init(__lowercase       :      List[str] ):
          __UpperCAmelCase    :						Optional[Any]       =						copy.deepcopy(_A )
          for key in configs_no_init.__dict__.keys():
            if "_range" in key or "_std" in key or "initializer_factor" in key or "layer_scale" in key:
              setattr(_A						,				_A						,				1e-10 )
            if isinstance(getattr(_A						,				_A						,				_A )						,				_A ):
              __UpperCAmelCase    :						Tuple       =						_config_zero_init(getattr(_A						,				_A ) )
              setattr(_A						,				_A						,				_A )
          return configs_no_init

        __UpperCAmelCase    :						Optional[Any]       =						self.model_tester.prepare_config_and_inputs_for_common()

        __UpperCAmelCase    :						int       =						_config_zero_init(_A )
        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						List[str]       =						model_class(config=_A )
          for name, param in model.named_parameters():
            if param.requires_grad:
              self.assertIn(
                  ((param.data.mean() * 1e9) / 1e9).round().item()						,				[0.0, 1.0]						,				msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''						,				)







      @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''' )
      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        pass


def       lowerCamelCase_       (   )						->List[str]:





  """simple docstring"""


  __UpperCAmelCase    :						List[Any]       =						Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''  )
  return image







@require_torch
@require_vision
class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      @cached_property
      def       A_					(							self       :      Union[str, Any] ):


        '''simple docstring'''




        return ViTImageProcessor.from_pretrained('''MBZUAI/swiftformer-xs''' ) if is_vision_available() else None







      @slow
      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						SwiftFormerForImageClassification.from_pretrained('''MBZUAI/swiftformer-xs''' ).to(_A )

        __UpperCAmelCase    :						List[Any]       =						self.default_image_processor
        __UpperCAmelCase    :						str       =						prepare_img()
        __UpperCAmelCase    :						int       =						image_processor(images=_A						,				return_tensors='''pt''' ).to(_A )

        # forward pass
        with torch.no_grad():
          __UpperCAmelCase    :						str       =						model(**_A )

        # verify the logits
        __UpperCAmelCase    :						Dict       =						torch.Size((1, 1_000) )
        self.assertEqual(outputs.logits.shape						,				_A )

        __UpperCAmelCase    :						Any       =						torch.tensor([[-2.1703e00, 2.1107e00, -2.0811e00]] ).to(_A )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]						,				_A						,				atol=1e-4 ) )






"""simple docstring"""





import argparse
import json
from pathlib import Path

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

from transformers import AutoImageProcessor, SwinvaConfig, SwinvaForImageClassification


def       lowerCamelCase_       (   UpperCAmelCase_  )						->Optional[Any]:





  """simple docstring"""


  __UpperCAmelCase    :						Tuple       =						SwinvaConfig()
  __UpperCAmelCase    :						Any       =						swinva_name.split('''_'''  )

  __UpperCAmelCase    :						Optional[Any]       =						name_split[1]
  if "to" in name_split[3]:
    __UpperCAmelCase    :						List[str]       =						int(name_split[3][-3:]  )
  else:
    __UpperCAmelCase    :						Union[str, Any]       =						int(name_split[3]  )
  if "to" in name_split[2]:
    __UpperCAmelCase    :						List[str]       =						int(name_split[2][-2:]  )
  else:
    __UpperCAmelCase    :						str       =						int(name_split[2][6:]  )

  if model_size == "tiny":
    __UpperCAmelCase    :						Union[str, Any]       =						96
    __UpperCAmelCase    :						str       =						(2, 2, 6, 2)
    __UpperCAmelCase    :						Union[str, Any]       =						(3, 6, 12, 24)
  elif model_size == "small":
    __UpperCAmelCase    :						List[Any]       =						96
    __UpperCAmelCase    :						Union[str, Any]       =						(2, 2, 18, 2)
    __UpperCAmelCase    :						str       =						(3, 6, 12, 24)
  elif model_size == "base":
    __UpperCAmelCase    :						Any       =						1_28
    __UpperCAmelCase    :						int       =						(2, 2, 18, 2)
    __UpperCAmelCase    :						Dict       =						(4, 8, 16, 32)
  else:
    __UpperCAmelCase    :						Union[str, Any]       =						1_92
    __UpperCAmelCase    :						Union[str, Any]       =						(2, 2, 18, 2)
    __UpperCAmelCase    :						str       =						(6, 12, 24, 48)

  if "to" in swinva_name:
    __UpperCAmelCase    :						Optional[int]       =						(12, 12, 12, 6)

  if ("22k" in swinva_name) and ("to" not in swinva_name):
    __UpperCAmelCase    :						int       =						2_18_41
    __UpperCAmelCase    :						Optional[Any]       =						'''huggingface/label-files'''
    __UpperCAmelCase    :						Optional[Any]       =						'''imagenet-22k-id2label.json'''
    __UpperCAmelCase    :						Any       =						json.load(open(hf_hub_download(UpperCAmelCase_  ,			UpperCAmelCase_  ,			repo_type='''dataset'''  )  ,			'''r'''  )  )
    __UpperCAmelCase    :						List[Any]       =						{int(UpperCAmelCase_  ): v for k, v in idalabel.items()}
    __UpperCAmelCase    :						Optional[Any]       =						idalabel
    __UpperCAmelCase    :						Optional[int]       =						{v: k for k, v in idalabel.items()}

  else:
    __UpperCAmelCase    :						List[Any]       =						10_00
    __UpperCAmelCase    :						str       =						'''huggingface/label-files'''
    __UpperCAmelCase    :						Tuple       =						'''imagenet-1k-id2label.json'''
    __UpperCAmelCase    :						Optional[int]       =						json.load(open(hf_hub_download(UpperCAmelCase_  ,			UpperCAmelCase_  ,			repo_type='''dataset'''  )  ,			'''r'''  )  )
    __UpperCAmelCase    :						str       =						{int(UpperCAmelCase_  ): v for k, v in idalabel.items()}
    __UpperCAmelCase    :						Any       =						idalabel
    __UpperCAmelCase    :						List[Any]       =						{v: k for k, v in idalabel.items()}

  __UpperCAmelCase    :						Optional[int]       =						img_size
  __UpperCAmelCase    :						int       =						num_classes
  __UpperCAmelCase    :						Tuple       =						embed_dim
  __UpperCAmelCase    :						str       =						depths
  __UpperCAmelCase    :						Optional[Any]       =						num_heads
  __UpperCAmelCase    :						Optional[int]       =						window_size

  return config


def       lowerCamelCase_       (   UpperCAmelCase_  )						->str:





  """simple docstring"""


  if "patch_embed.proj" in name:
    __UpperCAmelCase    :						Tuple       =						name.replace('''patch_embed.proj'''  ,			'''embeddings.patch_embeddings.projection'''  )
  if "patch_embed.norm" in name:
    __UpperCAmelCase    :						List[str]       =						name.replace('''patch_embed.norm'''  ,			'''embeddings.norm'''  )
  if "layers" in name:
    __UpperCAmelCase    :						Optional[int]       =						'''encoder.''' + name
  if "attn.proj" in name:
    __UpperCAmelCase    :						Optional[Any]       =						name.replace('''attn.proj'''  ,			'''attention.output.dense'''  )
  if "attn" in name:
    __UpperCAmelCase    :						List[Any]       =						name.replace('''attn'''  ,			'''attention.self'''  )
  if "norm1" in name:
    __UpperCAmelCase    :						List[str]       =						name.replace('''norm1'''  ,			'''layernorm_before'''  )
  if "norm2" in name:
    __UpperCAmelCase    :						Dict       =						name.replace('''norm2'''  ,			'''layernorm_after'''  )
  if "mlp.fc1" in name:
    __UpperCAmelCase    :						str       =						name.replace('''mlp.fc1'''  ,			'''intermediate.dense'''  )
  if "mlp.fc2" in name:
    __UpperCAmelCase    :						int       =						name.replace('''mlp.fc2'''  ,			'''output.dense'''  )
  if "q_bias" in name:
    __UpperCAmelCase    :						Tuple       =						name.replace('''q_bias'''  ,			'''query.bias'''  )
  if "k_bias" in name:
    __UpperCAmelCase    :						List[str]       =						name.replace('''k_bias'''  ,			'''key.bias'''  )
  if "v_bias" in name:
    __UpperCAmelCase    :						str       =						name.replace('''v_bias'''  ,			'''value.bias'''  )
  if "cpb_mlp" in name:
    __UpperCAmelCase    :						Tuple       =						name.replace('''cpb_mlp'''  ,			'''continuous_position_bias_mlp'''  )
  if name == "norm.weight":
    __UpperCAmelCase    :						int       =						'''layernorm.weight'''
  if name == "norm.bias":
    __UpperCAmelCase    :						Optional[Any]       =						'''layernorm.bias'''

  if "head" in name:
    __UpperCAmelCase    :						List[Any]       =						name.replace('''head'''  ,			'''classifier'''  )
  else:
    __UpperCAmelCase    :						List[Any]       =						'''swinv2.''' + name

  return name


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->int:





  """simple docstring"""


  for key in orig_state_dict.copy().keys():
    __UpperCAmelCase    :						Any       =						orig_state_dict.pop(UpperCAmelCase_  )

    if "mask" in key:
      continue
    elif "qkv" in key:
      __UpperCAmelCase    :						int       =						key.split('''.'''  )
      __UpperCAmelCase    :						int       =						int(key_split[1]  )
      __UpperCAmelCase    :						str       =						int(key_split[3]  )
      __UpperCAmelCase    :						List[str]       =						model.swinva.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

      if "weight" in key:
        __UpperCAmelCase    :						Optional[Any]       =						val[:dim, :]
        __UpperCAmelCase    :						Union[str, Any]       =						val[dim : dim * 2, :]
        __UpperCAmelCase    :						List[Any]       =						val[-dim:, :]
      else:
        __UpperCAmelCase    :						Optional[int]       =						val[:dim]
        __UpperCAmelCase    :						Any       =						val[
            dim : dim * 2
        ]
        __UpperCAmelCase    :						List[str]       =						val[-dim:]
    else:
      __UpperCAmelCase    :						Tuple       =						val

  return orig_state_dict


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->List[str]:





  """simple docstring"""


  __UpperCAmelCase    :						str       =						timm.create_model(UpperCAmelCase_  ,			pretrained=UpperCAmelCase_  )
  timm_model.eval()

  __UpperCAmelCase    :						Any       =						get_swinva_config(UpperCAmelCase_  )
  __UpperCAmelCase    :						List[str]       =						SwinvaForImageClassification(UpperCAmelCase_  )
  model.eval()

  __UpperCAmelCase    :						int       =						convert_state_dict(timm_model.state_dict()  ,			UpperCAmelCase_  )
  model.load_state_dict(UpperCAmelCase_  )

  __UpperCAmelCase    :						Optional[Any]       =						'''http://images.cocodataset.org/val2017/000000039769.jpg'''

  __UpperCAmelCase    :						Optional[Any]       =						AutoImageProcessor.from_pretrained('''microsoft/{}'''.format(swinva_name.replace('''_'''  ,			'''-'''  )  )  )
  __UpperCAmelCase    :						List[Any]       =						Image.open(requests.get(UpperCAmelCase_  ,			stream=UpperCAmelCase_  ).raw  )
  __UpperCAmelCase    :						Union[str, Any]       =						image_processor(images=UpperCAmelCase_  ,			return_tensors='''pt'''  )

  __UpperCAmelCase    :						List[Any]       =						timm_model(inputs['''pixel_values''']  )
  __UpperCAmelCase    :						Tuple       =						model(**UpperCAmelCase_  ).logits

  assert torch.allclose(UpperCAmelCase_  ,			UpperCAmelCase_  ,			atol=1e-3  )

  print(f'''Saving model {swinva_name} to {pytorch_dump_folder_path}'''  )
  model.save_pretrained(UpperCAmelCase_  )

  print(f'''Saving image processor to {pytorch_dump_folder_path}'''  )
  image_processor.save_pretrained(UpperCAmelCase_  )

  model.push_to_hub(
      repo_path_or_name=Path(UpperCAmelCase_  ,			UpperCAmelCase_  )  ,			organization='''nandwalritik'''  ,			commit_message='''Add model'''  ,			)


if __name__ == "__main__":
  lowercase__				:List[Any]				   =       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--swinv2_name',
      default='swinv2_tiny_patch4_window8_256',
      type=str,
      help='Name of the Swinv2 timm model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
  )

  lowercase__				:List[Any]				   =       parser.parse_args()
  convert_swinva_checkpoint(args.swinva_name, args.pytorch_dump_folder_path)