
from __future__ import annotations

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import TFViTForImageClassification, TFViTModel


if is_vision_available():
						from PIL import Image

						from transformers import ViTImageProcessor





class   UpperCamelCase_      :







						'''simple docstring'''



						def __init__(			self							:List[str]		,    lowerCAmelCase__							:List[Any]		,    lowerCAmelCase__							:Union[str, Any]=13		,    lowerCAmelCase__							:List[str]=30		,    lowerCAmelCase__							:List[str]=2		,    lowerCAmelCase__							:List[str]=3		,    lowerCAmelCase__							:str=True		,    lowerCAmelCase__							:List[Any]=True		,    lowerCAmelCase__							:Union[str, Any]=32		,    lowerCAmelCase__							:int=2		,    lowerCAmelCase__							:List[str]=4		,    lowerCAmelCase__							:Any=37		,    lowerCAmelCase__							:Union[str, Any]="gelu"		,    lowerCAmelCase__							:List[Any]=0.1		,    lowerCAmelCase__							:Tuple=0.1		,    lowerCAmelCase__							:int=10		,    lowerCAmelCase__							:Union[str, Any]=0.02		,    lowerCAmelCase__							:Optional[int]=3		,    lowerCAmelCase__							:int=None		,    )  ->Union[str, Any]:
										lowercase         =      parent
										lowercase         =      batch_size
										lowercase         =      image_size
										lowercase         =      patch_size
										lowercase         =      num_channels
										lowercase         =      is_training
										lowercase         =      use_labels
										lowercase         =      hidden_size
										lowercase         =      num_hidden_layers
										lowercase         =      num_attention_heads
										lowercase         =      intermediate_size
										lowercase         =      hidden_act
										lowercase         =      hidden_dropout_prob
										lowercase         =      attention_probs_dropout_prob
										lowercase         =      type_sequence_label_size
										lowercase         =      initializer_range
										lowercase         =      scope

										# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
										lowercase         =      (image_size // patch_size) ** 2
										lowercase         =      num_patches + 1





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->Tuple:
										lowercase         =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

										lowercase         =      None
										if self.use_labels:
														lowercase         =      ids_tensor([self.batch_size]		,    self.type_sequence_label_size     )

										lowercase         =      self.get_config()

										return config, pixel_values, labels





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->List[Any]:
										return ViTConfig(
										    image_size=self.image_size		,    patch_size=self.patch_size		,    num_channels=self.num_channels		,    hidden_size=self.hidden_size		,    num_hidden_layers=self.num_hidden_layers		,    num_attention_heads=self.num_attention_heads		,    intermediate_size=self.intermediate_size		,    hidden_act=self.hidden_act		,    hidden_dropout_prob=self.hidden_dropout_prob		,    attention_probs_dropout_prob=self.attention_probs_dropout_prob		,    is_decoder=SCREAMING_SNAKE_CASE_		,    initializer_range=self.initializer_range		,    )





						def        SCREAMING_SNAKE_CASE(			self							:List[str]		,    lowerCAmelCase__							:List[str]		,    lowerCAmelCase__							:Union[str, Any]		,    lowerCAmelCase__							:str     )  ->Dict:
										lowercase         =      TFViTModel(config=SCREAMING_SNAKE_CASE_     )
										lowercase         =      model(SCREAMING_SNAKE_CASE_		,    training=SCREAMING_SNAKE_CASE_     )
										self.parent.assertEqual(result.last_hidden_state.shape		,    (self.batch_size, self.seq_length, self.hidden_size)     )

										# Test with an image with different size than the one specified in config.
										lowercase         =      self.image_size // 2
										lowercase         =      pixel_values[:, :, :image_size, :image_size]
										lowercase         =      model(SCREAMING_SNAKE_CASE_		,    interpolate_pos_encoding=SCREAMING_SNAKE_CASE_		,    training=SCREAMING_SNAKE_CASE_     )
										lowercase         =      (image_size // self.patch_size) ** 2 + 1
										self.parent.assertEqual(result.last_hidden_state.shape		,    (self.batch_size, seq_length, self.hidden_size)     )





						def        SCREAMING_SNAKE_CASE(			self							:List[str]		,    lowerCAmelCase__							:Tuple		,    lowerCAmelCase__							:Dict		,    lowerCAmelCase__							:Dict     )  ->Any:
										lowercase         =      self.type_sequence_label_size
										lowercase         =      TFViTForImageClassification(SCREAMING_SNAKE_CASE_     )
										lowercase         =      model(SCREAMING_SNAKE_CASE_		,    labels=SCREAMING_SNAKE_CASE_		,    training=SCREAMING_SNAKE_CASE_     )
										self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.type_sequence_label_size)     )

										# Test with an image with different size than the one specified in config.
										lowercase         =      self.image_size // 2
										lowercase         =      pixel_values[:, :, :image_size, :image_size]
										lowercase         =      model(SCREAMING_SNAKE_CASE_		,    interpolate_pos_encoding=SCREAMING_SNAKE_CASE_		,    training=SCREAMING_SNAKE_CASE_     )
										self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.type_sequence_label_size)     )

										# test greyscale images
										lowercase         =      1
										lowercase         =      TFViTForImageClassification(SCREAMING_SNAKE_CASE_     )
										lowercase         =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
										lowercase         =      model(SCREAMING_SNAKE_CASE_     )
										self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.type_sequence_label_size)     )





						def        SCREAMING_SNAKE_CASE(			self							:List[Any]     )  ->Dict:
										lowercase         =      self.prepare_config_and_inputs()
										lowercase         =      config_and_inputs
										lowercase         =      {"""pixel_values""": pixel_values}
										return config, inputs_dict





