




import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    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 ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
   from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import ViTImageProcessor





class 	lowerCamelCase_							:


 '''simple docstring'''







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

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



 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

       SCREAMING_SNAKE_CASE_     						=      self.get_config()

       return config, pixel_values, labels



 def 				lowerCAmelCase_  (							self  :		Union[str, 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=_lowerCAmelCase    ,    initializer_range=self.initializer_range    ,    encoder_stride=self.encoder_stride    ,    )



 def 				lowerCAmelCase_  (							self  :		str    ,    _lowerCAmelCase  :		List[str]    ,    _lowerCAmelCase  :		Tuple    ,    _lowerCAmelCase  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      ViTModel(config=_lowerCAmelCase						)
       model.to(_lowerCAmelCase						)
       model.eval()
       SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase						)
       self.parent.assertEqual(result.last_hidden_state.shape    ,    (self.batch_size, self.seq_length, self.hidden_size)						)



 def 				lowerCAmelCase_  (							self  :		List[Any]    ,    _lowerCAmelCase  :		int    ,    _lowerCAmelCase  :		Any    ,    _lowerCAmelCase  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_     						=      ViTForMaskedImageModeling(config=_lowerCAmelCase						)
       model.to(_lowerCAmelCase						)
       model.eval()
       SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase						)
       self.parent.assertEqual(
           result.reconstruction.shape    ,    (self.batch_size, self.num_channels, self.image_size, self.image_size)						)

       # test greyscale images
       SCREAMING_SNAKE_CASE_     						=      1
       SCREAMING_SNAKE_CASE_     						=      ViTForMaskedImageModeling(_lowerCAmelCase						)
       model.to(_lowerCAmelCase						)
       model.eval()

       SCREAMING_SNAKE_CASE_     						=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
       SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase						)
       self.parent.assertEqual(result.reconstruction.shape    ,    (self.batch_size, 1, self.image_size, self.image_size)						)



 def 				lowerCAmelCase_  (							self  :		Tuple    ,    _lowerCAmelCase  :		List[Any]    ,    _lowerCAmelCase  :		int    ,    _lowerCAmelCase  :		Dict						):
       SCREAMING_SNAKE_CASE_     						=      self.type_sequence_label_size
       SCREAMING_SNAKE_CASE_     						=      ViTForImageClassification(_lowerCAmelCase						)
       model.to(_lowerCAmelCase						)
       model.eval()
       SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase    ,    labels=_lowerCAmelCase						)
       self.parent.assertEqual(result.logits.shape    ,    (self.batch_size, self.type_sequence_label_size)						)

       # test greyscale images
       SCREAMING_SNAKE_CASE_     						=      1
       SCREAMING_SNAKE_CASE_     						=      ViTForImageClassification(_lowerCAmelCase						)
       model.to(_lowerCAmelCase						)
       model.eval()

       SCREAMING_SNAKE_CASE_     						=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
       SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase						)
       self.parent.assertEqual(result.logits.shape    ,    (self.batch_size, self.type_sequence_label_size)						)




 def 				lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_     						=      self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE_
       )						,			(
           SCREAMING_SNAKE_CASE_
       )						,			(
           SCREAMING_SNAKE_CASE_
       )						,			
       )     						=      config_and_inputs
       SCREAMING_SNAKE_CASE_     						=      {'pixel_values': pixel_values}
       return config, inputs_dict





