
import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
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 transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
		from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor







class 				lowerCAmelCase_				:



						'''simple docstring'''


						def __init__( self     :      List[str]  ,A_     :      Optional[int]  ,A_     :      Optional[int]=13  ,A_     :      str=32  ,A_     :      str=3  ,A_     :      int=4  ,A_     :      List[str]=[10, 20, 30, 40]  ,A_     :      Any=[2, 2, 3, 2]  ,A_     :      Any=True  ,A_     :      int=True  ,A_     :      str=37  ,A_     :      List[Any]="gelu"  ,A_     :      int=10  ,A_     :      str=0.02  ,A_     :      int=["stage2", "stage3", "stage4"]  ,A_     :      List[str]=[2, 3, 4]  ,A_     :      str=None  ,)     ->   Optional[Any]:
							A						   =   parent
							A						   =   batch_size
							A						   =   image_size
							A						   =   num_channels
							A						   =   num_stages
							A						   =   hidden_sizes
							A						   =   depths
							A						   =   is_training
							A						   =   use_labels
							A						   =   intermediate_size
							A						   =   hidden_act
							A						   =   num_labels
							A						   =   initializer_range
							A						   =   out_features
							A						   =   out_indices
							A						   =   scope


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   int:
							A						   =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

							A						   =   None
							if self.use_labels:
								A						   =   ids_tensor([self.batch_size]  ,self.num_labels      )

							A						   =   self.get_config()

							return config, pixel_values, labels


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]      )     ->   Optional[int]:
							return ConvNextVaConfig(
							    num_channels=self.num_channels  ,hidden_sizes=self.hidden_sizes  ,depths=self.depths  ,num_stages=self.num_stages  ,hidden_act=self.hidden_act  ,is_decoder=__lowerCAmelCase  ,initializer_range=self.initializer_range  ,out_features=self.out_features  ,out_indices=self.out_indices  ,num_labels=self.num_labels  ,)


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any  ,A_     :      List[str]  ,A_     :      Optional[Any]  ,A_     :      Dict      )     ->   Optional[int]:
							A						   =   ConvNextVaModel(config=__lowerCAmelCase      )
							model.to(__lowerCAmelCase      )
							model.eval()
							A						   =   model(__lowerCAmelCase      )
							# expected last hidden states: B, C, H // 32, W // 32
							self.parent.assertEqual(
							    result.last_hidden_state.shape  ,(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)  ,)


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int  ,A_     :      Dict  ,A_     :      Dict  ,A_     :      Union[str, Any]      )     ->   Dict:
							A						   =   ConvNextVaForImageClassification(__lowerCAmelCase      )
							model.to(__lowerCAmelCase      )
							model.eval()
							A						   =   model(__lowerCAmelCase  ,labels=__lowerCAmelCase      )
							self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.num_labels)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]  ,A_     :      Optional[Any]  ,A_     :      Any  ,A_     :      Optional[int]      )     ->   int:
							A						   =   ConvNextVaBackbone(config=__lowerCAmelCase      )
							model.to(__lowerCAmelCase      )
							model.eval()
							A						   =   model(__lowerCAmelCase      )

							# 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, self.hidden_sizes[1], 4, 4]      )

							# verify channels
							self.parent.assertEqual(len(model.channels      )  ,len(config.out_features      )      )
							self.parent.assertListEqual(model.channels  ,config.hidden_sizes[1:]      )

							# verify backbone works with out_features=None
							A						   =   None
							A						   =   ConvNextVaBackbone(config=__lowerCAmelCase      )
							model.to(__lowerCAmelCase      )
							model.eval()
							A						   =   model(__lowerCAmelCase      )

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

							# verify channels
							self.parent.assertEqual(len(model.channels      )  ,1      )
							self.parent.assertListEqual(model.channels  ,[config.hidden_sizes[-1]]      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]      )     ->   Dict:
							A						   =   self.prepare_config_and_inputs()
							A   ,      A   ,      A						   =   config_and_inputs
							A						   =   {'pixel_values': pixel_values}
							return config, inputs_dict




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Optional[int]:
							A						   =   self.prepare_config_and_inputs()
							A   ,      A   ,      A						   =   config_and_inputs
							A						   =   {'pixel_values': pixel_values, 'labels': labels}
							return config, inputs_dict







