


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class     __UpperCAmelCase							(       __A ,					unittest.TestCase							):


     """simple docstring"""
     _lowerCamelCase   =      UnCLIPImageVariationPipeline
     _lowerCamelCase   =      IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
     _lowerCamelCase   =      IMAGE_VARIATION_BATCH_PARAMS

     _lowerCamelCase   =      [
         """generator""",
         """return_dict""",
         """decoder_num_inference_steps""",
         """super_res_num_inference_steps""",
     ]
     _lowerCamelCase   =      False


     @property
     def 	snake_case_  (	self    ):
        return 32


     @property
     def 	snake_case_  (	self    ):
        return 32


     @property
     def 	snake_case_  (	self    ):
        return self.time_input_dim


     @property
     def 	snake_case_  (	self    ):
        return self.time_input_dim * 4


     @property
     def 	snake_case_  (	self    ):
        return 100


     @property
     def 	snake_case_  (	self    ):
        __a					=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )
        return tokenizer


     @property
     def 	snake_case_  (	self    ):
        torch.manual_seed(0    )
        __a					=  CLIPTextConfig(
            bos_token_id=0  ,       eos_token_id=2  ,       hidden_size=self.text_embedder_hidden_size  ,       projection_dim=self.text_embedder_hidden_size  ,       intermediate_size=37  ,       layer_norm_eps=1E-05  ,       num_attention_heads=4  ,       num_hidden_layers=5  ,       pad_token_id=1  ,       vocab_size=1000  ,       )
        return CLIPTextModelWithProjection(__A    )


     @property
     def 	snake_case_  (	self    ):
        torch.manual_seed(0    )
        __a					=  CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size  ,       projection_dim=self.text_embedder_hidden_size  ,       num_hidden_layers=5  ,       num_attention_heads=4  ,       image_size=32  ,       intermediate_size=37  ,       patch_size=1  ,       )
        return CLIPVisionModelWithProjection(__A    )


     @property
     def 	snake_case_  (	self    ):
        torch.manual_seed(0    )

        __a					=  {
            """clip_embeddings_dim""": self.text_embedder_hidden_size,
            """time_embed_dim""": self.time_embed_dim,
            """cross_attention_dim""": self.cross_attention_dim,
        }

        __a					=  UnCLIPTextProjModel(**__A    )
        return model


     @property
     def 	snake_case_  (	self    ):
        torch.manual_seed(0    )

        __a					=  {
            """sample_size""": 32,
            # RGB in channels
            """in_channels""": 3,
            # Out channels is double in channels because predicts mean and variance
            """out_channels""": 6,
            """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
            """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
            """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
            """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
            """layers_per_block""": 1,
            """cross_attention_dim""": self.cross_attention_dim,
            """attention_head_dim""": 4,
            """resnet_time_scale_shift""": """scale_shift""",
            """class_embed_type""": """identity""",
        }

        __a					=  UNetaDConditionModel(**__A    )
        return model


     @property
     def 	snake_case_  (	self    ):
        return {
            "sample_size": 64,
            "layers_per_block": 1,
            "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
            "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "in_channels": 6,
            "out_channels": 3,
        }


     @property
     def 	snake_case_  (	self    ):
        torch.manual_seed(0    )

        __a					=  UNetaDModel(**self.dummy_super_res_kwargs    )
        return model


     @property
     def 	snake_case_  (	self    ):
        # seeded differently to get different unet than `self.dummy_super_res_first`
        torch.manual_seed(1    )

        __a					=  UNetaDModel(**self.dummy_super_res_kwargs    )
        return model


     def 	snake_case_  (	self    ):
        __a					=  self.dummy_decoder
        __a					=  self.dummy_text_proj
        __a					=  self.dummy_text_encoder
        __a					=  self.dummy_tokenizer
        __a					=  self.dummy_super_res_first
        __a					=  self.dummy_super_res_last

        __a					=  UnCLIPScheduler(
            variance_type="""learned_range"""  ,       prediction_type="""epsilon"""  ,       num_train_timesteps=1000  ,       )

        __a					=  UnCLIPScheduler(
            variance_type="""fixed_small_log"""  ,       prediction_type="""epsilon"""  ,       num_train_timesteps=1000  ,       )

        __a					=  CLIPImageProcessor(crop_size=32  ,       size=32    )

        __a					=  self.dummy_image_encoder

        return {
            "decoder": decoder,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "text_proj": text_proj,
            "feature_extractor": feature_extractor,
            "image_encoder": image_encoder,
            "super_res_first": super_res_first,
            "super_res_last": super_res_last,
            "decoder_scheduler": decoder_scheduler,
            "super_res_scheduler": super_res_scheduler,
        }


     def 	snake_case_  (	self  ,       __A  ,       __A=0  ,       __A=True    ):
        __a					=  floats_tensor((1, 3, 32, 32)  ,       rng=random.Random(__A    )    ).to(__A    )
        if str(__A    ).startswith("""mps"""    ):
           __a					=  torch.manual_seed(__A    )
        else:
           __a					=  torch.Generator(device=__A    ).manual_seed(__A    )

        if pil_image:
           __a					=  input_image * 0.5 + 0.5
           __a					=  input_image.clamp(0  ,       1    )
           __a					=  input_image.cpu().permute(0  ,       2  ,       3  ,       1    ).float().numpy()
           __a					=  DiffusionPipeline.numpy_to_pil(__A    )[0]

        return {
            "image": input_image,
            "generator": generator,
            "decoder_num_inference_steps": 2,
            "super_res_num_inference_steps": 2,
            "output_type": "np",
        }


     def 	snake_case_  (	self    ):
        __a					=  """cpu"""

        __a					=  self.get_dummy_components()

        __a					=  self.pipeline_class(**__A    )
        __a					=  pipe.to(__A    )

        pipe.set_progress_bar_config(disable=__A    )

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )

        __a					=  pipe(**__A    )
        __a					=  output.images

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )

        __a					=  pipe(
            **__A  ,       return_dict=__A  ,       )[0]

        __a					=  image[0, -3:, -3:, -1]
        __a					=  image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        __a					=  np.array(
            [
                0.9997,
                0.0002,
                0.9997,
                0.9997,
                0.9969,
                0.0023,
                0.9997,
                0.9969,
                0.9970,
            ]    )

        assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1E-2


     def 	snake_case_  (	self    ):
        __a					=  """cpu"""

        __a					=  self.get_dummy_components()

        __a					=  self.pipeline_class(**__A    )
        __a					=  pipe.to(__A    )

        pipe.set_progress_bar_config(disable=__A    )

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )

        __a					=  pipe(**__A    )
        __a					=  output.images

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )

        __a					=  pipe(
            **__A  ,       return_dict=__A  ,       )[0]

        __a					=  image[0, -3:, -3:, -1]
        __a					=  image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        __a					=  np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]    )

        assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1E-2


     def 	snake_case_  (	self    ):
        __a					=  """cpu"""

        __a					=  self.get_dummy_components()

        __a					=  self.pipeline_class(**__A    )
        __a					=  pipe.to(__A    )

        pipe.set_progress_bar_config(disable=__A    )

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )
        __a					=  [
            pipeline_inputs["""image"""],
            pipeline_inputs["""image"""],
        ]

        __a					=  pipe(**__A    )
        __a					=  output.images

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )
        __a					=  [
            tuple_pipeline_inputs["""image"""],
            tuple_pipeline_inputs["""image"""],
        ]

        __a					=  pipe(
            **__A  ,       return_dict=__A  ,       )[0]

        __a					=  image[0, -3:, -3:, -1]
        __a					=  image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (2, 64, 64, 3)

        __a					=  np.array(
            [
                0.9997,
                0.9989,
                0.0008,
                0.0021,
                0.9960,
                0.0018,
                0.0014,
                0.0002,
                0.9933,
            ]    )

        assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1E-2


     def 	snake_case_  (	self    ):
        __a					=  torch.device("""cpu"""    )


        class     __UpperCAmelCase							:


             """simple docstring"""
             _lowerCamelCase   =      1







        __a					=  self.get_dummy_components()

        __a					=  self.pipeline_class(**__A    )
        __a					=  pipe.to(__A    )

        pipe.set_progress_bar_config(disable=__A    )

        __a					=  torch.Generator(device=__A    ).manual_seed(0    )
        __a					=  pipe.decoder.dtype
        __a					=  1

        __a					=  (
            batch_size,
            pipe.decoder.config.in_channels,
            pipe.decoder.config.sample_size,
            pipe.decoder.config.sample_size,
        )
        __a					=  pipe.prepare_latents(
            __A  ,       dtype=__A  ,       device=__A  ,       generator=__A  ,       latents=__A  ,       scheduler=DummyScheduler()    )

        __a					=  (
            batch_size,
            pipe.super_res_first.config.in_channels // 2,
            pipe.super_res_first.config.sample_size,
            pipe.super_res_first.config.sample_size,
        )
        __a					=  pipe.prepare_latents(
            __A  ,       dtype=__A  ,       device=__A  ,       generator=__A  ,       latents=__A  ,       scheduler=DummyScheduler()    )

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )

        __a					=  pipe(
            **__A  ,       decoder_latents=__A  ,       super_res_latents=__A    ).images

        __a					=  self.get_dummy_inputs(__A  ,       pil_image=__A    )
        # Don't pass image, instead pass embedding
        __a					=  pipeline_inputs.pop("""image"""    )
        __a					=  pipe.image_encoder(__A    ).image_embeds

        __a					=  pipe(
            **__A  ,       decoder_latents=__A  ,       super_res_latents=__A  ,       image_embeddings=__A  ,       ).images

        # make sure passing text embeddings manually is identical
        assert np.abs(img_out_a - img_out_a    ).max() < 1E-4


     @skip_mps
     def 	snake_case_  (	self    ):
        __a					=  torch_device == """cpu"""

        # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
        __a					=  1E-2

        self._test_attention_slicing_forward_pass(
            test_max_difference=__A  ,       expected_max_diff=__A    )


     @skip_mps
     def 	snake_case_  (	self    ):
        __a					=  torch_device == """cpu"""
        __a					=  True
        __a					=  [
            """decoder_num_inference_steps""",
            """super_res_num_inference_steps""",
        ]

        self._test_inference_batch_single_identical(
            test_max_difference=__A  ,       relax_max_difference=__A  ,       additional_params_copy_to_batched_inputs=__A  ,       )


     def 	snake_case_  (	self    ):
        __a					=  [
            """decoder_num_inference_steps""",
            """super_res_num_inference_steps""",
        ]

        if torch_device == "mps":
           # TODO: MPS errors with larger batch sizes
           __a					=  [2, 3]
           self._test_inference_batch_consistent(
               batch_sizes=__A  ,       additional_params_copy_to_batched_inputs=__A  ,       )
        else:
           self._test_inference_batch_consistent(
               additional_params_copy_to_batched_inputs=__A    )


     @skip_mps
     def 	snake_case_  (	self    ):
        return super().test_dict_tuple_outputs_equivalent()


     @skip_mps
     def 	snake_case_  (	self    ):
        return super().test_save_load_local()


     @skip_mps
     def 	snake_case_  (	self    ):
        return super().test_save_load_optional_components()