@require_torch
class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE	,       _SCREAMING_SNAKE_CASE	,       unittest.TestCase ):


 '''simple docstring'''







 lowercase_       =   (
     (
         ViTModel,
         ViTForImageClassification,
         ViTForMaskedImageModeling,
     )
     if is_torch_available()
     else ()
 )
 lowercase_       =   (
     {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
     if is_torch_available()
     else {}
 )
 lowercase_       =   True

 lowercase_       =   False
 lowercase_       =   False
 lowercase_       =   False



 def 				lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_     						=      ViTModelTester(self						)
       SCREAMING_SNAKE_CASE_     						=      ConfigTester(self    ,    config_class=_lowerCAmelCase    ,    has_text_modality=_lowerCAmelCase    ,    hidden_size=37						)



 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       self.config_tester.run_common_tests()



 @unittest.skip(reason='ViT does not use inputs_embeds'						)
 def 				lowerCAmelCase_  (							self  :		Tuple						):
       pass



 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             SCREAMING_SNAKE_CASE_     						=      model_class(_lowerCAmelCase						)
             self.assertIsInstance(model.get_input_embeddings()    ,    (nn.Module)						)
             SCREAMING_SNAKE_CASE_     						=      model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(_lowerCAmelCase    ,    nn.Linear						)						)



 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs_for_common()

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

             SCREAMING_SNAKE_CASE_     						=      ['pixel_values']
             self.assertListEqual(arg_names[:1]    ,    _lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Union[str, Any]						):
       SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Tuple						):
       SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*_lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase						)




 @slow
 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             SCREAMING_SNAKE_CASE_     						=      ViTModel.from_pretrained(_lowerCAmelCase						)
             self.assertIsNotNone(_lowerCAmelCase						)


def  UpperCAmelCase_							(				) ->			List[str]:
      SCREAMING_SNAKE_CASE_     						=      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'							)
      return image





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


 '''simple docstring'''







 @cached_property
 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'						) if is_vision_available() else None



 @slow
 def 				lowerCAmelCase_  (							self  :		Any						):
       SCREAMING_SNAKE_CASE_     						=      ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'						).to(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      self.default_image_processor
       SCREAMING_SNAKE_CASE_     						=      prepare_img()
       SCREAMING_SNAKE_CASE_     						=      image_processor(images=_lowerCAmelCase    ,    return_tensors='pt'						).to(_lowerCAmelCase						)

       # forward pass
       with torch.no_grad():
             SCREAMING_SNAKE_CASE_     						=      model(**_lowerCAmelCase						)

       # verify the logits
       SCREAMING_SNAKE_CASE_     						=      torch.Size((1, 1_000)						)
       self.assertEqual(outputs.logits.shape    ,    _lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      torch.tensor([-0.2744, 0.8215, -0.0836]						).to(_lowerCAmelCase						)

       self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,    _lowerCAmelCase    ,    atol=1E-4						)						)



 @slow
 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       # ViT models have an `interpolate_pos_encoding` argument in their forward method,
       # allowing to interpolate the pre-trained position embeddings in order to use
       # the model on higher resolutions. The DINO model by Facebook AI leverages this
       # to visualize self-attention on higher resolution images.
       SCREAMING_SNAKE_CASE_     						=      ViTModel.from_pretrained('facebook/dino-vits8'						).to(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      ViTImageProcessor.from_pretrained('facebook/dino-vits8'    ,    size=480						)
       SCREAMING_SNAKE_CASE_     						=      prepare_img()
       SCREAMING_SNAKE_CASE_     						=      image_processor(images=_lowerCAmelCase    ,    return_tensors='pt'						)
       SCREAMING_SNAKE_CASE_     						=      inputs.pixel_values.to(_lowerCAmelCase						)

       # forward pass
       with torch.no_grad():
             SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase    ,    interpolate_pos_encoding=_lowerCAmelCase						)

       # verify the logits
       SCREAMING_SNAKE_CASE_     						=      torch.Size((1, 3_601, 384)						)
       self.assertEqual(outputs.last_hidden_state.shape    ,    _lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      torch.tensor(
           [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]						).to(_lowerCAmelCase						)

       self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]    ,    _lowerCAmelCase    ,    atol=1E-4						)						)




 @slow
 @require_accelerate
 @require_torch_gpu
 def 				lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_     						=      ViTModel.from_pretrained('facebook/dino-vits8'    ,    torch_dtype=torch.floataa    ,    device_map='auto'						)
       SCREAMING_SNAKE_CASE_     						=      self.default_image_processor

       SCREAMING_SNAKE_CASE_     						=      prepare_img()
       SCREAMING_SNAKE_CASE_     						=      image_processor(images=_lowerCAmelCase    ,    return_tensors='pt'						)
       SCREAMING_SNAKE_CASE_     						=      inputs.pixel_values.to(_lowerCAmelCase						)

       # forward pass to make sure inference works in fp16
       with torch.no_grad():
             SCREAMING_SNAKE_CASE_     						=      model(_lowerCAmelCase						)