@require_torch
class 				lowerCAmelCase_				(      UpperCAmelCase__		,					UpperCAmelCase__		,					unittest.TestCase    ):



						'''simple docstring'''


						_lowerCamelCase:    List[Any]      =    (
						    (
						        ConvNextVaModel,
						        ConvNextVaForImageClassification,
						        ConvNextVaBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						_lowerCamelCase:    Optional[int]      =    (
						    {'''feature-extraction''': ConvNextVaModel, '''image-classification''': ConvNextVaForImageClassification}
						    if is_torch_available()
						    else {}
						)

						_lowerCamelCase:    Tuple      =    False
						_lowerCamelCase:    int      =    False
						_lowerCamelCase:    List[Any]      =    False
						_lowerCamelCase:    List[str]      =    False
						_lowerCamelCase:    Any      =    False


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   List[Any]:
							A						   =   ConvNextVaModelTester(self      )
							A						   =   ConfigTester(self  ,config_class=__lowerCAmelCase  ,has_text_modality=__lowerCAmelCase  ,hidden_size=37      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   Dict:
							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 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   List[Any]:
							return


						@unittest.skip(reason='ConvNextV2 does not use inputs_embeds'      )
						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   List[Any]:
							pass


						@unittest.skip(reason='ConvNextV2 does not support input and output embeddings'      )
						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   List[Any]:
							pass


						@unittest.skip(reason='ConvNextV2 does not use feedforward chunking'      )
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]      )     ->   Dict:
							pass


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   List[str]:






							if not self.model_tester.is_training:
								return

							for model_class in self.all_model_classes:
								A   ,      A						   =   self.model_tester.prepare_config_and_inputs_with_labels()
								A						   =   True

								if model_class.__name__ in [
								    *get_values(__lowerCAmelCase      ),
								    *get_values(__lowerCAmelCase      ),
								]:
									continue

								A						   =   model_class(__lowerCAmelCase      )
								model.to(__lowerCAmelCase      )
								model.train()
								A						   =   self._prepare_for_class(__lowerCAmelCase  ,__lowerCAmelCase  ,return_labels=__lowerCAmelCase      )
								A						   =   model(**__lowerCAmelCase      ).loss
								loss.backward()


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]      )     ->   List[str]:






							if not self.model_tester.is_training:
								return

							for model_class in self.all_model_classes:
								A   ,      A						   =   self.model_tester.prepare_config_and_inputs_with_labels()
								A						   =   False
								A						   =   True

								if (
								    model_class.__name__
								    in [*get_values(__lowerCAmelCase      ), *get_values(__lowerCAmelCase      )]
								    or not model_class.supports_gradient_checkpointing
								):
									continue

								A						   =   model_class(__lowerCAmelCase      )
								model.to(__lowerCAmelCase      )
								model.gradient_checkpointing_enable()
								model.train()
								A						   =   self._prepare_for_class(__lowerCAmelCase  ,__lowerCAmelCase  ,return_labels=__lowerCAmelCase      )
								A						   =   model(**__lowerCAmelCase      ).loss
								loss.backward()


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   Any:
							A   ,      A						   =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								A						   =   model_class(__lowerCAmelCase      )
								A						   =   inspect.signature(model.forward      )
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								A						   =   [*signature.parameters.keys()]

								A						   =   ['pixel_values']
								self.assertListEqual(arg_names[:1]  ,__lowerCAmelCase      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Dict:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__lowerCAmelCase      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   List[Any]:






							def check_hidden_states_output(A_     :      Union[str, Any]  ,A_     :      Dict  ,A_     :      Tuple      ):
								A						   =   model_class(__lowerCAmelCase      )
								model.to(__lowerCAmelCase      )
								model.eval()

								with torch.no_grad():
									A						   =   model(**self._prepare_for_class(__lowerCAmelCase  ,__lowerCAmelCase      )      )

								A						   =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

								A						   =   self.model_tester.num_stages
								self.assertEqual(len(__lowerCAmelCase      )  ,expected_num_stages + 1      )

								# ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
								self.assertListEqual(
								    list(hidden_states[0].shape[-2:]      )  ,[self.model_tester.image_size // 4, self.model_tester.image_size // 4]  ,)

							A   ,      A						   =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								A						   =   True
								check_hidden_states_output(__lowerCAmelCase  ,__lowerCAmelCase  ,__lowerCAmelCase      )

								# check that output_hidden_states also work using config
								del inputs_dict["output_hidden_states"]
								A						   =   True

								check_hidden_states_output(__lowerCAmelCase  ,__lowerCAmelCase  ,__lowerCAmelCase      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Optional[Any]:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*__lowerCAmelCase      )




						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   List[str]:






							for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								A						   =   ConvNextVaModel.from_pretrained(__lowerCAmelCase      )
								self.assertIsNotNone(__lowerCAmelCase      )






def       _snake_case				(		):
	A						   =   Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
	return image







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



						'''simple docstring'''


						@cached_property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]      )     ->   str:
							return AutoImageProcessor.from_pretrained('facebook/convnextv2-tiny-1k-224'      ) if is_vision_available() else None




						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   Union[str, Any]:
							A						   =   ConvNextVaForImageClassification.from_pretrained('facebook/convnextv2-tiny-1k-224'      ).to(__lowerCAmelCase      )

							A						   =   self.default_image_processor
							A						   =   prepare_img()
							A						   =   preprocessor(images=__lowerCAmelCase  ,return_tensors='pt'      ).to(__lowerCAmelCase      )

							# forward pass
							with torch.no_grad():
								A						   =   model(**__lowerCAmelCase      )

							# verify the logits
							A						   =   torch.Size((1, 1000)      )
							self.assertEqual(outputs.logits.shape  ,__lowerCAmelCase      )

							A						   =   torch.tensor([0.99_96, 0.19_66, -0.43_86]      ).to(__lowerCAmelCase      )
							self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,__lowerCAmelCase  ,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       _snake_case				(		snake_case__   :     int     ):
	A						   =   SwinvaConfig()
	A						   =   swinva_name.split('_'     )

	A						   =   name_split[1]
	if "to" in name_split[3]:
		A						   =   int(name_split[3][-3:]     )
	else:
		A						   =   int(name_split[3]     )
	if "to" in name_split[2]:
		A						   =   int(name_split[2][-2:]     )
	else:
		A						   =   int(name_split[2][6:]     )

	if model_size == "tiny":
		A						   =   96
		A						   =   (2, 2, 6, 2)
		A						   =   (3, 6, 12, 24)
	elif model_size == "small":
		A						   =   96
		A						   =   (2, 2, 18, 2)
		A						   =   (3, 6, 12, 24)
	elif model_size == "base":
		A						   =   128
		A						   =   (2, 2, 18, 2)
		A						   =   (4, 8, 16, 32)
	else:
		A						   =   192
		A						   =   (2, 2, 18, 2)
		A						   =   (6, 12, 24, 48)

	if "to" in swinva_name:
		A						   =   (12, 12, 12, 6)

	if ("22k" in swinva_name) and ("to" not in swinva_name):
		A						   =   2_1841
		A						   =   'huggingface/label-files'
		A						   =   'imagenet-22k-id2label.json'
		A						   =   json.load(open(hf_hub_download(snake_case__	,      snake_case__	,      repo_type='dataset'     )	,      'r'     )     )
		A						   =   {int(snake_case__     ): v for k, v in idalabel.items()}
		A						   =   idalabel
		A						   =   {v: k for k, v in idalabel.items()}

	else:
		A						   =   1000
		A						   =   'huggingface/label-files'
		A						   =   'imagenet-1k-id2label.json'
		A						   =   json.load(open(hf_hub_download(snake_case__	,      snake_case__	,      repo_type='dataset'     )	,      'r'     )     )
		A						   =   {int(snake_case__     ): v for k, v in idalabel.items()}
		A						   =   idalabel
		A						   =   {v: k for k, v in idalabel.items()}

	A						   =   img_size
	A						   =   num_classes
	A						   =   embed_dim
	A						   =   depths
	A						   =   num_heads
	A						   =   window_size

	return config






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

	if "head" in name:
		A						   =   name.replace('head'	,      'classifier'     )
	else:
		A						   =   'swinv2.' + name

	return name






def       _snake_case				(		snake_case__   :     List[Any]	,      snake_case__   :     List[Any]     ):
	for key in orig_state_dict.copy().keys():
		A						   =   orig_state_dict.pop(snake_case__     )

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

			if "weight" in key:
				A						   =   val[:dim, :]
				A						   =   val[dim : dim * 2, :]
				A						   =   val[-dim:, :]
			else:
				A						   =   val[:dim]
				A						   =   val[
				    dim : dim * 2
				]
				A						   =   val[-dim:]
		else:
			A						   =   val

	return orig_state_dict






def       _snake_case				(		snake_case__   :     Optional[int]	,      snake_case__   :     Tuple     ):
	A						   =   timm.create_model(snake_case__	,      pretrained=snake_case__     )
	timm_model.eval()

	A						   =   get_swinva_config(snake_case__     )
	A						   =   SwinvaForImageClassification(snake_case__     )
	model.eval()

	A						   =   convert_state_dict(timm_model.state_dict()	,      snake_case__     )
	model.load_state_dict(snake_case__     )

	A						   =   'http://images.cocodataset.org/val2017/000000039769.jpg'

	A						   =   AutoImageProcessor.from_pretrained('microsoft/{}'.format(swinva_name.replace('_'	,      '-'     )     )     )
	A						   =   Image.open(requests.get(snake_case__	,      stream=snake_case__     ).raw     )
	A						   =   image_processor(images=snake_case__	,      return_tensors='pt'     )

	A						   =   timm_model(inputs['pixel_values']     )
	A						   =   model(**snake_case__     ).logits

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

	print(F'Saving model {swinva_name} to {pytorch_dump_folder_path}'     )
	model.save_pretrained(snake_case__     )

	print(F'Saving image processor to {pytorch_dump_folder_path}'     )
	image_processor.save_pretrained(snake_case__     )

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


if __name__ == "__main__":
		_lowercase							   =				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							   =				parser.parse_args()
		convert_swinva_checkpoint(args.swinva_name, args.pytorch_dump_folder_path)