





"""simple docstring"""





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import 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_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





@skip_mps
class        _lowerCamelCase      (       _lowercase  ,			_lowercase  ,			unittest.TestCase ):
						UpperCAmelCase_ =			StableDiffusionPanoramaPipeline
						UpperCAmelCase_ =			TEXT_TO_IMAGE_PARAMS
						UpperCAmelCase_ =			TEXT_TO_IMAGE_BATCH_PARAMS
						UpperCAmelCase_ =			TEXT_TO_IMAGE_IMAGE_PARAMS
						UpperCAmelCase_ =			TEXT_TO_IMAGE_IMAGE_PARAMS


						def     snake_case_   (self		) ->   int:
								torch.manual_seed(0		)
								UpperCamelCase   =	UNetaDConditionModel(
								    block_out_channels=(32, 64)   ,							layers_per_block=1   ,							sample_size=32   ,							in_channels=4   ,							out_channels=4   ,							down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,							up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,							cross_attention_dim=32   ,							)
								UpperCamelCase   =	DDIMScheduler()
								torch.manual_seed(0		)
								UpperCamelCase   =	AutoencoderKL(
								    block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,							up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,							latent_channels=4   ,							)
								torch.manual_seed(0		)
								UpperCamelCase   =	CLIPTextConfig(
								    bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1e-0_5   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=10_00   ,							)
								UpperCamelCase   =	CLIPTextModel(__a		)
								UpperCamelCase   =	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

								UpperCamelCase   =	{
								    "unet": unet,
								    "scheduler": scheduler,
								    "vae": vae,
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "safety_checker": None,
								    "feature_extractor": None,
								}
								return components


						def     snake_case_   (self   ,							__a   ,							__a=0		) ->   List[str]:
								UpperCamelCase   =	torch.manual_seed(__a		)
								UpperCamelCase   =	{
								    "prompt": "a photo of the dolomites",
								    "generator": generator,
								    # Setting height and width to None to prevent OOMs on CPU.
								    "height": None,
								    "width": None,
								    "num_inference_steps": 1,
								    "guidance_scale": 6.0,
								    "output_type": "numpy",
								}
								return inputs


						def     snake_case_   (self		) ->   int:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								UpperCamelCase   =	np.array([0.6186, 0.5374, 0.4915, 0.4135, 0.4114, 0.4563, 0.5128, 0.4977, 0.4757]		)

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


						def     snake_case_   (self		) ->   Optional[int]:
								super().test_inference_batch_consistent(batch_sizes=[1, 2]		)


						def     snake_case_   (self		) ->   Tuple:
								super().test_inference_batch_single_identical(batch_size=2   ,							expected_max_diff=3.2_5e-3		)


						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	"french fries"
								UpperCamelCase   =	sd_pipe(**__a   ,							negative_prompt=__a		)
								UpperCamelCase   =	output.images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]		)

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


						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a   ,							view_batch_size=2		)
								UpperCamelCase   =	output.images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]		)

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


						def     snake_case_   (self		) ->   Optional[int]:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	EulerAncestralDiscreteScheduler(
								    beta_start=0.00085   ,							beta_end=0.012   ,							beta_schedule="scaled_linear"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.4024, 0.6510, 0.4901, 0.5378, 0.5813, 0.5622, 0.4795, 0.4467, 0.4952]		)

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




						def     snake_case_   (self		) ->   Optional[Any]:
								UpperCamelCase   =	"cpu"  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	PNDMScheduler(
								    beta_start=0.00085   ,							beta_end=0.012   ,							beta_schedule="scaled_linear"   ,							skip_prk_steps=__a		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline(**__a		)
								UpperCamelCase   =	sd_pipe.to(__a		)
								sd_pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	self.get_dummy_inputs(__a		)
								UpperCamelCase   =	sd_pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1]

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

								UpperCamelCase   =	np.array([0.6391, 0.6291, 0.4861, 0.5134, 0.5552, 0.4578, 0.5032, 0.5023, 0.4539]		)

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





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


						def     snake_case_   (self		) ->   List[str]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


						def     snake_case_   (self   ,							__a=0		) ->   List[str]:
								UpperCamelCase   =	torch.manual_seed(__a		)
								UpperCamelCase   =	{
								    "prompt": "a photo of the dolomites",
								    "generator": generator,
								    "num_inference_steps": 3,
								    "guidance_scale": 7.5,
								    "output_type": "numpy",
								}
								return inputs


						def     snake_case_   (self		) ->   Union[str, Any]:
								UpperCamelCase   =	"stabilityai/stable-diffusion-2-base"
								UpperCamelCase   =	DDIMScheduler.from_pretrained(__a   ,							subfolder="scheduler"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(__a   ,							scheduler=__a   ,							safety_checker=__a		)
								pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing()

								UpperCamelCase   =	self.get_inputs()
								UpperCamelCase   =	pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 20_48, 3)

								UpperCamelCase   =	np.array(
								    [
								        0.36968392,
								        0.27025372,
								        0.32446766,
								        0.28379387,
								        0.36363274,
								        0.30733347,
								        0.27100027,
								        0.27054125,
								        0.25536096,
								    ]		)

								assert np.abs(expected_slice - image_slice		).max() < 1e-2


						def     snake_case_   (self		) ->   Union[str, Any]:
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(
								    "stabilityai/stable-diffusion-2-base"   ,							safety_checker=__a		)
								UpperCamelCase   =	LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing()

								UpperCamelCase   =	self.get_inputs()
								UpperCamelCase   =	pipe(**__a		).images
								UpperCamelCase   =	image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 20_48, 3)

								UpperCamelCase   =	np.array(
								    [
								        [
								            0.0,
								            0.0,
								            0.0,
								            0.0,
								            0.0,
								            0.0,
								            0.0,
								            0.0,
								            0.0,
								        ]
								    ]		)

								assert np.abs(expected_slice - image_slice		).max() < 1e-3


						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	0

								def callback_fn(__a   ,							__a   ,							__a		) -> None:
										UpperCamelCase   =	True
										nonlocal number_of_steps
										number_of_steps += 1
										if step == 1:
												UpperCamelCase   =	latents.detach().cpu().numpy()
												assert latents.shape == (1, 4, 64, 2_56)
												UpperCamelCase   =	latents[0, -3:, -3:, -1]

												UpperCamelCase   =	np.array(
												    [
												        0.18681869,
												        0.33907816,
												        0.5361276,
												        0.14432865,
												        -0.02856611,
												        -0.73941123,
												        0.23397987,
												        0.47322682,
												        -0.37823164,
												    ]		)
												assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5e-2
										elif step == 2:
												UpperCamelCase   =	latents.detach().cpu().numpy()
												assert latents.shape == (1, 4, 64, 2_56)
												UpperCamelCase   =	latents[0, -3:, -3:, -1]

												UpperCamelCase   =	np.array(
												    [
												        0.18539645,
												        0.33987248,
												        0.5378559,
												        0.14437142,
												        -0.02455261,
												        -0.7338317,
												        0.23990755,
												        0.47356272,
												        -0.3786505,
												    ]		)

												assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5e-2

								UpperCamelCase   =	False

								UpperCamelCase   =	"stabilityai/stable-diffusion-2-base"
								UpperCamelCase   =	DDIMScheduler.from_pretrained(__a   ,							subfolder="scheduler"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(__a   ,							scheduler=__a   ,							safety_checker=__a		)
								UpperCamelCase   =	pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing()

								UpperCamelCase   =	self.get_inputs()
								pipe(**__a   ,							callback=__a   ,							callback_steps=1		)
								assert callback_fn.has_been_called
								assert number_of_steps == 3




						def     snake_case_   (self		) ->   List[Any]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								UpperCamelCase   =	"stabilityai/stable-diffusion-2-base"
								UpperCamelCase   =	DDIMScheduler.from_pretrained(__a   ,							subfolder="scheduler"		)
								UpperCamelCase   =	StableDiffusionPanoramaPipeline.from_pretrained(__a   ,							scheduler=__a   ,							safety_checker=__a		)
								UpperCamelCase   =	pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)
								pipe.enable_attention_slicing(1		)
								pipe.enable_sequential_cpu_offload()

								UpperCamelCase   =	self.get_inputs()
								UpperCamelCase   =	pipe(**__a		)

								UpperCamelCase   =	torch.cuda.max_memory_allocated()
								# make sure that less than 5.2 GB is allocated
								assert mem_bytes < 5.5 * 10**9