@slow
@require_torch_gpu
class     __UpperCAmelCase							(       unittest.TestCase							):


     """simple docstring"""


     def 	snake_case_  (	self    ):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()


     def 	snake_case_  (	self    ):
        __a					=  load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"""    )
        __a					=  load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/unclip/karlo_v1_alpha_cat_variation_fp16.npy"""    )

        __a					=  UnCLIPImageVariationPipeline.from_pretrained(
            """kakaobrain/karlo-v1-alpha-image-variations"""  ,       torch_dtype=torch.floataa    )
        __a					=  pipeline.to(__A    )
        pipeline.set_progress_bar_config(disable=__A    )

        __a					=  torch.Generator(device="""cpu"""    ).manual_seed(0    )
        __a					=  pipeline(
            __A  ,       generator=__A  ,       output_type="""np"""  ,       )

        __a					=  output.images[0]

        assert image.shape == (256, 256, 3)

        assert_mean_pixel_difference(__A  ,       __A  ,       15    )







import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)




def  a   (lowerCAmelCase__   ):
   __a					=  VideoMAEConfig()

   set_architecture_configs(lowerCAmelCase__		,				lowerCAmelCase__   )

   if "finetuned" not in model_name:
      __a					=  False

   if "finetuned" in model_name:
      __a					=  """huggingface/label-files"""
      if "kinetics" in model_name:
         __a					=  400
         __a					=  """kinetics400-id2label.json"""
      elif "ssv2" in model_name:
         __a					=  174
         __a					=  """something-something-v2-id2label.json"""
      else:
         raise ValueError("""Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."""   )
      __a					=  json.load(open(hf_hub_download(lowerCAmelCase__		,				lowerCAmelCase__		,				repo_type="""dataset"""   )		,				"""r"""   )   )
      __a					=  {int(lowerCAmelCase__   ): v for k, v in idalabel.items()}
      __a					=  idalabel
      __a					=  {v: k for k, v in idalabel.items()}

   return config




