




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       a__   :






						def __init__(       self		,      _A		,      _A=1_3		,      _A=3_2		,      _A=3		,      _A=4		,      _A=[1_0, 2_0, 3_0, 4_0]		,      _A=[2, 2, 3, 2]		,      _A=True		,      _A=True		,      _A=3_7		,      _A="gelu"		,      _A=1_0		,      _A=0.02		,      _A=["stage2", "stage3", "stage4"]		,      _A=[2, 3, 4]		,      _A=None		,      ):




											"""simple docstring"""



											__lowerCAmelCase           =	parent
											__lowerCAmelCase           =	batch_size
											__lowerCAmelCase           =	image_size
											__lowerCAmelCase           =	num_channels
											__lowerCAmelCase           =	num_stages
											__lowerCAmelCase           =	hidden_sizes
											__lowerCAmelCase           =	depths
											__lowerCAmelCase           =	is_training
											__lowerCAmelCase           =	use_labels
											__lowerCAmelCase           =	intermediate_size
											__lowerCAmelCase           =	hidden_act
											__lowerCAmelCase           =	num_labels
											__lowerCAmelCase           =	initializer_range
											__lowerCAmelCase           =	out_features
											__lowerCAmelCase           =	out_indices
											__lowerCAmelCase           =	scope






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

											__lowerCAmelCase           =	self.get_config()

											return config, pixel_values, labels






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											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=_A		,      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		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	ConvNextVaModel(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )
											# 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 // 3_2, self.image_size // 3_2)		,      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	ConvNextVaForImageClassification(_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A		,      labels=_A      )
											self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	ConvNextVaBackbone(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )

											# 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
											__lowerCAmelCase           =	None
											__lowerCAmelCase           =	ConvNextVaBackbone(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )

											# 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      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.prepare_config_and_inputs()
											__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	config_and_inputs
											__lowerCAmelCase           =	{"pixel_values": pixel_values}
											return config, inputs_dict






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.prepare_config_and_inputs()
											__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	config_and_inputs
											__lowerCAmelCase           =	{"pixel_values": pixel_values, "labels": labels}
											return config, inputs_dict






