
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_ (							snake_case							:   List[str] ,      snake_case							:   Union[str, Any] )      ->    Optional[Any]:



						'''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=snake_case ,      num_frames=snake_case )

						if "large" in model_name:
												__UpperCamelCase       =			768
												__UpperCamelCase       =			3072
												__UpperCamelCase       =			12

												__UpperCamelCase       =			1024
												__UpperCamelCase       =			4096
												__UpperCamelCase       =			16
												__UpperCamelCase       =			24
												__UpperCamelCase       =			768
												__UpperCamelCase       =			3072

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

						__UpperCamelCase       =			XCLIPConfig.from_text_vision_configs(snake_case ,      snake_case )

						if "large" in model_name:
												__UpperCamelCase       =			768

						return config




def 	A_ (							snake_case							:   Optional[int] )      ->    int:



						'''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_ (							snake_case							:   Optional[int] ,      snake_case							:   Optional[Any] )      ->    Dict:



						'''simple docstring'''

						for key in orig_state_dict.copy().keys():
												__UpperCamelCase       =			orig_state_dict.pop(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(snake_case )
																		if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
																								__UpperCamelCase       =			val.T
																		__UpperCamelCase       =			val

						return orig_state_dict




def 	A_ (							snake_case							:   int )      ->    Union[str, Any]:



						'''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=snake_case ,      repo_type='''dataset''' ,      )
						__UpperCamelCase       =			np.load(snake_case )
						return list(snake_case )




def 	A_ (							snake_case							:   Optional[Any] ,      snake_case							:   Dict=None ,      snake_case							:   Optional[Any]=False )      ->    List[str]:



						'''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(snake_case ,      snake_case )
						__UpperCamelCase       =			XCLIPModel(snake_case )
						model.eval()

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

						__UpperCamelCase       =			convert_state_dict(snake_case ,      snake_case )

						__UpperCamelCase       =			XCLIPModel(snake_case )
						__UpperCamelCase  ,							__UpperCamelCase       =			model.load_state_dict(snake_case ,      strict=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=snake_case )
						__UpperCamelCase       =			CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32''' )
						__UpperCamelCase       =			CLIPTokenizerFast.from_pretrained('''openai/clip-vit-base-patch32''' )
						__UpperCamelCase       =			XCLIPProcessor(image_processor=snake_case ,      tokenizer=snake_case )

						__UpperCamelCase       =			prepare_video(snake_case )
						__UpperCamelCase       =			processor(
						    text=['''playing sports''', '''eating spaghetti''', '''go shopping'''] ,      videos=snake_case ,      return_tensors='''pt''' ,      padding=snake_case )

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

						with torch.no_grad():
												__UpperCamelCase       =			model(**snake_case )

						# Verify outputs
						__UpperCamelCase       =			outputs.logits_per_video
						__UpperCamelCase       =			logits_per_video.softmax(dim=1 )
						print('''Probs:''' ,      snake_case )
						# kinetics-400
						if model_name == "xclip-base-patch32":
												__UpperCamelCase       =			torch.tensor([[0.0019, 0.9951, 0.0030]] )
						elif model_name == "xclip-base-patch32-16-frames":
												__UpperCamelCase       =			torch.tensor([[7.09_99e-04, 9.98_83e-01, 4.55_80e-04]] )
						elif model_name == "xclip-base-patch16":
												__UpperCamelCase       =			torch.tensor([[0.0083, 0.9681, 0.0236]] )
						elif model_name == "xclip-base-patch16-16-frames":
												__UpperCamelCase       =			torch.tensor([[7.69_37e-04, 9.97_28e-01, 1.94_73e-03]] )
						elif model_name == "xclip-large-patch14":
												__UpperCamelCase       =			torch.tensor([[0.0062, 0.9864, 0.0075]] )
						elif model_name == "xclip-large-patch14-16-frames":
												__UpperCamelCase       =			torch.tensor([[3.38_77e-04, 9.99_37e-01, 2.88_88e-04]] )
						# kinetics-600
						elif model_name == "xclip-base-patch16-kinetics-600":
												__UpperCamelCase       =			torch.tensor([[0.0555, 0.8914, 0.0531]] )
						elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
												__UpperCamelCase       =			torch.tensor([[3.85_54e-04, 9.99_29e-01, 3.27_54e-04]] )
						elif model_name == "xclip-large-patch14-kinetics-600":
												__UpperCamelCase       =			torch.tensor([[0.0036, 0.9920, 0.0045]] )
						# few shot
						elif model_name == "xclip-base-patch16-hmdb-2-shot":
												__UpperCamelCase       =			torch.tensor([[7.18_90e-06, 9.99_94e-01, 5.65_59e-05]] )
						elif model_name == "xclip-base-patch16-hmdb-4-shot":
												__UpperCamelCase       =			torch.tensor([[1.03_20e-05, 9.99_93e-01, 6.24_35e-05]] )
						elif model_name == "xclip-base-patch16-hmdb-8-shot":
												__UpperCamelCase       =			torch.tensor([[4.13_77e-06, 9.99_90e-01, 9.83_86e-05]] )
						elif model_name == "xclip-base-patch16-hmdb-16-shot":
												__UpperCamelCase       =			torch.tensor([[4.13_47e-05, 9.99_62e-01, 3.34_11e-04]] )
						elif model_name == "xclip-base-patch16-ucf-2-shot":
												__UpperCamelCase       =			torch.tensor([[8.58_57e-05, 9.99_28e-01, 6.32_91e-04]] )
						elif model_name == "xclip-base-patch16-ucf-4-shot":
												__UpperCamelCase       =			torch.tensor([[8.58_57e-05, 9.99_28e-01, 6.32_91e-04]] )
						elif model_name == "xclip-base-patch16-ucf-8-shot":
												__UpperCamelCase       =			torch.tensor([[0.0027, 0.9904, 0.0070]] )
						elif model_name == "xclip-base-patch16-ucf-16-shot":
												__UpperCamelCase       =			torch.tensor([[9.82_19e-04, 9.95_93e-01, 3.08_63e-03]] )
						# zero shot
						elif model_name == "xclip-base-patch16-zero-shot":
												__UpperCamelCase       =			torch.tensor([[3.50_82e-04, 9.97_85e-01, 1.79_66e-03]] )
						else:
												raise ValueError(f"Model name {model_name} not supported" )
						assert torch.allclose(snake_case ,      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(snake_case )

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


if __name__ == "__main__":
				lowercase__							:							str									=			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."
				)

				lowercase__							:							List[Any]									=			parser.parse_args()
				convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






from typing import List, Optional, Tuple, Union

import torch

from ...schedulers import DDIMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput






class        SCREAMING_SNAKE_CASE__    (							SCREAMING_SNAKE_CASE_    ):

							"""simple docstring"""






							def __init__(  self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			)->   str:





													'''simple docstring'''






													super().__init__()

													# make sure scheduler can always be converted to DDIM
													__UpperCamelCase       =			DDIMScheduler.from_config(scheduler.config			)

													self.register_modules(unet=SCREAMING_SNAKE_CASE_     ,	scheduler=SCREAMING_SNAKE_CASE_			)



							@torch.no_grad()
							def __call__(  self     ,	SCREAMING_SNAKE_CASE_ = 1     ,	SCREAMING_SNAKE_CASE_ = None     ,	SCREAMING_SNAKE_CASE_ = 0.0     ,	SCREAMING_SNAKE_CASE_ = 50     ,	SCREAMING_SNAKE_CASE_ = None     ,	SCREAMING_SNAKE_CASE_ = "pil"     ,	SCREAMING_SNAKE_CASE_ = True     ,	)->   Union[ImagePipelineOutput, Tuple]:





													'''simple docstring'''






													if isinstance(self.unet.config.sample_size     ,	SCREAMING_SNAKE_CASE_			):
																			__UpperCamelCase       =			(
																			    batch_size,
																			    self.unet.config.in_channels,
																			    self.unet.config.sample_size,
																			    self.unet.config.sample_size,
																			)
													else:
																			__UpperCamelCase       =			(batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)

													if isinstance(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			) and len(SCREAMING_SNAKE_CASE_			) != batch_size:
																			raise ValueError(
																			    F"You have passed a list of generators of length {len(SCREAMING_SNAKE_CASE_			)}, but requested an effective batch"
																			    F" size of {batch_size}. Make sure the batch size matches the length of the generators."			)

													__UpperCamelCase       =			randn_tensor(SCREAMING_SNAKE_CASE_     ,	generator=SCREAMING_SNAKE_CASE_     ,	device=self.device     ,	dtype=self.unet.dtype			)

													# set step values
													self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE_			)

													for t in self.progress_bar(self.scheduler.timesteps			):
																			# 1. predict noise model_output
																			__UpperCamelCase       =			self.unet(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			).sample

																			# 2. predict previous mean of image x_t-1 and add variance depending on eta
																			# eta corresponds to η in paper and should be between [0, 1]
																			# do x_t -> x_t-1
																			__UpperCamelCase       =			self.scheduler.step(
																			    SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	eta=SCREAMING_SNAKE_CASE_     ,	use_clipped_model_output=SCREAMING_SNAKE_CASE_     ,	generator=SCREAMING_SNAKE_CASE_			).prev_sample

													__UpperCamelCase       =			(image / 2 + 0.5).clamp(0     ,	1			)
													__UpperCamelCase       =			image.cpu().permute(0     ,	2     ,	3     ,	1			).numpy()
													if output_type == "pil":
																			__UpperCamelCase       =			self.numpy_to_pil(SCREAMING_SNAKE_CASE_			)

													if not return_dict:
																			return (image,)

													return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE_			)