@require_tf
class   UpperCamelCase_      (      UpperCAmelCase_		,    UpperCAmelCase_		,    unittest.TestCase		):







						'''simple docstring'''



						UpperCamelCase      :  Dict				   =					(TFViTModel, TFViTForImageClassification) if is_tf_available() else ()
						UpperCamelCase      :  Dict				   =					(
						    {'feature-extraction': TFViTModel, 'image-classification': TFViTForImageClassification}
						    if is_tf_available()
						    else {}
						)

						UpperCamelCase      :  Tuple				   =					False
						UpperCamelCase      :  List[str]				   =					False
						UpperCamelCase      :  Union[str, Any]				   =					False





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]     )  ->Optional[Any]:
										lowercase         =      TFViTModelTester(self     )
										lowercase         =      ConfigTester(self		,    config_class=SCREAMING_SNAKE_CASE_		,    has_text_modality=SCREAMING_SNAKE_CASE_		,    hidden_size=37     )





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->Dict:
										self.config_tester.run_common_tests()





						@unittest.skip(reason="ViT does not use inputs_embeds"     )
						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->int:
										pass





						@unittest.skip(reason="ViT does not use inputs_embeds"     )
						def        SCREAMING_SNAKE_CASE(			self							:Dict     )  ->Union[str, Any]:
										pass





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Union[str, Any]:
										lowercase         =      self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														lowercase         =      model_class(SCREAMING_SNAKE_CASE_     )
														self.assertIsInstance(model.get_input_embeddings()		,    (tf.keras.layers.Layer)     )
														lowercase         =      model.get_output_embeddings()
														self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE_		,    tf.keras.layers.Layer     )     )





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Optional[int]:
										lowercase         =      self.model_tester.prepare_config_and_inputs_for_common()

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

														lowercase         =      ["""pixel_values"""]
														self.assertListEqual(arg_names[:1]		,    SCREAMING_SNAKE_CASE_     )





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->Dict:
										lowercase         =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_     )





						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->Dict:
										lowercase         =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_     )





						@slow
						def        SCREAMING_SNAKE_CASE(			self							:Tuple     )  ->Union[str, Any]:
										lowercase         =      TFViTModel.from_pretrained("google/vit-base-patch16-224"     )
										self.assertIsNotNone(SCREAMING_SNAKE_CASE_     )




def  __snake_case							(	):







				'''simple docstring'''


				lowercase         =      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
				return image