def  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):
   if "small" in model_name:
      __a					=  384
      __a					=  1_536
      __a					=  12
      __a					=  16
      __a					=  12
      __a					=  3
      __a					=  192
      __a					=  768
   elif "large" in model_name:
      __a					=  1_024
      __a					=  4_096
      __a					=  24
      __a					=  16
      __a					=  12
      __a					=  8
      __a					=  512
      __a					=  2_048
   elif "huge" in model_name:
      __a					=  1_280
      __a					=  5_120
      __a					=  32
      __a					=  16
      __a					=  12
      __a					=  8
      __a					=  640
      __a					=  2_560
   elif "base" not in model_name:
      raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""   )




def  a   (lowerCAmelCase__   ):
   if "encoder." in name:
      __a					=  name.replace("""encoder."""		,				""""""   )
   if "cls_token" in name:
      __a					=  name.replace("""cls_token"""		,				"""videomae.embeddings.cls_token"""   )
   if "decoder_pos_embed" in name:
      __a					=  name.replace("""decoder_pos_embed"""		,				"""decoder.decoder_pos_embed"""   )
   if "pos_embed" in name and "decoder" not in name:
      __a					=  name.replace("""pos_embed"""		,				"""videomae.embeddings.position_embeddings"""   )
   if "patch_embed.proj" in name:
      __a					=  name.replace("""patch_embed.proj"""		,				"""videomae.embeddings.patch_embeddings.projection"""   )
   if "patch_embed.norm" in name:
      __a					=  name.replace("""patch_embed.norm"""		,				"""videomae.embeddings.norm"""   )
   if "decoder.blocks" in name:
      __a					=  name.replace("""decoder.blocks"""		,				"""decoder.decoder_layers"""   )
   if "blocks" in name:
      __a					=  name.replace("""blocks"""		,				"""videomae.encoder.layer"""   )
   if "attn.proj" in name:
      __a					=  name.replace("""attn.proj"""		,				"""attention.output.dense"""   )
   if "attn" in name and "bias" not in name:
      __a					=  name.replace("""attn"""		,				"""attention.self"""   )
   if "attn" in name:
      __a					=  name.replace("""attn"""		,				"""attention.attention"""   )
   if "norm1" in name:
      __a					=  name.replace("""norm1"""		,				"""layernorm_before"""   )
   if "norm2" in name:
      __a					=  name.replace("""norm2"""		,				"""layernorm_after"""   )
   if "mlp.fc1" in name:
      __a					=  name.replace("""mlp.fc1"""		,				"""intermediate.dense"""   )
   if "mlp.fc2" in name:
      __a					=  name.replace("""mlp.fc2"""		,				"""output.dense"""   )
   if "decoder_embed" in name:
      __a					=  name.replace("""decoder_embed"""		,				"""decoder.decoder_embed"""   )
   if "decoder_norm" in name:
      __a					=  name.replace("""decoder_norm"""		,				"""decoder.decoder_norm"""   )
   if "decoder_pred" in name:
      __a					=  name.replace("""decoder_pred"""		,				"""decoder.decoder_pred"""   )
   if "norm.weight" in name and "decoder" not in name and "fc" not in name:
      __a					=  name.replace("""norm.weight"""		,				"""videomae.layernorm.weight"""   )
   if "norm.bias" in name and "decoder" not in name and "fc" not in name:
      __a					=  name.replace("""norm.bias"""		,				"""videomae.layernorm.bias"""   )
   if "head" in name and "decoder" not in name:
      __a					=  name.replace("""head"""		,				"""classifier"""   )

   return name