import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel

def  UpperCAmelCase_							(				__UpperCAmelCase						:     List[str]							) ->			List[Any]:
      # vision encoder
      if "img_encoder.pos_embed" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.pos_embed'      ,      'vision_model.embeddings.position_embeddings'							)
      if "img_encoder.patch_embed.proj" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.patch_embed.proj'      ,      'vision_model.embeddings.patch_embeddings.projection'							)
      if "img_encoder.patch_embed.norm" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.patch_embed.norm'      ,      'vision_model.embeddings.layernorm'							)
      if "img_encoder.layers" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.layers'      ,      'vision_model.encoder.stages'							)
      if "blocks" in name and "res" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('blocks'      ,      'layers'							)
      if "attn" in name and "pre_assign" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('attn'      ,      'self_attn'							)
      if "proj" in name and "self_attn" in name and "text" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('proj'      ,      'out_proj'							)
      if "pre_assign_attn.attn.proj" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('pre_assign_attn.attn.proj'      ,      'pre_assign_attn.attn.out_proj'							)
      if "norm1" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('norm1'      ,      'layer_norm1'							)
      if "norm2" in name and "pre_assign" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('norm2'      ,      'layer_norm2'							)
      if "img_encoder.norm" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.norm'      ,      'vision_model.layernorm'							)
      # text encoder
      if "text_encoder.token_embedding" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder.token_embedding'      ,      'text_model.embeddings.token_embedding'							)
      if "text_encoder.positional_embedding" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder.positional_embedding'      ,      'text_model.embeddings.position_embedding.weight'							)
      if "text_encoder.transformer.resblocks." in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder.transformer.resblocks.'      ,      'text_model.encoder.layers.'							)
      if "ln_1" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('ln_1'      ,      'layer_norm1'							)
      if "ln_2" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('ln_2'      ,      'layer_norm2'							)
      if "c_fc" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('c_fc'      ,      'fc1'							)
      if "c_proj" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('c_proj'      ,      'fc2'							)
      if "text_encoder" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder'      ,      'text_model'							)
      if "ln_final" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('ln_final'      ,      'final_layer_norm'							)
      # projection layers
      if "img_projector.linear_hidden." in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_projector.linear_hidden.'      ,      'visual_projection.'							)
      if "img_projector.linear_out." in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_projector.linear_out.'      ,      'visual_projection.3.'							)
      if "text_projector.linear_hidden" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_projector.linear_hidden'      ,      'text_projection'							)
      if "text_projector.linear_out" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_projector.linear_out'      ,      'text_projection.3'							)

      return name

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Union[str, Any]      ,      __UpperCAmelCase						:     int							) ->			List[str]:
      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE_     						=      orig_state_dict.pop(__UpperCAmelCase							)

            if "qkv" in key:
                  # weights and biases of the key, value and query projections of vision encoder's attention layers require special treatment:
                  # we need to split them up into separate matrices/vectors
                  SCREAMING_SNAKE_CASE_     						=      key.split('.'							)
                  SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      int(key_split[2]							), int(key_split[4]							)
                  SCREAMING_SNAKE_CASE_     						=      config.vision_config.hidden_size
                  if "weight" in key:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim, :]
                        SCREAMING_SNAKE_CASE_     						=      val[dim : dim * 2, :]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:, :]
                  else:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim]
                        SCREAMING_SNAKE_CASE_     						=      val[dim : dim * 2]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:]
            elif "in_proj" in key:
                  # weights and biases of the key, value and query projections of text encoder's attention layers require special treatment:
                  # we need to split them up into separate matrices/vectors
                  SCREAMING_SNAKE_CASE_     						=      key.split('.'							)
                  SCREAMING_SNAKE_CASE_     						=      int(key_split[3]							)
                  SCREAMING_SNAKE_CASE_     						=      config.text_config.hidden_size
                  if "weight" in key:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim, :]
                        SCREAMING_SNAKE_CASE_     						=      val[
                            dim : dim * 2, :
                        ]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:, :]
                  else:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim]
                        SCREAMING_SNAKE_CASE_     						=      val[dim : dim * 2]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:]
            else:
                  SCREAMING_SNAKE_CASE_     						=      rename_key(__UpperCAmelCase							)
                  # squeeze if necessary
                  if (
                      "text_projection.0" in new_name
                      or "text_projection.3" in new_name
                      or "visual_projection.0" in new_name
                      or "visual_projection.3" in new_name
                  ):
                        SCREAMING_SNAKE_CASE_     						=      val.squeeze_()
                  else:
                        SCREAMING_SNAKE_CASE_     						=      val

      return orig_state_dict