@require_tf
@require_vision
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						@cached_property
						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Dict:
										return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224"     ) if is_vision_available() else None





						@slow
						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Dict:
										lowercase         =      TFViTForImageClassification.from_pretrained("google/vit-base-patch16-224"     )

										lowercase         =      self.default_image_processor
										lowercase         =      prepare_img()
										lowercase         =      image_processor(images=SCREAMING_SNAKE_CASE_		,    return_tensors="tf"     )

										# forward pass
										lowercase         =      model(**SCREAMING_SNAKE_CASE_     )

										# verify the logits
										lowercase         =      tf.TensorShape((1, 1000)     )
										self.assertEqual(outputs.logits.shape		,    SCREAMING_SNAKE_CASE_     )

										lowercase         =      tf.constant([-0.27_44, 0.82_15, -0.08_36]     )

										tf.debugging.assert_near(outputs.logits[0, :3]		,    SCREAMING_SNAKE_CASE_		,    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 		UpperCamelCase						(   lowercase_ )		->						Union[str, Any]:


  '''simple docstring'''






  return sum(param.float().sum() if """encoder.embeddings""" not in key else 0 for key, param in state_dict.items() )

def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						List[Any]:


  '''simple docstring'''






  lowercase__	: int     		=   {}

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

    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.mim_head.cls.predictions"""			,			"""mmm_image_head""" )
    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.mlm_head.cls.predictions"""			,			"""mmm_text_head""" )
    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.itm_head.cls"""			,			"""itm_head""" )
    lowercase__	: Tuple     		=   key.replace("""heads.cmd.itm_head.pooler"""			,			"""itm_head.pooler""" )
    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.clip_head.logit_scale"""			,			"""flava.logit_scale""" )
    lowercase__	: Optional[int]     		=   key.replace("""heads.fairseq_mlm.cls.predictions"""			,			"""mlm_head""" )
    lowercase__	: List[Any]     		=   key.replace("""heads.imagenet.mim_head.cls.predictions"""			,			"""mim_head""" )
    lowercase__	: int     		=   key.replace("""mm_text_projection"""			,			"""flava.text_to_mm_projection""" )
    lowercase__	: Optional[Any]     		=   key.replace("""mm_image_projection"""			,			"""flava.image_to_mm_projection""" )
    lowercase__	: Optional[Any]     		=   key.replace("""image_encoder.module"""			,			"""flava.image_model""" )
    lowercase__	: Any     		=   key.replace("""text_encoder.module"""			,			"""flava.text_model""" )
    lowercase__	: Optional[Any]     		=   key.replace("""mm_encoder.module.encoder.cls_token"""			,			"""flava.multimodal_model.cls_token""" )
    lowercase__	: Tuple     		=   key.replace("""mm_encoder.module"""			,			"""flava.multimodal_model""" )
    lowercase__	: Any     		=   key.replace("""text_projection"""			,			"""flava.text_projection""" )
    lowercase__	: List[Any]     		=   key.replace("""image_projection"""			,			"""flava.image_projection""" )

    lowercase__	: str     		=   value.float()

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

  return upgrade

@torch.no_grad()
def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_			,			lowercase_=None )		->						Union[str, Any]:


  '''simple docstring'''






  if config_path is not None:
    lowercase__	: int     		=   FlavaConfig.from_pretrained(lowercase_ )
  else:
    lowercase__	: Optional[int]     		=   FlavaConfig()

  lowercase__	: List[Any]     		=   FlavaForPreTraining(lowercase_ ).eval()

  lowercase__	: Dict     		=   convert_dalle_checkpoint(lowercase_			,			lowercase_			,			save_checkpoint=lowercase_ )

  if os.path.exists(lowercase_ ):
    lowercase__	: Dict     		=   torch.load(lowercase_			,			map_location="""cpu""" )
  else:
    lowercase__	: Dict     		=   torch.hub.load_state_dict_from_url(lowercase_			,			map_location="""cpu""" )

  lowercase__	: int     		=   upgrade_state_dict(lowercase_			,			lowercase_ )
  hf_model.load_state_dict(lowercase_ )
  lowercase__	: Optional[int]     		=   hf_model.state_dict()
  lowercase__	: Optional[int]     		=   count_parameters(lowercase_ )
  lowercase__	: Any     		=   count_parameters(lowercase_ ) + count_parameters(lowercase_ )

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

  hf_model.save_pretrained(lowercase_ )


if __name__ == "__main__":
     lowerCamelCase__     :		int        = 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""")
     lowerCamelCase__     :		List[str]        = parser.parse_args()

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