def  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):
   for key in orig_state_dict.copy().keys():
      __a					=  orig_state_dict.pop(lowerCAmelCase__   )

      if key.startswith("""encoder."""   ):
         __a					=  key.replace("""encoder."""		,				""""""   )

      if "qkv" in key:
         __a					=  key.split("""."""   )
         if key.startswith("""decoder.blocks"""   ):
            __a					=  config.decoder_hidden_size
            __a					=  int(key_split[2]   )
            __a					=  """decoder.decoder_layers."""
            if "weight" in key:
               __a					=  val[:dim, :]
               __a					=  val[dim : dim * 2, :]
               __a					=  val[-dim:, :]
         else:
            __a					=  config.hidden_size
            __a					=  int(key_split[1]   )
            __a					=  """videomae.encoder.layer."""
            if "weight" in key:
               __a					=  val[:dim, :]
               __a					=  val[dim : dim * 2, :]
               __a					=  val[-dim:, :]
      else:
         __a					=  val

   return orig_state_dict




def  a   ():
   __a					=  hf_hub_download(
       repo_id="""hf-internal-testing/spaghetti-video"""		,				filename="""eating_spaghetti.npy"""		,				repo_type="""dataset"""   )
   __a					=  np.load(lowerCAmelCase__   )
   return list(lowerCAmelCase__   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   ):
   __a					=  get_videomae_config(lowerCAmelCase__   )

   if "finetuned" in model_name:
      __a					=  VideoMAEForVideoClassification(lowerCAmelCase__   )
   else:
      __a					=  VideoMAEForPreTraining(lowerCAmelCase__   )

   # download original checkpoint, hosted on Google Drive
   __a					=  """pytorch_model.bin"""
   gdown.cached_download(lowerCAmelCase__		,				lowerCAmelCase__		,				quiet=lowerCAmelCase__   )
   __a					=  torch.load(lowerCAmelCase__		,				map_location="""cpu"""   )
   if "model" in files:
      __a					=  files["""model"""]
   else:
      __a					=  files["""module"""]
   __a					=  convert_state_dict(lowerCAmelCase__		,				lowerCAmelCase__   )

   model.load_state_dict(lowerCAmelCase__   )
   model.eval()

   # verify model on basic input
   __a					=  VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]		,				image_std=[0.5, 0.5, 0.5]   )
   __a					=  prepare_video()
   __a					=  image_processor(lowerCAmelCase__		,				return_tensors="""pt"""   )

   if "finetuned" not in model_name:
      __a					=  hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""		,				filename="""bool_masked_pos.pt"""   )
      __a					=  torch.load(lowerCAmelCase__   )

   __a					=  model(**lowerCAmelCase__   )
   __a					=  outputs.logits

   __a					=  [
       """videomae-small-finetuned-kinetics""",
       """videomae-small-finetuned-ssv2""",
       # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
       """videomae-base-short""",
       """videomae-base-short-finetuned-kinetics""",
       """videomae-base""",
       """videomae-base-finetuned-kinetics""",
       """videomae-large""",
       """videomae-large-finetuned-kinetics""",
       """videomae-huge-finetuned-kinetics""",
       # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
       """videomae-base-short-ssv2""",
       """videomae-base-short-finetuned-ssv2""",
       """videomae-base-ssv2""",
       """videomae-base-finetuned-ssv2""",
   ]

   # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
   if model_name == "videomae-small-finetuned-kinetics":
      __a					=  torch.Size([1, 400]   )
      __a					=  torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]   )
   elif model_name == "videomae-small-finetuned-ssv2":
      __a					=  torch.Size([1, 174]   )
      __a					=  torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]   )
   elif model_name == "videomae-base":
      __a					=  torch.Size([1, 1_408, 1_536]   )
      __a					=  torch.tensor([[0.7_7_3_9, 0.7_9_6_8, 0.7_0_8_9], [0.6_7_0_1, 0.7_4_8_7, 0.6_2_0_9], [0.4_2_8_7, 0.5_1_5_8, 0.4_7_7_3]]   )
   elif model_name == "videomae-base-short":
      __a					=  torch.Size([1, 1_408, 1_536]   )
      __a					=  torch.tensor([[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]   )
      # we verified the loss both for normalized and unnormalized targets for this one
      __a					=  torch.tensor([0.5_1_4_2]   ) if config.norm_pix_loss else torch.tensor([0.6_4_6_9]   )
   elif model_name == "videomae-large":
      __a					=  torch.Size([1, 1_408, 1_536]   )
      __a					=  torch.tensor([[0.7_1_4_9, 0.7_9_9_7, 0.6_9_6_6], [0.6_7_6_8, 0.7_8_6_9, 0.6_9_4_8], [0.5_1_3_9, 0.6_2_2_1, 0.5_6_0_5]]   )
   elif model_name == "videomae-large-finetuned-kinetics":
      __a					=  torch.Size([1, 400]   )
      __a					=  torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]   )
   elif model_name == "videomae-huge-finetuned-kinetics":
      __a					=  torch.Size([1, 400]   )
      __a					=  torch.tensor([0.2_4_3_3, 0.1_6_3_2, -0.4_8_9_4]   )
   elif model_name == "videomae-base-short-finetuned-kinetics":
      __a					=  torch.Size([1, 400]   )
      __a					=  torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]   )
   elif model_name == "videomae-base-finetuned-kinetics":
      __a					=  torch.Size([1, 400]   )
      __a					=  torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]   )
   elif model_name == "videomae-base-short-ssv2":
      __a					=  torch.Size([1, 1_408, 1_536]   )
      __a					=  torch.tensor([[0.4_7_1_2, 0.5_2_9_6, 0.5_7_8_6], [0.2_2_7_8, 0.2_7_2_9, 0.4_0_2_6], [0.0_3_5_2, 0.0_7_3_0, 0.2_5_0_6]]   )
   elif model_name == "videomae-base-short-finetuned-ssv2":
      __a					=  torch.Size([1, 174]   )
      __a					=  torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]   )
   elif model_name == "videomae-base-ssv2":
      __a					=  torch.Size([1, 1_408, 1_536]   )
      __a					=  torch.tensor([[0.8_1_3_1, 0.8_7_2_7, 0.8_5_4_6], [0.7_3_6_6, 0.9_3_7_7, 0.8_8_7_0], [0.5_9_3_5, 0.8_8_7_4, 0.8_5_6_4]]   )
   elif model_name == "videomae-base-finetuned-ssv2":
      __a					=  torch.Size([1, 174]   )
      __a					=  torch.tensor([0.1_9_6_1, -0.8_3_3_7, -0.6_3_8_9]   )
   else:
      raise ValueError(f'''Model name not supported. Should be one of {model_names}'''   )

   # verify logits
   assert logits.shape == expected_shape
   if "finetuned" in model_name:
      assert torch.allclose(logits[0, :3]		,				lowerCAmelCase__		,				atol=1E-4   )
   else:
      print("""Logits:"""		,				logits[0, :3, :3]   )
      assert torch.allclose(logits[0, :3, :3]		,				lowerCAmelCase__		,				atol=1E-4   )
   print("""Logits ok!"""   )

   # verify loss, if applicable
   if model_name == "videomae-base-short":
      __a					=  outputs.loss
      assert torch.allclose(lowerCAmelCase__		,				lowerCAmelCase__		,				atol=1E-4   )
      print("""Loss ok!"""   )

   if pytorch_dump_folder_path is not None:
      print(f'''Saving model and image processor to {pytorch_dump_folder_path}'''   )
      image_processor.save_pretrained(lowerCAmelCase__   )
      model.save_pretrained(lowerCAmelCase__   )

   if push_to_hub:
      print("""Pushing to the hub..."""   )
      model.push_to_hub(lowerCAmelCase__		,				organization="""nielsr"""   )


if __name__ == "__main__":
  SCREAMING_SNAKE_CASE							    =							argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--checkpoint_url',
      default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
      type=str,
      help=(
          'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
          ' download link.'
      ),
  )
  parser.add_argument(
      '--pytorch_dump_folder_path',
      default='/Users/nielsrogge/Documents/VideoMAE/Test',
      type=str,
      help='Path to the output PyTorch model directory.',
  )
  parser.add_argument('--model_name', default='videomae-base', type=str, help='Name of the model.')
  parser.add_argument(
      '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
  )

  SCREAMING_SNAKE_CASE							    =							parser.parse_args()
  convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)