"""simple docstring"""





import argparse

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

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	XCLIPTextConfig()

		# derive patch size from model name
		UpperCamelCase   =	model_name.find("patch"    )
		UpperCamelCase   =	int(model_name[start_idx + len("patch"    ) : start_idx + len("patch"    ) + 2]    )
		UpperCamelCase   =	XCLIPVisionConfig(patch_size=_SCREAMING_SNAKE_CASE  ,	num_frames=_SCREAMING_SNAKE_CASE    )

		if "large" in model_name:
				UpperCamelCase   =	768
				UpperCamelCase   =	3_072
				UpperCamelCase   =	12

				UpperCamelCase   =	1_024
				UpperCamelCase   =	4_096
				UpperCamelCase   =	16
				UpperCamelCase   =	24
				UpperCamelCase   =	768
				UpperCamelCase   =	3_072

		if model_name == "xclip-large-patch14-16-frames":
				UpperCamelCase   =	336

		UpperCamelCase   =	XCLIPConfig.from_text_vision_configs(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

		if "large" in model_name:
				UpperCamelCase   =	768

		return config


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		if name == "token_embedding.weight":
				UpperCamelCase   =	name.replace("token_embedding.weight"  ,	"text_model.embeddings.token_embedding.weight"    )
		if name == "positional_embedding":
				UpperCamelCase   =	name.replace("positional_embedding"  ,	"text_model.embeddings.position_embedding.weight"    )
		if "ln_1" in name:
				UpperCamelCase   =	name.replace("ln_1"  ,	"layer_norm1"    )
		if "ln_2" in name:
				UpperCamelCase   =	name.replace("ln_2"  ,	"layer_norm2"    )
		if "c_fc" in name:
				UpperCamelCase   =	name.replace("c_fc"  ,	"fc1"    )
		if "c_proj" in name:
				UpperCamelCase   =	name.replace("c_proj"  ,	"fc2"    )
		if name.startswith("transformer.resblocks"    ):
				UpperCamelCase   =	name.replace("transformer.resblocks"  ,	"text_model.encoder.layers"    )
		if "attn.out_proj" in name and "message" not in name:
				UpperCamelCase   =	name.replace("attn.out_proj"  ,	"self_attn.out_proj"    )
		if "ln_final" in name:
				UpperCamelCase   =	name.replace("ln_final"  ,	"text_model.final_layer_norm"    )
		# visual encoder
		if name == "visual.class_embedding":
				UpperCamelCase   =	name.replace("visual.class_embedding"  ,	"vision_model.embeddings.class_embedding"    )
		if name == "visual.positional_embedding":
				UpperCamelCase   =	name.replace("visual.positional_embedding"  ,	"vision_model.embeddings.position_embedding.weight"    )
		if name.startswith("visual.transformer.resblocks"    ):
				UpperCamelCase   =	name.replace("visual.transformer.resblocks"  ,	"vision_model.encoder.layers"    )
		if "visual.conv1" in name:
				UpperCamelCase   =	name.replace("visual.conv1"  ,	"vision_model.embeddings.patch_embedding"    )
		if "visual.ln_pre" in name:
				UpperCamelCase   =	name.replace("visual.ln_pre"  ,	"vision_model.pre_layernorm"    )
		if "visual.ln_post" in name:
				UpperCamelCase   =	name.replace("visual.ln_post"  ,	"vision_model.post_layernorm"    )
		if "visual.proj" in name:
				UpperCamelCase   =	name.replace("visual.proj"  ,	"visual_projection.weight"    )
		if "text_projection" in name:
				UpperCamelCase   =	name.replace("text_projection"  ,	"text_projection.weight"    )
		# things on top
		if "prompts_visual_proj" in name:
				UpperCamelCase   =	name.replace("prompts_visual_proj"  ,	"prompts_visual_projection"    )
		if "prompts_visual_ln" in name:
				UpperCamelCase   =	name.replace("prompts_visual_ln"  ,	"prompts_visual_layernorm"    )
		# mit
		if name == "mit.positional_embedding":
				UpperCamelCase   =	name.replace("positional"  ,	"position"    )
		if name.startswith("mit.resblocks"    ):
				UpperCamelCase   =	name.replace("mit.resblocks"  ,	"mit.encoder.layers"    )
		# prompts generator
		if name.startswith("prompts_generator.norm"    ):
				UpperCamelCase   =	name.replace("prompts_generator.norm"  ,	"prompts_generator.layernorm"    )

		return name


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		for key in orig_state_dict.copy().keys():
				UpperCamelCase   =	orig_state_dict.pop(_SCREAMING_SNAKE_CASE    )

				if "attn.in_proj" in key:
						UpperCamelCase   =	key.split("."    )
						if key.startswith("visual"    ):
								UpperCamelCase   =	key_split[3]
								UpperCamelCase   =	config.vision_config.hidden_size
								if "message_attn" in key:
										if "weight" in key:
												UpperCamelCase   =	val[
												    :dim, :
												]
												UpperCamelCase   =	val[
												    dim : dim * 2, :
												]
												UpperCamelCase   =	val[
												    -dim:, :
												]
										else:
												UpperCamelCase   =	val[
												    :dim
												]
												UpperCamelCase   =	val[
												    dim : dim * 2
												]
												UpperCamelCase   =	val[
												    -dim:
												]
								else:
										if "weight" in key:
												UpperCamelCase   =	val[
												    :dim, :
												]
												UpperCamelCase   =	val[
												    dim : dim * 2, :
												]
												UpperCamelCase   =	val[
												    -dim:, :
												]
										else:
												UpperCamelCase   =	val[:dim]
												UpperCamelCase   =	val[
												    dim : dim * 2
												]
												UpperCamelCase   =	val[-dim:]
						elif key.startswith("mit"    ):
								UpperCamelCase   =	key_split[2]
								UpperCamelCase   =	config.vision_config.mit_hidden_size
								if "weight" in key:
										UpperCamelCase   =	val[:dim, :]
										UpperCamelCase   =	val[dim : dim * 2, :]
										UpperCamelCase   =	val[-dim:, :]
								else:
										UpperCamelCase   =	val[:dim]
										UpperCamelCase   =	val[dim : dim * 2]
										UpperCamelCase   =	val[-dim:]
						else:
								UpperCamelCase   =	key_split[2]
								UpperCamelCase   =	config.text_config.hidden_size
								if "weight" in key:
										UpperCamelCase   =	val[:dim, :]
										UpperCamelCase   =	val[
										    dim : dim * 2, :
										]
										UpperCamelCase   =	val[-dim:, :]
								else:
										UpperCamelCase   =	val[:dim]
										UpperCamelCase   =	val[
										    dim : dim * 2
										]
										UpperCamelCase   =	val[-dim:]
				else:
						UpperCamelCase   =	rename_key(_SCREAMING_SNAKE_CASE    )
						if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
								UpperCamelCase   =	val.T
						UpperCamelCase   =	val

		return orig_state_dict


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		if num_frames == 8:
				UpperCamelCase   =	"eating_spaghetti_8_frames.npy"
		elif num_frames == 16:
				UpperCamelCase   =	"eating_spaghetti.npy"
		elif num_frames == 32:
				UpperCamelCase   =	"eating_spaghetti_32_frames.npy"
		UpperCamelCase   =	hf_hub_download(
		    repo_id="hf-internal-testing/spaghetti-video"  ,	filename=_SCREAMING_SNAKE_CASE  ,	repo_type="dataset"  ,	)
		UpperCamelCase   =	np.load(_SCREAMING_SNAKE_CASE    )
		return list(_SCREAMING_SNAKE_CASE    )


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE=None  ,	_SCREAMING_SNAKE_CASE=False    ):




		"""simple docstring"""


		UpperCamelCase   =	{
		    # fully supervised kinetics-400 checkpoints
		    "xclip-base-patch32": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth",
		    "xclip-base-patch32-16-frames": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth"
		    ),
		    "xclip-base-patch16": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth",
		    "xclip-base-patch16-16-frames": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth"
		    ),
		    "xclip-large-patch14": "https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb",
		    "xclip-large-patch14-16-frames": "https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f",
		    # fully supervised kinetics-600 checkpoints
		    "xclip-base-patch16-kinetics-600": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth"
		    ),
		    "xclip-base-patch16-kinetics-600-16-frames": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth"
		    ),
		    "xclip-large-patch14-kinetics-600": "https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be",
		    # few shot
		    "xclip-base-patch16-hmdb-2-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth"
		    ),
		    "xclip-base-patch16-hmdb-4-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth"
		    ),
		    "xclip-base-patch16-hmdb-8-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth"
		    ),
		    "xclip-base-patch16-hmdb-16-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth"
		    ),
		    "xclip-base-patch16-ucf-2-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth"
		    ),
		    "xclip-base-patch16-ucf-4-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth"
		    ),
		    "xclip-base-patch16-ucf-8-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth"
		    ),
		    "xclip-base-patch16-ucf-16-shot": (
		        "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth"
		    ),
		    # zero shot
		    "xclip-base-patch16-zero-shot": "https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth",
		}

		UpperCamelCase   =	model_to_url[model_name]
		UpperCamelCase   =	8
		if "16-frames" in model_name:
				UpperCamelCase   =	16
		elif "shot" in model_name:
				UpperCamelCase   =	32

		UpperCamelCase   =	get_xclip_config(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	XCLIPModel(_SCREAMING_SNAKE_CASE    )
		model.eval()

		if "drive" in checkpoint_url:
				UpperCamelCase   =	"pytorch_model.bin"
				gdown.cached_download(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	quiet=_SCREAMING_SNAKE_CASE    )
				UpperCamelCase   =	torch.load(_SCREAMING_SNAKE_CASE  ,	map_location="cpu"    )["model"]
		else:
				UpperCamelCase   =	torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE    )["model"]

		UpperCamelCase   =	convert_state_dict(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

		UpperCamelCase   =	XCLIPModel(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   ,				UpperCamelCase   =	model.load_state_dict(_SCREAMING_SNAKE_CASE  ,	strict=_SCREAMING_SNAKE_CASE    )
		assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
		model.eval()

		UpperCamelCase   =	336 if model_name == "xclip-large-patch14-16-frames" else 224
		UpperCamelCase   =	VideoMAEImageProcessor(size=_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"    )
		UpperCamelCase   =	CLIPTokenizerFast.from_pretrained("openai/clip-vit-base-patch32"    )
		UpperCamelCase   =	XCLIPProcessor(image_processor=_SCREAMING_SNAKE_CASE  ,	tokenizer=_SCREAMING_SNAKE_CASE    )

		UpperCamelCase   =	prepare_video(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	processor(
		    text=["playing sports", "eating spaghetti", "go shopping"]  ,	videos=_SCREAMING_SNAKE_CASE  ,	return_tensors="pt"  ,	padding=_SCREAMING_SNAKE_CASE    )

		print("Shape of pixel values:"  ,	inputs.pixel_values.shape    )

		with torch.no_grad():
				UpperCamelCase   =	model(**_SCREAMING_SNAKE_CASE    )

		# Verify outputs
		UpperCamelCase   =	outputs.logits_per_video
		UpperCamelCase   =	logits_per_video.softmax(dim=1    )
		print("Probs:"  ,	_SCREAMING_SNAKE_CASE    )
		# kinetics-400
		if model_name == "xclip-base-patch32":
				UpperCamelCase   =	torch.tensor([[0.00_19, 0.99_51, 0.00_30]]    )
		elif model_name == "xclip-base-patch32-16-frames":
				UpperCamelCase   =	torch.tensor([[7.0999e-04, 9.9883e-01, 4.5580e-04]]    )
		elif model_name == "xclip-base-patch16":
				UpperCamelCase   =	torch.tensor([[0.00_83, 0.96_81, 0.02_36]]    )
		elif model_name == "xclip-base-patch16-16-frames":
				UpperCamelCase   =	torch.tensor([[7.6937e-04, 9.9728e-01, 1.9473e-03]]    )
		elif model_name == "xclip-large-patch14":
				UpperCamelCase   =	torch.tensor([[0.00_62, 0.98_64, 0.00_75]]    )
		elif model_name == "xclip-large-patch14-16-frames":
				UpperCamelCase   =	torch.tensor([[3.3877e-04, 9.9937e-01, 2.8888e-04]]    )
		# kinetics-600
		elif model_name == "xclip-base-patch16-kinetics-600":
				UpperCamelCase   =	torch.tensor([[0.05_55, 0.89_14, 0.05_31]]    )
		elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
				UpperCamelCase   =	torch.tensor([[3.8554e-04, 9.9929e-01, 3.2754e-04]]    )
		elif model_name == "xclip-large-patch14-kinetics-600":
				UpperCamelCase   =	torch.tensor([[0.00_36, 0.99_20, 0.00_45]]    )
		# few shot
		elif model_name == "xclip-base-patch16-hmdb-2-shot":
				UpperCamelCase   =	torch.tensor([[7.1890e-06, 9.9994e-01, 5.6559e-05]]    )
		elif model_name == "xclip-base-patch16-hmdb-4-shot":
				UpperCamelCase   =	torch.tensor([[1.0320e-05, 9.9993e-01, 6.2435e-05]]    )
		elif model_name == "xclip-base-patch16-hmdb-8-shot":
				UpperCamelCase   =	torch.tensor([[4.1377e-06, 9.9990e-01, 9.8386e-05]]    )
		elif model_name == "xclip-base-patch16-hmdb-16-shot":
				UpperCamelCase   =	torch.tensor([[4.1347e-05, 9.9962e-01, 3.3411e-04]]    )
		elif model_name == "xclip-base-patch16-ucf-2-shot":
				UpperCamelCase   =	torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]    )
		elif model_name == "xclip-base-patch16-ucf-4-shot":
				UpperCamelCase   =	torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]    )
		elif model_name == "xclip-base-patch16-ucf-8-shot":
				UpperCamelCase   =	torch.tensor([[0.00_27, 0.99_04, 0.00_70]]    )
		elif model_name == "xclip-base-patch16-ucf-16-shot":
				UpperCamelCase   =	torch.tensor([[9.8219e-04, 9.9593e-01, 3.0863e-03]]    )
		# zero shot
		elif model_name == "xclip-base-patch16-zero-shot":
				UpperCamelCase   =	torch.tensor([[3.5082e-04, 9.9785e-01, 1.7966e-03]]    )
		else:
				raise ValueError(F"Model name {model_name} not supported"    )
		assert torch.allclose(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	atol=1e-3    )
		print("Looks ok!"    )

		if pytorch_dump_folder_path is not None:
				print(F"Saving model {model_name} to {pytorch_dump_folder_path}"    )
				model.save_pretrained(_SCREAMING_SNAKE_CASE    )

		if push_to_hub:
				print("Pushing model, processor and slow tokenizer files to the hub..."    )
				model.push_to_hub(_SCREAMING_SNAKE_CASE  ,	organization="nielsr"    )
				processor.push_to_hub(_SCREAMING_SNAKE_CASE  ,	organization="nielsr"    )
				slow_tokenizer.push_to_hub(_SCREAMING_SNAKE_CASE  ,	organization="nielsr"    )


if __name__ == "__main__":
						lowerCAmelCase__     			=			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--model_name''',
						    default='''xclip-base-patch32''',
						    type=str,
						    help='''Name of the model.''',
						)
						parser.add_argument(
						    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
						)
						parser.add_argument(
						    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
						)

						lowerCAmelCase__     			=			parser.parse_args()
						convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






