


"""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   (	a,		a       ):





       """simple docstring"""







       _a							 =					XCLIPTextConfig()

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

       if "large" in model_name:
              _a							 =					7_6_8
              _a							 =					3_0_7_2
              _a							 =					1_2

              _a							 =					1_0_2_4
              _a							 =					4_0_9_6
              _a							 =					1_6
              _a							 =					2_4
              _a							 =					7_6_8
              _a							 =					3_0_7_2

       if model_name == "xclip-large-patch14-16-frames":
              _a							 =					3_3_6

       _a							 =					XCLIPConfig.from_text_vision_configs(_snake_case,		_snake_case       )

       if "large" in model_name:
              _a							 =					7_6_8

       return config



def     __a   (	a       ):





       """simple docstring"""







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

       return name



def     __a   (	a,		a       ):





       """simple docstring"""







       for key in orig_state_dict.copy().keys():
              _a							 =					orig_state_dict.pop(_snake_case       )

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

       return orig_state_dict



def     __a   (	a       ):





       """simple docstring"""







       if num_frames == 8:
              _a							 =					"eating_spaghetti_8_frames.npy"
       elif num_frames == 1_6:
              _a							 =					"eating_spaghetti.npy"
       elif num_frames == 3_2:
              _a							 =					"eating_spaghetti_32_frames.npy"
       _a							 =					hf_hub_download(
           repo_id="hf-internal-testing/spaghetti-video",		filename=_snake_case,		repo_type="dataset",		)
       _a							 =					np.load(_snake_case       )
       return list(_snake_case       )



def     __a   (	a,		a=None,		a=False       ):





       """simple docstring"""







       _a							 =					{
           # 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",
       }

       _a							 =					model_to_url[model_name]
       _a							 =					8
       if "16-frames" in model_name:
              _a							 =					1_6
       elif "shot" in model_name:
              _a							 =					3_2

       _a							 =					get_xclip_config(_snake_case,		_snake_case       )
       _a							 =					XCLIPModel(_snake_case       )
       model.eval()

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

       _a							 =					convert_state_dict(_snake_case,		_snake_case       )

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

       _a							 =					3_3_6 if model_name == "xclip-large-patch14-16-frames" else 2_2_4
       _a							 =					VideoMAEImageProcessor(size=_snake_case       )
       _a							 =					CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"       )
       _a							 =					CLIPTokenizerFast.from_pretrained("openai/clip-vit-base-patch32"       )
       _a							 =					XCLIPProcessor(image_processor=_snake_case,		tokenizer=_snake_case       )

       _a							 =					prepare_video(_snake_case       )
       _a							 =					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():
              _a							 =					model(**_snake_case       )

       # Verify outputs
       _a							 =					outputs.logits_per_video
       _a							 =					logits_per_video.softmax(dim=1       )
       print("Probs:",		_snake_case       )
       # kinetics-400
       if model_name == "xclip-base-patch32":
              _a							 =					torch.tensor([[0.0019, 0.9951, 0.0030]]       )
       elif model_name == "xclip-base-patch32-16-frames":
              _a							 =					torch.tensor([[7.0_999e-04, 9.9_883e-01, 4.5_580e-04]]       )
       elif model_name == "xclip-base-patch16":
              _a							 =					torch.tensor([[0.0083, 0.9681, 0.0236]]       )
       elif model_name == "xclip-base-patch16-16-frames":
              _a							 =					torch.tensor([[7.6_937e-04, 9.9_728e-01, 1.9_473e-03]]       )
       elif model_name == "xclip-large-patch14":
              _a							 =					torch.tensor([[0.0062, 0.9864, 0.0075]]       )
       elif model_name == "xclip-large-patch14-16-frames":
              _a							 =					torch.tensor([[3.3_877e-04, 9.9_937e-01, 2.8_888e-04]]       )
       # kinetics-600
       elif model_name == "xclip-base-patch16-kinetics-600":
              _a							 =					torch.tensor([[0.0555, 0.8914, 0.0531]]       )
       elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
              _a							 =					torch.tensor([[3.8_554e-04, 9.9_929e-01, 3.2_754e-04]]       )
       elif model_name == "xclip-large-patch14-kinetics-600":
              _a							 =					torch.tensor([[0.0036, 0.9920, 0.0045]]       )
       # few shot
       elif model_name == "xclip-base-patch16-hmdb-2-shot":
              _a							 =					torch.tensor([[7.1_890e-06, 9.9_994e-01, 5.6_559e-05]]       )
       elif model_name == "xclip-base-patch16-hmdb-4-shot":
              _a							 =					torch.tensor([[1.0_320e-05, 9.9_993e-01, 6.2_435e-05]]       )
       elif model_name == "xclip-base-patch16-hmdb-8-shot":
              _a							 =					torch.tensor([[4.1_377e-06, 9.9_990e-01, 9.8_386e-05]]       )
       elif model_name == "xclip-base-patch16-hmdb-16-shot":
              _a							 =					torch.tensor([[4.1_347e-05, 9.9_962e-01, 3.3_411e-04]]       )
       elif model_name == "xclip-base-patch16-ucf-2-shot":
              _a							 =					torch.tensor([[8.5_857e-05, 9.9_928e-01, 6.3_291e-04]]       )
       elif model_name == "xclip-base-patch16-ucf-4-shot":
              _a							 =					torch.tensor([[8.5_857e-05, 9.9_928e-01, 6.3_291e-04]]       )
       elif model_name == "xclip-base-patch16-ucf-8-shot":
              _a							 =					torch.tensor([[0.0027, 0.9904, 0.0070]]       )
       elif model_name == "xclip-base-patch16-ucf-16-shot":
              _a							 =					torch.tensor([[9.8_219e-04, 9.9_593e-01, 3.0_863e-03]]       )
       # zero shot
       elif model_name == "xclip-base-patch16-zero-shot":
              _a							 =					torch.tensor([[3.5_082e-04, 9.9_785e-01, 1.7_966e-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__":
   __SCREAMING_SNAKE_CASE      		=    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."""
   )

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








from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available






@dataclass
class 					SCREAMING_SNAKE_CASE_  (       _a ):
       """simple docstring"""


       __lowerCAmelCase    :		Union[List[PIL.Image.Image], np.ndarray]
       __lowerCAmelCase    :		Optional[List[bool]]


if is_transformers_available() and is_torch_available():
    from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline


