





'''simple docstring'''




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
 import torch

if is_vision_available():
 from PIL import Image

 from transformers import ChineseCLIPImageProcessor
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     def __init__(   self							:			Dict   ,       __snake_case							:			int   ,       __snake_case							:			Optional[Any]=7   ,       __snake_case							:			Optional[int]=3   ,       __snake_case							:			Union[str, Any]=18   ,       __snake_case							:			Union[str, Any]=30   ,       __snake_case							:			int=4_00   ,       __snake_case							:			Any=True   ,       __snake_case							:			List[str]=None   ,       __snake_case							:			List[Any]=True   ,       __snake_case							:			Optional[int]=None   ,       __snake_case							:			Optional[Any]=True   ,       __snake_case							:			Tuple=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]   ,       __snake_case							:			str=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]   ,       __snake_case							:			Tuple=True   ,       )->	List[Any]:
        snake_case 						=   size if size is not None else {"""height""": 2_24, """width""": 2_24}
        snake_case 						=   crop_size if crop_size is not None else {"""height""": 18, """width""": 18}
        snake_case 						=   parent
        snake_case 						=   batch_size
        snake_case 						=   num_channels
        snake_case 						=   image_size
        snake_case 						=   min_resolution
        snake_case 						=   max_resolution
        snake_case 						=   do_resize
        snake_case 						=   size
        snake_case 						=   do_center_crop
        snake_case 						=   crop_size
        snake_case 						=   do_normalize
        snake_case 						=   image_mean
        snake_case 						=   image_std
        snake_case 						=   do_convert_rgb



     def      lowerCAmelCase			(   self							:			Dict						)->	Any:
        return {
            "do_resize": self.do_resize,
            "size": self.size,
            "do_center_crop": self.do_center_crop,
            "crop_size": self.crop_size,
            "do_normalize": self.do_normalize,
            "image_mean": self.image_mean,
            "image_std": self.image_std,
            "do_convert_rgb": self.do_convert_rgb,
        }



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			Tuple=False   ,       __snake_case							:			Dict=False   ,       __snake_case							:			List[Any]=False						)->	Union[str, Any]:
        assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

        if equal_resolution:
           snake_case 						=   []
           for i in range(self.batch_size						):
              image_inputs.append(
                  np.random.randint(
                      2_55   ,       size=(self.num_channels, self.max_resolution, self.max_resolution)   ,       dtype=np.uinta						)						)
        else:
           snake_case 						=   []
           for i in range(self.batch_size						):
              snake_case						,    snake_case 						=   np.random.choice(np.arange(self.min_resolution   ,       self.max_resolution						)   ,       2						)
              image_inputs.append(np.random.randint(2_55   ,       size=(self.num_channels, width, height)   ,       dtype=np.uinta						)						)

        if not numpify and not torchify:
           # PIL expects the channel dimension as last dimension
           snake_case 						=   [Image.fromarray(np.moveaxis(__snake_case   ,       0   ,       -1						)						) for x in image_inputs]

        if torchify:
           snake_case 						=   [torch.from_numpy(__snake_case						) for x in image_inputs]

        return image_inputs




