





"""simple docstring"""
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class lowerCamelCase__  (    __UpperCamelCase  ,				unittest.TestCase				):
   __UpperCAmelCase        =					UnCLIPPipeline
   __UpperCAmelCase        =					TEXT_TO_IMAGE_PARAMS - {
       """negative_prompt""",
       """height""",
       """width""",
       """negative_prompt_embeds""",
       """guidance_scale""",
       """prompt_embeds""",
       """cross_attention_kwargs""",
   }
   __UpperCAmelCase        =					TEXT_TO_IMAGE_BATCH_PARAMS
   __UpperCAmelCase        =					[
       """generator""",
       """return_dict""",
       """prior_num_inference_steps""",
       """decoder_num_inference_steps""",
       """super_res_num_inference_steps""",
   ]
   __UpperCAmelCase        =					False



   @property
   def    _UpperCAmelCase	(       self	) ->    Union[str, Any]:







      """simple docstring"""
      return 3_2



   @property
   def    _UpperCAmelCase	(       self	) ->    Any:







      """simple docstring"""
      return 3_2



   @property
   def    _UpperCAmelCase	(       self	) ->    Dict:







      """simple docstring"""
      return self.time_input_dim



   @property
   def    _UpperCAmelCase	(       self	) ->    Any:







      """simple docstring"""
      return self.time_input_dim * 4



   @property
   def    _UpperCAmelCase	(       self	) ->    List[str]:







      """simple docstring"""
      return 1_0_0



   @property
   def    _UpperCAmelCase	(       self	) ->    Any:







      """simple docstring"""
      lowercase    :			Any  =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)
      return tokenizer



   @property
   def    _UpperCAmelCase	(       self	) ->    Any:







      """simple docstring"""
      torch.manual_seed(0	)
      lowercase    :			str  =       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=3_7      ,   layer_norm_eps=1E-05      ,   num_attention_heads=4      ,   num_hidden_layers=5      ,   pad_token_id=1      ,   vocab_size=1_0_0_0      ,   )
      return CLIPTextModelWithProjection(snake_case	)



   @property
   def    _UpperCAmelCase	(       self	) ->    List[str]:







      """simple docstring"""
      torch.manual_seed(0	)

      lowercase    :			Dict  =       {
          """num_attention_heads""": 2,
          """attention_head_dim""": 1_2,
          """embedding_dim""": self.text_embedder_hidden_size,
          """num_layers""": 1,
      }

      lowercase    :			str  =       PriorTransformer(**snake_case	)
      return model



   @property
   def    _UpperCAmelCase	(       self	) ->    str:







      """simple docstring"""
      torch.manual_seed(0	)

      lowercase    :			List[Any]  =       {
          """clip_embeddings_dim""": self.text_embedder_hidden_size,
          """time_embed_dim""": self.time_embed_dim,
          """cross_attention_dim""": self.cross_attention_dim,
      }

      lowercase    :			Union[str, Any]  =       UnCLIPTextProjModel(**snake_case	)
      return model



   @property
   def    _UpperCAmelCase	(       self	) ->    Tuple:







      """simple docstring"""
      torch.manual_seed(0	)

      lowercase    :			Any  =       {
          """sample_size""": 3_2,
          # 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""",
      }

      lowercase    :			Optional[Any]  =       UNetaDConditionModel(**snake_case	)
      return model



   @property
   def    _UpperCAmelCase	(       self	) ->    Dict:







      """simple docstring"""
      return {
          "sample_size": 6_4,
          "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    _UpperCAmelCase	(       self	) ->    Tuple:







      """simple docstring"""
      torch.manual_seed(0	)

      lowercase    :			str  =       UNetaDModel(**self.dummy_super_res_kwargs	)
      return model



   @property
   def    _UpperCAmelCase	(       self	) ->    Optional[Any]:







      """simple docstring"""
      # seeded differently to get different unet than `self.dummy_super_res_first`
      torch.manual_seed(1	)

      lowercase    :			List[str]  =       UNetaDModel(**self.dummy_super_res_kwargs	)
      return model



   def    _UpperCAmelCase	(       self	) ->    Any:







      """simple docstring"""
      lowercase    :			str  =       self.dummy_prior
      lowercase    :			str  =       self.dummy_decoder
      lowercase    :			Optional[int]  =       self.dummy_text_proj
      lowercase    :			List[Any]  =       self.dummy_text_encoder
      lowercase    :			Optional[Any]  =       self.dummy_tokenizer
      lowercase    :			Any  =       self.dummy_super_res_first
      lowercase    :			List[Any]  =       self.dummy_super_res_last

      lowercase    :			List[str]  =       UnCLIPScheduler(
          variance_type="""fixed_small_log"""      ,   prediction_type="""sample"""      ,   num_train_timesteps=1_0_0_0      ,   clip_sample_range=5.0      ,   )

      lowercase    :			int  =       UnCLIPScheduler(
          variance_type="""learned_range"""      ,   prediction_type="""epsilon"""      ,   num_train_timesteps=1_0_0_0      ,   )

      lowercase    :			int  =       UnCLIPScheduler(
          variance_type="""fixed_small_log"""      ,   prediction_type="""epsilon"""      ,   num_train_timesteps=1_0_0_0      ,   )

      lowercase    :			Union[str, Any]  =       {
          """prior""": prior,
          """decoder""": decoder,
          """text_proj""": text_proj,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """super_res_first""": super_res_first,
          """super_res_last""": super_res_last,
          """prior_scheduler""": prior_scheduler,
          """decoder_scheduler""": decoder_scheduler,
          """super_res_scheduler""": super_res_scheduler,
      }

      return components



   def    _UpperCAmelCase	(       self      ,   snake_case      ,   snake_case=0	) ->    Any:







      """simple docstring"""
      if str(snake_case	).startswith("""mps"""	):
         lowercase    :			Union[str, Any]  =       torch.manual_seed(snake_case	)
      else:
         lowercase    :			str  =       torch.Generator(device=snake_case	).manual_seed(snake_case	)
      lowercase    :			Union[str, Any]  =       {
          """prompt""": """horse""",
          """generator""": generator,
          """prior_num_inference_steps""": 2,
          """decoder_num_inference_steps""": 2,
          """super_res_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs



   def    _UpperCAmelCase	(       self	) ->    Any:







      """simple docstring"""
      lowercase    :			Any  =       """cpu"""

      lowercase    :			List[str]  =       self.get_dummy_components()

      lowercase    :			str  =       self.pipeline_class(**snake_case	)
      lowercase    :			Optional[int]  =       pipe.to(snake_case	)

      pipe.set_progress_bar_config(disable=snake_case	)

      lowercase    :			Any  =       pipe(**self.get_dummy_inputs(snake_case	)	)
      lowercase    :			Any  =       output.images

      lowercase    :			Optional[Any]  =       pipe(
          **self.get_dummy_inputs(snake_case	)      ,   return_dict=snake_case      ,   )[0]

      lowercase    :			Tuple  =       image[0, -3:, -3:, -1]
      lowercase    :			Union[str, Any]  =       image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)

      lowercase    :			List[str]  =       np.array(
          [
              0.99_97,
              0.99_88,
              0.00_28,
              0.99_97,
              0.99_84,
              0.99_65,
              0.00_29,
              0.99_86,
              0.00_25,
          ]	)

      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    _UpperCAmelCase	(       self	) ->    List[Any]:







      """simple docstring"""
      lowercase    :			List[str]  =       torch.device("""cpu"""	)
      class lowerCamelCase__  :
         __UpperCAmelCase        =					1




      lowercase    :			List[Any]  =       self.get_dummy_components()

      lowercase    :			int  =       self.pipeline_class(**snake_case	)
      lowercase    :			int  =       pipe.to(snake_case	)

      lowercase    :			Optional[Any]  =       components["""prior"""]
      lowercase    :			str  =       components["""decoder"""]
      lowercase    :			Optional[int]  =       components["""super_res_first"""]
      lowercase    :			str  =       components["""tokenizer"""]
      lowercase    :			Tuple  =       components["""text_encoder"""]

      lowercase    :			int  =       torch.Generator(device=snake_case	).manual_seed(0	)
      lowercase    :			Optional[int]  =       prior.dtype
      lowercase    :			Optional[int]  =       1

      lowercase    :			Optional[Any]  =       (batch_size, prior.config.embedding_dim)
      lowercase    :			Dict  =       pipe.prepare_latents(
          snake_case      ,   dtype=snake_case      ,   device=snake_case      ,   generator=snake_case      ,   latents=snake_case      ,   scheduler=DummyScheduler()	)
      lowercase    :			Optional[Any]  =       (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
      lowercase    :			Any  =       pipe.prepare_latents(
          snake_case      ,   dtype=snake_case      ,   device=snake_case      ,   generator=snake_case      ,   latents=snake_case      ,   scheduler=DummyScheduler()	)

      lowercase    :			List[str]  =       (
          batch_size,
          super_res_first.config.in_channels // 2,
          super_res_first.config.sample_size,
          super_res_first.config.sample_size,
      )
      lowercase    :			int  =       pipe.prepare_latents(
          snake_case      ,   dtype=snake_case      ,   device=snake_case      ,   generator=snake_case      ,   latents=snake_case      ,   scheduler=DummyScheduler()	)

      pipe.set_progress_bar_config(disable=snake_case	)

      lowercase    :			Tuple  =       """this is a prompt example"""

      lowercase    :			Optional[int]  =       torch.Generator(device=snake_case	).manual_seed(0	)
      lowercase    :			Optional[int]  =       pipe(
          [prompt]      ,   generator=snake_case      ,   prior_num_inference_steps=2      ,   decoder_num_inference_steps=2      ,   super_res_num_inference_steps=2      ,   prior_latents=snake_case      ,   decoder_latents=snake_case      ,   super_res_latents=snake_case      ,   output_type="""np"""      ,   )
      lowercase    :			Union[str, Any]  =       output.images

      lowercase    :			Optional[Any]  =       tokenizer(
          snake_case      ,   padding="""max_length"""      ,   max_length=tokenizer.model_max_length      ,   return_tensors="""pt"""      ,   )
      lowercase    :			str  =       text_encoder(text_inputs.input_ids	)
      lowercase    :			int  =       text_inputs.attention_mask

      lowercase    :			List[str]  =       torch.Generator(device=snake_case	).manual_seed(0	)
      lowercase    :			Optional[int]  =       pipe(
          generator=snake_case      ,   prior_num_inference_steps=2      ,   decoder_num_inference_steps=2      ,   super_res_num_inference_steps=2      ,   prior_latents=snake_case      ,   decoder_latents=snake_case      ,   super_res_latents=snake_case      ,   text_model_output=snake_case      ,   text_attention_mask=snake_case      ,   output_type="""np"""      ,   )[0]

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



   @skip_mps
   def    _UpperCAmelCase	(       self	) ->    Tuple:







      """simple docstring"""
      lowercase    :			str  =       torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=snake_case      ,   expected_max_diff=0.01	)



   @skip_mps
   def    _UpperCAmelCase	(       self	) ->    str:







      """simple docstring"""
      lowercase    :			List[str]  =       torch_device == """cpu"""
      lowercase    :			Optional[Any]  =       True
      lowercase    :			Union[str, Any]  =       [
          """prior_num_inference_steps""",
          """decoder_num_inference_steps""",
          """super_res_num_inference_steps""",
      ]

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



   def    _UpperCAmelCase	(       self	) ->    Optional[int]:







      """simple docstring"""
      lowercase    :			List[Any]  =       [
          """prior_num_inference_steps""",
          """decoder_num_inference_steps""",
          """super_res_num_inference_steps""",
      ]

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



   @skip_mps
   def    _UpperCAmelCase	(       self	) ->    Union[str, Any]:







      """simple docstring"""
      return super().test_dict_tuple_outputs_equivalent()



   @skip_mps
   def    _UpperCAmelCase	(       self	) ->    Tuple:







      """simple docstring"""
      return super().test_save_load_local()





   @skip_mps
   def    _UpperCAmelCase	(       self	) ->    Optional[int]:







      """simple docstring"""
      return super().test_save_load_optional_components()




@nightly
class lowerCamelCase__  (    unittest.TestCase				):



   def    _UpperCAmelCase	(       self	) ->    List[Any]:







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





   def    _UpperCAmelCase	(       self	) ->    Dict:







      """simple docstring"""
      lowercase    :			int  =       load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/unclip/karlo_v1_alpha_horse_cpu.npy"""	)

      lowercase    :			Optional[int]  =       UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""	)
      pipeline.set_progress_bar_config(disable=snake_case	)

      lowercase    :			str  =       torch.manual_seed(0	)
      lowercase    :			Any  =       pipeline(
          """horse"""      ,   num_images_per_prompt=1      ,   generator=snake_case      ,   output_type="""np"""      ,   )

      lowercase    :			Union[str, Any]  =       output.images[0]

      assert image.shape == (2_5_6, 2_5_6, 3)
      assert np.abs(expected_image - image	).max() < 1E-1




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



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







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



   def    _UpperCAmelCase	(       self	) ->    Tuple:







      """simple docstring"""
      lowercase    :			Dict  =       load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/unclip/karlo_v1_alpha_horse_fp16.npy"""	)

      lowercase    :			Dict  =       UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""      ,   torch_dtype=torch.floataa	)
      lowercase    :			Optional[Any]  =       pipeline.to(snake_case	)
      pipeline.set_progress_bar_config(disable=snake_case	)

      lowercase    :			Optional[int]  =       torch.Generator(device="""cpu"""	).manual_seed(0	)
      lowercase    :			Dict  =       pipeline(
          """horse"""      ,   generator=snake_case      ,   output_type="""np"""      ,   )

      lowercase    :			Optional[Any]  =       output.images[0]

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

      assert_mean_pixel_difference(snake_case      ,   snake_case	)





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







      """simple docstring"""
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      lowercase    :			Dict  =       UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""      ,   torch_dtype=torch.floataa	)
      lowercase    :			int  =       pipe.to(snake_case	)
      pipe.set_progress_bar_config(disable=snake_case	)
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      lowercase    :			Optional[int]  =       pipe(
          """horse"""      ,   num_images_per_prompt=1      ,   prior_num_inference_steps=2      ,   decoder_num_inference_steps=2      ,   super_res_num_inference_steps=2      ,   output_type="""np"""      ,   )

      lowercase    :			Optional[Any]  =       torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 1_0**9









"""simple docstring"""
import argparse
import json

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

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor




def 			__snake_case						(      __A    )		->				Optional[int]:
   lowercase    :			int  =       TimesformerConfig()

   if "large" in model_name:
      lowercase    :			str  =       96

   if "hr" in model_name:
      lowercase    :			Tuple  =       16
      lowercase    :			Any  =       448

   lowercase    :			List[str]  =       """huggingface/label-files"""
   if "k400" in model_name:
      lowercase    :			Optional[Any]  =       400
      lowercase    :			Dict  =       """kinetics400-id2label.json"""
   elif "k600" in model_name:
      lowercase    :			Tuple  =       600
      lowercase    :			int  =       """kinetics600-id2label.json"""
   elif "ssv2" in model_name:
      lowercase    :			int  =       174
      lowercase    :			Union[str, Any]  =       """something-something-v2-id2label.json"""
   else:
      raise ValueError("""Model name should either contain 'k400', 'k600' or 'ssv2'."""    )
   lowercase    :			Optional[int]  =       json.load(open(hf_hub_download(__A ,__A ,repo_type="""dataset"""    ) ,"""r"""    )    )
   lowercase    :			Optional[int]  =       {int(__A    ): v for k, v in idalabel.items()}
   lowercase    :			int  =       idalabel
   lowercase    :			Tuple  =       {v: k for k, v in idalabel.items()}

   return config




def 			__snake_case						(      __A    )		->				Optional[int]:
   if "encoder." in name:
      lowercase    :			str  =       name.replace("""encoder.""" ,""""""    )
   if "cls_token" in name:
      lowercase    :			Optional[int]  =       name.replace("""cls_token""" ,"""timesformer.embeddings.cls_token"""    )
   if "pos_embed" in name:
      lowercase    :			Union[str, Any]  =       name.replace("""pos_embed""" ,"""timesformer.embeddings.position_embeddings"""    )
   if "time_embed" in name:
      lowercase    :			str  =       name.replace("""time_embed""" ,"""timesformer.embeddings.time_embeddings"""    )
   if "patch_embed.proj" in name:
      lowercase    :			Any  =       name.replace("""patch_embed.proj""" ,"""timesformer.embeddings.patch_embeddings.projection"""    )
   if "patch_embed.norm" in name:
      lowercase    :			Union[str, Any]  =       name.replace("""patch_embed.norm""" ,"""timesformer.embeddings.norm"""    )
   if "blocks" in name:
      lowercase    :			str  =       name.replace("""blocks""" ,"""timesformer.encoder.layer"""    )
   if "attn.proj" in name:
      lowercase    :			Dict  =       name.replace("""attn.proj""" ,"""attention.output.dense"""    )
   if "attn" in name and "bias" not in name and "temporal" not in name:
      lowercase    :			Tuple  =       name.replace("""attn""" ,"""attention.self"""    )
   if "attn" in name and "temporal" not in name:
      lowercase    :			str  =       name.replace("""attn""" ,"""attention.attention"""    )
   if "temporal_norm1" in name:
      lowercase    :			List[str]  =       name.replace("""temporal_norm1""" ,"""temporal_layernorm"""    )
   if "temporal_attn.proj" in name:
      lowercase    :			str  =       name.replace("""temporal_attn""" ,"""temporal_attention.output.dense"""    )
   if "temporal_fc" in name:
      lowercase    :			Any  =       name.replace("""temporal_fc""" ,"""temporal_dense"""    )
   if "norm1" in name and "temporal" not in name:
      lowercase    :			List[Any]  =       name.replace("""norm1""" ,"""layernorm_before"""    )
   if "norm2" in name:
      lowercase    :			Dict  =       name.replace("""norm2""" ,"""layernorm_after"""    )
   if "mlp.fc1" in name:
      lowercase    :			List[Any]  =       name.replace("""mlp.fc1""" ,"""intermediate.dense"""    )
   if "mlp.fc2" in name:
      lowercase    :			Dict  =       name.replace("""mlp.fc2""" ,"""output.dense"""    )
   if "norm.weight" in name and "fc" not in name and "temporal" not in name:
      lowercase    :			Optional[Any]  =       name.replace("""norm.weight""" ,"""timesformer.layernorm.weight"""    )
   if "norm.bias" in name and "fc" not in name and "temporal" not in name:
      lowercase    :			Union[str, Any]  =       name.replace("""norm.bias""" ,"""timesformer.layernorm.bias"""    )
   if "head" in name:
      lowercase    :			Any  =       name.replace("""head""" ,"""classifier"""    )

   return name




def 			__snake_case						(      __A ,__A    )		->				List[str]:
   for key in orig_state_dict.copy().keys():
      lowercase    :			str  =       orig_state_dict.pop(__A    )

      if key.startswith("""model."""    ):
         lowercase    :			Union[str, Any]  =       key.replace("""model.""" ,""""""    )

      if "qkv" in key:
         lowercase    :			List[Any]  =       key.split("""."""    )
         lowercase    :			Optional[int]  =       int(key_split[1]    )
         lowercase    :			Optional[Any]  =       """timesformer.encoder.layer."""
         if "temporal" in key:
            lowercase    :			List[Any]  =       """.temporal_attention.attention.qkv."""
         else:
            lowercase    :			str  =       """.attention.attention.qkv."""
         if "weight" in key:
            lowercase    :			str  =       val
         else:
            lowercase    :			Any  =       val
      else:
         lowercase    :			Optional[Any]  =       val

   return orig_state_dict




def 			__snake_case						(      )		->				Dict:
   lowercase    :			str  =       hf_hub_download(
       repo_id="""hf-internal-testing/spaghetti-video""" ,filename="""eating_spaghetti.npy""" ,repo_type="""dataset"""    )
   lowercase    :			List[Any]  =       np.load(__A    )
   return list(__A    )




def 			__snake_case						(      __A ,__A ,__A ,__A    )		->				Union[str, Any]:
   lowercase    :			Union[str, Any]  =       get_timesformer_config(__A    )

   lowercase    :			Any  =       TimesformerForVideoClassification(__A    )

   # download original checkpoint, hosted on Google Drive
   lowercase    :			Dict  =       """pytorch_model.bin"""
   gdown.cached_download(__A ,__A ,quiet=__A    )
   lowercase    :			str  =       torch.load(__A ,map_location="""cpu"""    )
   if "model" in files:
      lowercase    :			Union[str, Any]  =       files["""model"""]
   elif "module" in files:
      lowercase    :			Tuple  =       files["""module"""]
   else:
      lowercase    :			Any  =       files["""model_state"""]
   lowercase    :			List[Any]  =       convert_state_dict(__A ,__A    )

   model.load_state_dict(__A    )
   model.eval()

   # verify model on basic input
   lowercase    :			Union[str, Any]  =       VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5] ,image_std=[0.5, 0.5, 0.5]    )
   lowercase    :			Dict  =       prepare_video()
   lowercase    :			Optional[int]  =       image_processor(video[:8] ,return_tensors="""pt"""    )

   lowercase    :			Dict  =       model(**__A    )
   lowercase    :			int  =       outputs.logits

   lowercase    :			List[str]  =       [
       # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
       """timesformer-base-finetuned-k400""",
       """timesformer-large-finetuned-k400""",
       """timesformer-hr-finetuned-k400""",
       # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
       """timesformer-base-finetuned-k600""",
       """timesformer-large-finetuned-k600""",
       """timesformer-hr-finetuned-k600""",
       # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
       """timesformer-base-finetuned-ssv2""",
       """timesformer-large-finetuned-ssv2""",
       """timesformer-hr-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 == "timesformer-base-finetuned-k400":
      lowercase    :			Optional[int]  =       torch.Size([1, 400]    )
      lowercase    :			Dict  =       torch.tensor([-0.30_16, -0.77_13, -0.42_05]    )
   elif model_name == "timesformer-base-finetuned-k600":
      lowercase    :			str  =       torch.Size([1, 600]    )
      lowercase    :			List[Any]  =       torch.tensor([-0.72_67, -0.74_66, 3.24_04]    )
   elif model_name == "timesformer-base-finetuned-ssv2":
      lowercase    :			Tuple  =       torch.Size([1, 174]    )
      lowercase    :			Optional[Any]  =       torch.tensor([-0.90_59, 0.64_33, -3.14_57]    )
   elif model_name == "timesformer-large-finetuned-k400":
      lowercase    :			str  =       torch.Size([1, 400]    )
      lowercase    :			List[str]  =       torch.tensor([0, 0, 0]    )
   elif model_name == "timesformer-large-finetuned-k600":
      lowercase    :			List[Any]  =       torch.Size([1, 600]    )
      lowercase    :			str  =       torch.tensor([0, 0, 0]    )
   elif model_name == "timesformer-large-finetuned-ssv2":
      lowercase    :			Dict  =       torch.Size([1, 174]    )
      lowercase    :			List[Any]  =       torch.tensor([0, 0, 0]    )
   elif model_name == "timesformer-hr-finetuned-k400":
      lowercase    :			Tuple  =       torch.Size([1, 400]    )
      lowercase    :			Union[str, Any]  =       torch.tensor([-0.96_17, -3.73_11, -3.77_08]    )
   elif model_name == "timesformer-hr-finetuned-k600":
      lowercase    :			Union[str, Any]  =       torch.Size([1, 600]    )
      lowercase    :			Any  =       torch.tensor([2.52_73, 0.71_27, 1.88_48]    )
   elif model_name == "timesformer-hr-finetuned-ssv2":
      lowercase    :			int  =       torch.Size([1, 174]    )
      lowercase    :			Optional[Any]  =       torch.tensor([-3.67_56, -0.75_13, 0.71_80]    )
   else:
      raise ValueError(F'''Model name not supported. Should be one of {model_names}'''    )

   # verify logits
   assert logits.shape == expected_shape
   assert torch.allclose(logits[0, :3] ,__A ,atol=1E-4    )
   print("""Logits 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(__A    )
      model.save_pretrained(__A    )

   if push_to_hub:
      print("""Pushing to the hub..."""    )
      model.push_to_hub(F'''fcakyon/{model_name}'''    )


if __name__ == "__main__":
   lowerCAmelCase:	Any							      =argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       "--checkpoint_url",
       default="https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download",
       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="",
       type=str,
       help="Path to the output PyTorch model directory.",
   )
   parser.add_argument("--model_name", default="timesformer-base-finetuned-k400", 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."
   )

   lowerCAmelCase:	int							      =parser.parse_args()
   convert_timesformer_checkpoint(
       args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
   )


