import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()



class  UpperCamelCase				( __lowercase   ,  __lowercase   ,  unittest.TestCase ):

     '''simple docstring'''

     A_										=  AutoencoderKL
     A_										=  'sample'
     A_										=  1E-2



     @property
     def 		UpperCamelCase_						(					self						) ->				Optional[int]:






           """simple docstring"""
           _lowerCamelCase			    =				4
           _lowerCamelCase			    =				3
           _lowerCamelCase			    =				(32, 32)

           _lowerCamelCase			    =				floats_tensor((batch_size, num_channels) + sizes						).to(A_						)

           return {"sample": image}



     @property
     def 		UpperCamelCase_						(					self						) ->				Any:






           """simple docstring"""
           return (3, 32, 32)



     @property
     def 		UpperCamelCase_						(					self						) ->				Dict:






           """simple docstring"""
           return (3, 32, 32)



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






           """simple docstring"""
           _lowerCamelCase			    =				{
               '''block_out_channels''': [32, 64],
               '''in_channels''': 3,
               '''out_channels''': 3,
               '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
               '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
               '''latent_channels''': 4,
           }
           _lowerCamelCase			    =				self.dummy_input
           return init_dict, inputs_dict



     def 		UpperCamelCase_						(					self						) ->				Dict:






           """simple docstring"""
           pass



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           pass



     @unittest.skipIf(torch_device == '''mps'''				,				'''Gradient checkpointing skipped on MPS'''						)
     def 		UpperCamelCase_						(					self						) ->				Optional[int]:






           """simple docstring"""
           # enable deterministic behavior for gradient checkpointing
           _lowerCamelCase    ,					_lowerCamelCase			    =				self.prepare_init_args_and_inputs_for_common()
           _lowerCamelCase			    =				self.model_class(**A_						)
           model.to(A_						)

           assert not model.is_gradient_checkpointing and model.training

           _lowerCamelCase			    =				model(**A_						).sample
           # run the backwards pass on the model. For backwards pass, for simplicity purpose,
           # we won't calculate the loss and rather backprop on out.sum()
           model.zero_grad()

           _lowerCamelCase			    =				torch.randn_like(A_						)
           _lowerCamelCase			    =				(out - labels).mean()
           loss.backward()

           # re-instantiate the model now enabling gradient checkpointing
           _lowerCamelCase			    =				self.model_class(**A_						)
           # clone model
           model_a.load_state_dict(model.state_dict()						)
           model_a.to(A_						)
           model_a.enable_gradient_checkpointing()

           assert model_a.is_gradient_checkpointing and model_a.training

           _lowerCamelCase			    =				model_a(**A_						).sample
           # run the backwards pass on the model. For backwards pass, for simplicity purpose,
           # we won't calculate the loss and rather backprop on out.sum()
           model_a.zero_grad()
           _lowerCamelCase			    =				(out_a - labels).mean()
           loss_a.backward()

           # compare the output and parameters gradients
           self.assertTrue((loss - loss_a).abs() < 1E-5						)
           _lowerCamelCase			    =				dict(model.named_parameters()						)
           _lowerCamelCase			    =				dict(model_a.named_parameters()						)
           for name, param in named_params.items():
                 self.assertTrue(torch_all_close(param.grad.data				,				named_params_a[name].grad.data				,				atol=5E-5						)						)



     def 		UpperCamelCase_						(					self						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase    ,					_lowerCamelCase			    =				AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''				,				output_loading_info=A_						)
           self.assertIsNotNone(A_						)
           self.assertEqual(len(loading_info['''missing_keys''']						)				,				0						)

           model.to(A_						)
           _lowerCamelCase			    =				model(**self.dummy_input						)

           assert image is not None, "Make sure output is not None"





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






           """simple docstring"""
           _lowerCamelCase			    =				AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''						)
           _lowerCamelCase			    =				model.to(A_						)
           model.eval()

           if torch_device == "mps":
                 _lowerCamelCase			    =				torch.manual_seed(0						)
           else:
                 _lowerCamelCase			    =				torch.Generator(device=A_						).manual_seed(0						)

           _lowerCamelCase			    =				torch.randn(
               1				,				model.config.in_channels				,				model.config.sample_size				,				model.config.sample_size				,				generator=torch.manual_seed(0						)				,				)
           _lowerCamelCase			    =				image.to(A_						)
           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_				,				sample_posterior=A_				,				generator=A_						).sample

           _lowerCamelCase			    =				output[0, -1, -3:, -3:].flatten().cpu()

           # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
           # the expected output slices are not the same for CPU and GPU.
           if torch_device == "mps":
                 _lowerCamelCase			    =				torch.tensor(
                     [
                         -4.0_0_7_8E-0_1,
                         -3.8_3_2_3E-0_4,
                         -1.2_6_8_1E-0_1,
                         -1.1_4_6_2E-0_1,
                         2.0_0_9_5E-0_1,
                         1.0_8_9_3E-0_1,
                         -8.8_2_4_7E-0_2,
                         -3.0_3_6_1E-0_1,
                         -9.8_6_4_4E-0_3,
                     ]						)
           elif torch_device == "cpu":
                 _lowerCamelCase			    =				torch.tensor(
                     [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]						)
           else:
                 _lowerCamelCase			    =				torch.tensor(
                     [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]						)

           self.assertTrue(torch_all_close(A_				,				A_				,				rtol=1E-2						)						)