@require_torch
@require_vision
class 			_lowerCAmelCase		(			A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   ChineseCLIPImageProcessor if is_vision_available() else None



     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Union[str, Any]:
        snake_case 						=   ChineseCLIPImageProcessingTester(self   ,       do_center_crop=__snake_case						)



     @property
     def      lowerCAmelCase			(   self							:			int						)->	List[str]:
        return self.image_processor_tester.prepare_image_processor_dict()



     def      lowerCAmelCase			(   self							:			List[str]						)->	int:
        snake_case 						=   self.image_processing_class(**self.image_processor_dict						)
        self.assertTrue(hasattr(__snake_case   ,       """do_resize"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """size"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """do_center_crop"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """center_crop"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """do_normalize"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """image_mean"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """image_std"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """do_convert_rgb"""						)						)



     def      lowerCAmelCase			(   self							:			Tuple						)->	Optional[int]:
        snake_case 						=   self.image_processing_class.from_dict(self.image_processor_dict						)
        self.assertEqual(image_processor.size   ,       {"""height""": 2_24, """width""": 2_24}						)
        self.assertEqual(image_processor.crop_size   ,       {"""height""": 18, """width""": 18}						)

        snake_case 						=   self.image_processing_class.from_dict(self.image_processor_dict   ,       size=42   ,       crop_size=84						)
        self.assertEqual(image_processor.size   ,       {"""shortest_edge""": 42}						)
        self.assertEqual(image_processor.crop_size   ,       {"""height""": 84, """width""": 84}						)



     def      lowerCAmelCase			(   self							:			Any						)->	Union[str, Any]:
        pass



     def      lowerCAmelCase			(   self							:			int						)->	List[Any]:
        # Initialize image_processing
        snake_case 						=   self.image_processing_class(**self.image_processor_dict						)
        # create random PIL images
        snake_case 						=   self.image_processor_tester.prepare_inputs(equal_resolution=__snake_case						)
        for image in image_inputs:
           self.assertIsInstance(__snake_case   ,       Image.Image						)

        # Test not batched input
        snake_case 						=   image_processing(image_inputs[0]   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                1,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )

        # Test batched
        snake_case 						=   image_processing(__snake_case   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )



     def      lowerCAmelCase			(   self							:			Tuple						)->	Optional[int]:
        # Initialize image_processing
        snake_case 						=   self.image_processing_class(**self.image_processor_dict						)
        # create random numpy tensors
        snake_case 						=   self.image_processor_tester.prepare_inputs(equal_resolution=__snake_case   ,       numpify=__snake_case						)
        for image in image_inputs:
           self.assertIsInstance(__snake_case   ,       np.ndarray						)

        # Test not batched input
        snake_case 						=   image_processing(image_inputs[0]   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                1,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )

        # Test batched
        snake_case 						=   image_processing(__snake_case   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )



     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[Any]:
        # Initialize image_processing
        snake_case 						=   self.image_processing_class(**self.image_processor_dict						)
        # create random PyTorch tensors
        snake_case 						=   self.image_processor_tester.prepare_inputs(equal_resolution=__snake_case   ,       torchify=__snake_case						)
        for image in image_inputs:
           self.assertIsInstance(__snake_case   ,       torch.Tensor						)

        # Test not batched input
        snake_case 						=   image_processing(image_inputs[0]   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                1,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )

        # Test batched
        snake_case 						=   image_processing(__snake_case   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )




@require_torch
@require_vision
class 			_lowerCAmelCase		(			A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   ChineseCLIPImageProcessor if is_vision_available() else None



     def      lowerCAmelCase			(   self							:			List[str]						)->	Optional[int]:
        snake_case 						=   ChineseCLIPImageProcessingTester(self   ,       num_channels=4   ,       do_center_crop=__snake_case						)
        snake_case 						=   3



     @property
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Any:
        return self.image_processor_tester.prepare_image_processor_dict()



     def      lowerCAmelCase			(   self							:			Dict						)->	int:
        snake_case 						=   self.image_processing_class(**self.image_processor_dict						)
        self.assertTrue(hasattr(__snake_case   ,       """do_resize"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """size"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """do_center_crop"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """center_crop"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """do_normalize"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """image_mean"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """image_std"""						)						)
        self.assertTrue(hasattr(__snake_case   ,       """do_convert_rgb"""						)						)



     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Any:
        pass



     def      lowerCAmelCase			(   self							:			Tuple						)->	int:
        # Initialize image_processing
        snake_case 						=   self.image_processing_class(**self.image_processor_dict						)
        # create random PIL images
        snake_case 						=   self.image_processor_tester.prepare_inputs(equal_resolution=__snake_case						)
        for image in image_inputs:
           self.assertIsInstance(__snake_case   ,       Image.Image						)

        # Test not batched input
        snake_case 						=   image_processing(image_inputs[0]   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                1,
                self.expected_encoded_image_num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )

        # Test batched
        snake_case 						=   image_processing(__snake_case   ,       return_tensors="""pt"""						).pixel_values
        self.assertEqual(
            encoded_images.shape   ,       (
                self.image_processor_tester.batch_size,
                self.expected_encoded_image_num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )   ,       )








'''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 DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[int]      ,    __lowerCAmelCase       :							List[Any]=False		)		->  List[str]:
   snake_case 						=   []
   for i in range(config.num_hidden_layers		):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append((F'''blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')		)
      rename_keys.append((F'''blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')		)
      rename_keys.append((F'''blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')		)
      rename_keys.append((F'''blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')		)
      rename_keys.append((F'''blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')		)
      rename_keys.append((F'''blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')		)
      rename_keys.append((F'''blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')		)
      rename_keys.append((F'''blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')		)
      rename_keys.append((F'''blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')		)
      rename_keys.append((F'''blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')		)

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ("""cls_token""", """vit.embeddings.cls_token"""),
           ("""patch_embed.proj.weight""", """vit.embeddings.patch_embeddings.projection.weight"""),
           ("""patch_embed.proj.bias""", """vit.embeddings.patch_embeddings.projection.bias"""),
           ("""pos_embed""", """vit.embeddings.position_embeddings"""),
       ]		)

   if base_model:
      # layernorm + pooler
      rename_keys.extend(
          [
              ("""norm.weight""", """layernorm.weight"""),
              ("""norm.bias""", """layernorm.bias"""),
              ("""pre_logits.fc.weight""", """pooler.dense.weight"""),
              ("""pre_logits.fc.bias""", """pooler.dense.bias"""),
          ]		)

      # if just the base model, we should remove "vit" from all keys that start with "vit"
      snake_case 						=   [(pair[0], pair[1][4:]) if pair[1].startswith("""vit"""		) else pair for pair in rename_keys]
   else:
      # layernorm + classification head
      rename_keys.extend(
          [
              ("""norm.weight""", """vit.layernorm.weight"""),
              ("""norm.bias""", """vit.layernorm.bias"""),
              ("""head.weight""", """classifier.weight"""),
              ("""head.bias""", """classifier.bias"""),
          ]		)

   return rename_keys





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Dict      ,    __lowerCAmelCase       :							Optional[Any]      ,    __lowerCAmelCase       :							int=False		)		->  Union[str, Any]:
   for i in range(config.num_hidden_layers		):
      if base_model:
         snake_case 						=   """"""
      else:
         snake_case 						=   """vit."""
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      snake_case 						=   state_dict.pop(F'''blocks.{i}.attn.qkv.weight'''		)
      snake_case 						=   state_dict.pop(F'''blocks.{i}.attn.qkv.bias'''		)
      # next, add query, keys and values (in that order) to the state dict
      snake_case 						=   in_proj_weight[
          : config.hidden_size, :
      ]
      snake_case 						=   in_proj_bias[: config.hidden_size]
      snake_case 						=   in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      snake_case 						=   in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      snake_case 						=   in_proj_weight[
          -config.hidden_size :, :
      ]
      snake_case 						=   in_proj_bias[-config.hidden_size :]





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]		)		->  List[str]:
   snake_case 						=   ["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(__lowerCAmelCase      ,    __lowerCAmelCase		)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[Any]      ,    __lowerCAmelCase       :							int      ,    __lowerCAmelCase       :							Tuple		)		->  int:
   snake_case 						=   dct.pop(__lowerCAmelCase		)
   snake_case 						=   val





def 	__lowerCamelCase			(      )		->  Optional[Any]:
   snake_case 						=   """http://images.cocodataset.org/val2017/000000039769.jpg"""
   snake_case 						=   Image.open(requests.get(__lowerCAmelCase      ,    stream=__lowerCAmelCase		).raw		)
   return im





@torch.no_grad()
def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]      ,    __lowerCAmelCase       :							Any		)		->  Dict:
   snake_case 						=   ViTConfig()
   snake_case 						=   False
   # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
   if vit_name[-5:] == "in21k":
      snake_case 						=   True
      snake_case 						=   int(vit_name[-12:-10]		)
      snake_case 						=   int(vit_name[-9:-6]		)
   else:
      snake_case 						=   10_00
      snake_case 						=   """huggingface/label-files"""
      snake_case 						=   """imagenet-1k-id2label.json"""
      snake_case 						=   json.load(open(hf_hub_download(__lowerCAmelCase      ,    __lowerCAmelCase      ,    repo_type="""dataset"""		)      ,    """r"""		)		)
      snake_case 						=   {int(__lowerCAmelCase		): v for k, v in idalabel.items()}
      snake_case 						=   idalabel
      snake_case 						=   {v: k for k, v in idalabel.items()}
      snake_case 						=   int(vit_name[-6:-4]		)
      snake_case 						=   int(vit_name[-3:]		)
   # size of the architecture
   if "deit" in vit_name:
      if vit_name[9:].startswith("""tiny"""		):
         snake_case 						=   1_92
         snake_case 						=   7_68
         snake_case 						=   12
         snake_case 						=   3
      elif vit_name[9:].startswith("""small"""		):
         snake_case 						=   3_84
         snake_case 						=   15_36
         snake_case 						=   12
         snake_case 						=   6
      else:
         pass
   else:
      if vit_name[4:].startswith("""small"""		):
         snake_case 						=   7_68
         snake_case 						=   23_04
         snake_case 						=   8
         snake_case 						=   8
      elif vit_name[4:].startswith("""base"""		):
         pass
      elif vit_name[4:].startswith("""large"""		):
         snake_case 						=   10_24
         snake_case 						=   40_96
         snake_case 						=   24
         snake_case 						=   16
      elif vit_name[4:].startswith("""huge"""		):
         snake_case 						=   12_80
         snake_case 						=   51_20
         snake_case 						=   32
         snake_case 						=   16

    # load original model from timm
   snake_case 						=   timm.create_model(__lowerCAmelCase      ,    pretrained=__lowerCAmelCase		)
   timm_model.eval()

   # load state_dict of original model, remove and rename some keys
   snake_case 						=   timm_model.state_dict()
   if base_model:
      remove_classification_head_(__lowerCAmelCase		)
   snake_case 						=   create_rename_keys(__lowerCAmelCase      ,    __lowerCAmelCase		)
   for src, dest in rename_keys:
      rename_key(__lowerCAmelCase      ,    __lowerCAmelCase      ,    __lowerCAmelCase		)
   read_in_q_k_v(__lowerCAmelCase      ,    __lowerCAmelCase      ,    __lowerCAmelCase		)

   # load HuggingFace model
   if vit_name[-5:] == "in21k":
      snake_case 						=   ViTModel(__lowerCAmelCase		).eval()
   else:
      snake_case 						=   ViTForImageClassification(__lowerCAmelCase		).eval()
   model.load_state_dict(__lowerCAmelCase		)

   # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
   if "deit" in vit_name:
      snake_case 						=   DeiTImageProcessor(size=config.image_size		)
   else:
      snake_case 						=   ViTImageProcessor(size=config.image_size		)
   snake_case 						=   image_processor(images=prepare_img()      ,    return_tensors="""pt"""		)
   snake_case 						=   encoding["""pixel_values"""]
   snake_case 						=   model(__lowerCAmelCase		)

   if base_model:
      snake_case 						=   timm_model.forward_features(__lowerCAmelCase		)
      assert timm_pooled_output.shape == outputs.pooler_output.shape
      assert torch.allclose(__lowerCAmelCase      ,    outputs.pooler_output      ,    atol=1e-3		)
   else:
      snake_case 						=   timm_model(__lowerCAmelCase		)
      assert timm_logits.shape == outputs.logits.shape
      assert torch.allclose(__lowerCAmelCase      ,    outputs.logits      ,    atol=1e-3		)

   Path(__lowerCAmelCase		).mkdir(exist_ok=__lowerCAmelCase		)
   print(F'''Saving model {vit_name} to {pytorch_dump_folder_path}'''		)
   model.save_pretrained(__lowerCAmelCase		)
   print(F'''Saving image processor to {pytorch_dump_folder_path}'''		)
   image_processor.save_pretrained(__lowerCAmelCase		)


if __name__ == "__main__":
 _SCREAMING_SNAKE_CASE							     =			argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--vit_name",
     default="vit_base_patch16_224",
     type=str,
     help="Name of the ViT 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."
 )

 _SCREAMING_SNAKE_CASE							     =			parser.parse_args()
 convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)