@require_torch
class       a__   ( snake_case__			,    snake_case__			,    unittest.TestCase				):
						_a							:   str =   (
						    (
						        ConvNextVaModel,
						        ConvNextVaForImageClassification,
						        ConvNextVaBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						_a							:   Optional[int] =   (
						    {"""feature-extraction""": ConvNextVaModel, """image-classification""": ConvNextVaForImageClassification}
						    if is_torch_available()
						    else {}
						)

						_a							:   Optional[int] =   False
						_a							:   Dict =   False
						_a							:   Tuple =   False
						_a							:   List[str] =   False
						_a							:   Union[str, Any] =   False






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	ConvNextVaModelTester(self      )
											__lowerCAmelCase           =	ConfigTester(self		,      config_class=_A		,      has_text_modality=_A		,      hidden_size=3_7      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""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     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return






						@unittest.skip(reason="ConvNextV2 does not use inputs_embeds"      )
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											pass






						@unittest.skip(reason="ConvNextV2 does not support input and output embeddings"      )
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											pass






						@unittest.skip(reason="ConvNextV2 does not use feedforward chunking"      )
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											pass






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											if not self.model_tester.is_training:
																return

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

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

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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											if not self.model_tester.is_training:
																return

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

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

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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()

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

																__lowerCAmelCase           =	["pixel_values"]
																self.assertListEqual(arg_names[:1]		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											def check_hidden_states_output(_A		,      _A		,      _A      ):
																__lowerCAmelCase           =	model_class(_A      )
																model.to(_A      )
																model.eval()

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

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

																__lowerCAmelCase           =	self.model_tester.num_stages
																self.assertEqual(len(_A      )		,      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]		,      )

											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()

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

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

																check_hidden_states_output(_A		,      _A		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_A      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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







def      _a				(						):
					__lowerCAmelCase           =	Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
					return image


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






						@cached_property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224"      ) if is_vision_available() else None






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	ConvNextVaForImageClassification.from_pretrained("facebook/convnextv2-tiny-1k-224"      ).to(_A      )

											__lowerCAmelCase           =	self.default_image_processor
											__lowerCAmelCase           =	prepare_img()
											__lowerCAmelCase           =	preprocessor(images=_A		,      return_tensors="pt"      ).to(_A      )

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

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

											__lowerCAmelCase           =	torch.tensor([0.99_96, 0.19_66, -0.43_86]      ).to(_A      )
											self.assertTrue(torch.allclose(outputs.logits[0, :3]		,      _A		,      atol=1E-4      )      )










import argparse
import os

import torch

from transformers import FlavaConfig, FlavaForPreTraining
from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint







def      _a				(						SCREAMING_SNAKE_CASE_							:    Any					):
					# encoder.embeddings are double copied in original FLAVA
					return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items()					)







def      _a				(						SCREAMING_SNAKE_CASE_							:    Dict				,				SCREAMING_SNAKE_CASE_							:    Dict					):
					__lowerCAmelCase           =	{}

					for key, value in state_dict.items():
										if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
															continue

										__lowerCAmelCase           =	key.replace("heads.cmd.mim_head.cls.predictions"				,				"mmm_image_head"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.mlm_head.cls.predictions"				,				"mmm_text_head"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.itm_head.cls"				,				"itm_head"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.itm_head.pooler"				,				"itm_head.pooler"					)
										__lowerCAmelCase           =	key.replace("heads.cmd.clip_head.logit_scale"				,				"flava.logit_scale"					)
										__lowerCAmelCase           =	key.replace("heads.fairseq_mlm.cls.predictions"				,				"mlm_head"					)
										__lowerCAmelCase           =	key.replace("heads.imagenet.mim_head.cls.predictions"				,				"mim_head"					)
										__lowerCAmelCase           =	key.replace("mm_text_projection"				,				"flava.text_to_mm_projection"					)
										__lowerCAmelCase           =	key.replace("mm_image_projection"				,				"flava.image_to_mm_projection"					)
										__lowerCAmelCase           =	key.replace("image_encoder.module"				,				"flava.image_model"					)
										__lowerCAmelCase           =	key.replace("text_encoder.module"				,				"flava.text_model"					)
										__lowerCAmelCase           =	key.replace("mm_encoder.module.encoder.cls_token"				,				"flava.multimodal_model.cls_token"					)
										__lowerCAmelCase           =	key.replace("mm_encoder.module"				,				"flava.multimodal_model"					)
										__lowerCAmelCase           =	key.replace("text_projection"				,				"flava.text_projection"					)
										__lowerCAmelCase           =	key.replace("image_projection"				,				"flava.image_projection"					)

										__lowerCAmelCase           =	value.float()

					for key, value in codebook_state_dict.items():
										__lowerCAmelCase           =	value

					return upgrade







@torch.no_grad()
def      _a				(						SCREAMING_SNAKE_CASE_							:    List[Any]				,				SCREAMING_SNAKE_CASE_							:    Optional[int]				,				SCREAMING_SNAKE_CASE_							:    List[str]				,				SCREAMING_SNAKE_CASE_							:    int=None					):
					if config_path is not None:
										__lowerCAmelCase           =	FlavaConfig.from_pretrained(SCREAMING_SNAKE_CASE_					)
					else:
										__lowerCAmelCase           =	FlavaConfig()

					__lowerCAmelCase           =	FlavaForPreTraining(SCREAMING_SNAKE_CASE_					).eval()

					__lowerCAmelCase           =	convert_dalle_checkpoint(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				save_checkpoint=SCREAMING_SNAKE_CASE_					)

					if os.path.exists(SCREAMING_SNAKE_CASE_					):
										__lowerCAmelCase           =	torch.load(SCREAMING_SNAKE_CASE_				,				map_location="cpu"					)
					else:
										__lowerCAmelCase           =	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_				,				map_location="cpu"					)

					__lowerCAmelCase           =	upgrade_state_dict(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_					)
					hf_model.load_state_dict(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	hf_model.state_dict()
					__lowerCAmelCase           =	count_parameters(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	count_parameters(SCREAMING_SNAKE_CASE_					) + count_parameters(SCREAMING_SNAKE_CASE_					)

					assert torch.allclose(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				atol=1E-3					)

					hf_model.save_pretrained(SCREAMING_SNAKE_CASE_					)


if __name__ == "__main__":
					UpperCamelCase__         =      argparse.ArgumentParser()
					parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
					parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to flava checkpoint""")
					parser.add_argument("""--codebook_path""", default=None, type=str, help="""Path to flava codebook checkpoint""")
					parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
					UpperCamelCase__         =      parser.parse_args()

					convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)