def  UpperCAmelCase_							(				) ->			List[Any]:
      SCREAMING_SNAKE_CASE_     						=      'http://images.cocodataset.org/val2017/000000039769.jpg'
      SCREAMING_SNAKE_CASE_     						=      Image.open(requests.get(__UpperCAmelCase      ,      stream=__UpperCAmelCase							).raw							)
      return im

@torch.no_grad()
def  UpperCAmelCase_							(				__UpperCAmelCase						:     int      ,      __UpperCAmelCase						:     str      ,      __UpperCAmelCase						:     Dict="groupvit-gcc-yfcc"      ,      __UpperCAmelCase						:     Any=False							) ->			Optional[int]:
      SCREAMING_SNAKE_CASE_     						=      GroupViTConfig()
      SCREAMING_SNAKE_CASE_     						=      GroupViTModel(__UpperCAmelCase							).eval()

      SCREAMING_SNAKE_CASE_     						=      torch.load(__UpperCAmelCase      ,      map_location='cpu'							)['model']
      SCREAMING_SNAKE_CASE_     						=      convert_state_dict(__UpperCAmelCase      ,      __UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      model.load_state_dict(__UpperCAmelCase      ,      strict=__UpperCAmelCase							)
      assert missing_keys == ["text_model.embeddings.position_ids"]
      assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(__UpperCAmelCase							) == 0)

      # verify result
      SCREAMING_SNAKE_CASE_     						=      CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32'							)
      SCREAMING_SNAKE_CASE_     						=      prepare_img()
      SCREAMING_SNAKE_CASE_     						=      processor(text=['a photo of a cat', 'a photo of a dog']      ,      images=__UpperCAmelCase      ,      padding=__UpperCAmelCase      ,      return_tensors='pt'							)

      with torch.no_grad():
            SCREAMING_SNAKE_CASE_     						=      model(**__UpperCAmelCase							)

      if model_name == "groupvit-gcc-yfcc":
            SCREAMING_SNAKE_CASE_     						=      torch.tensor([[1_3.3_5_2_3, 6.3_6_2_9]]							)
      elif model_name == "groupvit-gcc-redcaps":
            SCREAMING_SNAKE_CASE_     						=      torch.tensor([[1_6.1_8_7_3, 8.6_2_3_0]]							)
      else:
            raise ValueError(f"Model name {model_name} not supported."							)
      assert torch.allclose(outputs.logits_per_image      ,      __UpperCAmelCase      ,      atol=1E-3							)

      processor.save_pretrained(__UpperCAmelCase							)
      model.save_pretrained(__UpperCAmelCase							)
      print('Successfully saved processor and model to'      ,      __UpperCAmelCase							)

      if push_to_hub:
            print('Pushing to the hub...'							)
            processor.push_to_hub(__UpperCAmelCase      ,      organization='nielsr'							)
            model.push_to_hub(__UpperCAmelCase      ,      organization='nielsr'							)


if __name__ == "__main__":
   lowerCamelCase__			:       int         =							argparse.ArgumentParser()
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, help='Path to dump the processor and PyTorch model.'
   )
   parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to GroupViT checkpoint')
   parser.add_argument(
       '--model_name',
       default='groupvit-gccy-fcc',
       type=str,
       help='Name of the model. Expecting either \'groupvit-gcc-yfcc\' or \'groupvit-gcc-redcaps\'',
   )
   parser.add_argument(
       '--push_to_hub',
       action='store_true',
       help='Whether or not to push the converted model and processor to the 🤗 hub using the provided `model_name`.',
   )
   lowerCamelCase__			:       List[str]         =							parser.parse_args()

   convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)