@slow
class  UpperCamelCase				( unittest.TestCase ):

     '''simple docstring'''

     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				Optional[int]:






           """simple docstring"""
           return F'gaussian_noise_s={seed}_shape={"_".join([str(A_						) for s in shape]						)}.npy'



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






           """simple docstring"""
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



     def 		UpperCamelCase_						(					self				,				A_=0				,				A_=(4, 3, 5_12, 5_12)				,				A_=False						) ->				Optional[int]:






           """simple docstring"""
           _lowerCamelCase			    =				torch.floataa if fpaa else torch.floataa
           _lowerCamelCase			    =				torch.from_numpy(load_hf_numpy(self.get_file_format(A_				,				A_						)						)						).to(A_						).to(A_						)
           return image



     def 		UpperCamelCase_						(					self				,				A_="CompVis/stable-diffusion-v1-4"				,				A_=False						) ->				Any:






           """simple docstring"""
           _lowerCamelCase			    =				'''fp16''' if fpaa else None
           _lowerCamelCase			    =				torch.floataa if fpaa else torch.floataa

           _lowerCamelCase			    =				AutoencoderKL.from_pretrained(
               A_				,				subfolder='''vae'''				,				torch_dtype=A_				,				revision=A_				,				)
           model.to(A_						).eval()

           return model



     def 		UpperCamelCase_						(					self				,				A_=0						) ->				int:






           """simple docstring"""
           if torch_device == "mps":
                 return torch.manual_seed(A_						)
           return torch.Generator(device=A_						).manual_seed(A_						)



     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
             [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
             # fmt: on
         ]						)
     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model()
           _lowerCamelCase			    =				self.get_sd_image(A_						)
           _lowerCamelCase			    =				self.get_generator(A_						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_				,				generator=A_				,				sample_posterior=A_						).sample

           assert sample.shape == image.shape

           _lowerCamelCase			    =				sample[-1, -2:, -2:, :2].flatten().float().cpu()
           _lowerCamelCase			    =				torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice						)

           assert torch_all_close(A_				,				A_				,				atol=3E-3						)



     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
             [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
             # fmt: on
         ]						)
     @require_torch_gpu
     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				List[str]:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model(fpaa=A_						)
           _lowerCamelCase			    =				self.get_sd_image(A_				,				fpaa=A_						)
           _lowerCamelCase			    =				self.get_generator(A_						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_				,				generator=A_				,				sample_posterior=A_						).sample

           assert sample.shape == image.shape

           _lowerCamelCase			    =				sample[-1, -2:, :2, -2:].flatten().float().cpu()
           _lowerCamelCase			    =				torch.tensor(A_						)

           assert torch_all_close(A_				,				A_				,				atol=1E-2						)



     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
             [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
             # fmt: on
         ]						)
     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model()
           _lowerCamelCase			    =				self.get_sd_image(A_						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model(A_						).sample

           assert sample.shape == image.shape

           _lowerCamelCase			    =				sample[-1, -2:, -2:, :2].flatten().float().cpu()
           _lowerCamelCase			    =				torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice						)

           assert torch_all_close(A_				,				A_				,				atol=3E-3						)



     @parameterized.expand(
         [
             # fmt: off
             [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
             [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
             # fmt: on
         ]						)
     @require_torch_gpu
     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model()
           _lowerCamelCase			    =				self.get_sd_image(A_				,				shape=(3, 4, 64, 64)						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model.decode(A_						).sample

           assert list(sample.shape						) == [3, 3, 5_12, 5_12]

           _lowerCamelCase			    =				sample[-1, -2:, :2, -2:].flatten().cpu()
           _lowerCamelCase			    =				torch.tensor(A_						)

           assert torch_all_close(A_				,				A_				,				atol=1E-3						)



     @parameterized.expand(
         [
             # fmt: off
             [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
             [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
             # fmt: on
         ]						)
     @require_torch_gpu
     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				List[str]:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model(fpaa=A_						)
           _lowerCamelCase			    =				self.get_sd_image(A_				,				shape=(3, 4, 64, 64)				,				fpaa=A_						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model.decode(A_						).sample

           assert list(sample.shape						) == [3, 3, 5_12, 5_12]

           _lowerCamelCase			    =				sample[-1, -2:, :2, -2:].flatten().float().cpu()
           _lowerCamelCase			    =				torch.tensor(A_						)

           assert torch_all_close(A_				,				A_				,				atol=5E-3						)



     @parameterized.expand([(13,), (16,), (27,)]						)
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available()				,				reason='''xformers is not required when using PyTorch 2.0.'''						)
     def 		UpperCamelCase_						(					self				,				A_						) ->				str:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model(fpaa=A_						)
           _lowerCamelCase			    =				self.get_sd_image(A_				,				shape=(3, 4, 64, 64)				,				fpaa=A_						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model.decode(A_						).sample

           model.enable_xformers_memory_efficient_attention()
           with torch.no_grad():
                 _lowerCamelCase			    =				model.decode(A_						).sample

           assert list(sample.shape						) == [3, 3, 5_12, 5_12]

           assert torch_all_close(A_				,				A_				,				atol=1E-1						)



     @parameterized.expand([(13,), (16,), (37,)]						)
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available()				,				reason='''xformers is not required when using PyTorch 2.0.'''						)
     def 		UpperCamelCase_						(					self				,				A_						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model()
           _lowerCamelCase			    =				self.get_sd_image(A_				,				shape=(3, 4, 64, 64)						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model.decode(A_						).sample

           model.enable_xformers_memory_efficient_attention()
           with torch.no_grad():
                 _lowerCamelCase			    =				model.decode(A_						).sample

           assert list(sample.shape						) == [3, 3, 5_12, 5_12]

           assert torch_all_close(A_				,				A_				,				atol=1E-2						)





     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
             [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
             # fmt: on
         ]						)
     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				Optional[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				self.get_sd_vae_model()
           _lowerCamelCase			    =				self.get_sd_image(A_						)
           _lowerCamelCase			    =				self.get_generator(A_						)

           with torch.no_grad():
                 _lowerCamelCase			    =				model.encode(A_						).latent_dist
                 _lowerCamelCase			    =				dist.sample(generator=A_						)

           assert list(sample.shape						) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

           _lowerCamelCase			    =				sample[0, -1, -3:, -3:].flatten().cpu()
           _lowerCamelCase			    =				torch.tensor(A_						)

           _lowerCamelCase			    =				3E-3 if torch_device != '''mps''' else 1E-2
           assert torch_all_close(A_				,				A_				,				atol=A_						)
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
snake_case__				= logging.get_logger(__name__)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase=False		)						->    List[Any]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      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'''),
                ]		)

            # if just the base model, we should remove "vit" from all keys that start with "vit"
            _lowerCamelCase			    =				[(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       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=False		)						->    str:



      '''simple docstring'''

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



      '''simple docstring'''

      _lowerCamelCase			    =				['''head.weight''', '''head.bias''']
      for k in ignore_keys:
            state_dict.pop(__UpperCAmelCase ,						__UpperCAmelCase		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    List[Any]:



      '''simple docstring'''

      _lowerCamelCase			    =				dct.pop(__UpperCAmelCase		)
      _lowerCamelCase			    =				val
def       __magic_name__(  )						->    List[str]:



      '''simple docstring'''

      _lowerCamelCase			    =				'''http://images.cocodataset.org/val2017/000000039769.jpg'''
      _lowerCamelCase			    =				Image.open(requests.get(__UpperCAmelCase ,						stream=__UpperCAmelCase		).raw		)
      return im
@torch.no_grad()
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=True		)						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				ViTConfig()
      # patch_size
      if model_name[-1] == "8":
            _lowerCamelCase			    =				8
      # set labels if required
      if not base_model:
            _lowerCamelCase			    =				1000
            _lowerCamelCase			    =				'''huggingface/label-files'''
            _lowerCamelCase			    =				'''imagenet-1k-id2label.json'''
            _lowerCamelCase			    =				json.load(open(hf_hub_download(__UpperCAmelCase ,						__UpperCAmelCase ,						repo_type='''dataset'''		) ,						'''r'''		)		)
            _lowerCamelCase			    =				{int(__UpperCAmelCase		): v for k, v in idalabel.items()}
            _lowerCamelCase			    =				idalabel
            _lowerCamelCase			    =				{v: k for k, v in idalabel.items()}
      # size of the architecture
      if model_name in ["dino_vits8", "dino_vits16"]:
            _lowerCamelCase			    =				384
            _lowerCamelCase			    =				1536
            _lowerCamelCase			    =				12
            _lowerCamelCase			    =				6

      # load original model from torch hub
      _lowerCamelCase			    =				torch.hub.load('''facebookresearch/dino:main''' ,						__UpperCAmelCase		)
      original_model.eval()

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

      # load HuggingFace model
      if base_model:
            _lowerCamelCase			    =				ViTModel(__UpperCAmelCase ,						add_pooling_layer=__UpperCAmelCase		).eval()
      else:
            _lowerCamelCase			    =				ViTForImageClassification(__UpperCAmelCase		).eval()
      model.load_state_dict(__UpperCAmelCase		)

      # Check outputs on an image, prepared by ViTImageProcessor
      _lowerCamelCase			    =				ViTImageProcessor()
      _lowerCamelCase			    =				image_processor(images=prepare_img() ,						return_tensors='''pt'''		)
      _lowerCamelCase			    =				encoding['''pixel_values''']
      _lowerCamelCase			    =				model(__UpperCAmelCase		)

      if base_model:
            _lowerCamelCase			    =				original_model(__UpperCAmelCase		)
            assert torch.allclose(__UpperCAmelCase ,						outputs.last_hidden_state[:, 0, :] ,						atol=1E-1		)
      else:
            _lowerCamelCase			    =				original_model(__UpperCAmelCase		)
            assert logits.shape == outputs.logits.shape
            assert torch.allclose(__UpperCAmelCase ,						outputs.logits ,						atol=1E-3		)

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


if __name__ == "__main__":
     snake_case__				= argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default='dino_vitb16',
         type=str,
         help='Name of the model trained with DINO you\'d like to convert.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )
     parser.add_argument(
         '--base_model',
         action='store_true',
         help='Whether to only convert the base model (no projection head weights).',
     )

     parser.set_defaults(base_model=True)
     snake_case__				= parser.parse_args()
     convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)