




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 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_      )       ->	Dict:

    """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=lowercase_					,     num_frames=lowercase_      )

    if "large" in model_name:
        __UpperCamelCase  		=						7_68
        __UpperCamelCase  		=						30_72
        __UpperCamelCase  		=						12

        __UpperCamelCase  		=						10_24
        __UpperCamelCase  		=						40_96
        __UpperCamelCase  		=						16
        __UpperCamelCase  		=						24
        __UpperCamelCase  		=						7_68
        __UpperCamelCase  		=						30_72

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

    __UpperCamelCase  		=						XCLIPConfig.from_text_vision_configs(lowercase_					,     lowercase_      )

    if "large" in model_name:
        __UpperCamelCase  		=						7_68

    return config




def 					__SCREAMING_SNAKE_CASE   (  lowercase_      )       ->	Optional[Any]:

    """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 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_      )       ->	List[str]:

    """simple docstring"""


    for key in orig_state_dict.copy().keys():
        __UpperCamelCase  		=						orig_state_dict.pop(lowercase_      )

        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(lowercase_      )
            if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
                __UpperCamelCase  		=						val.T
            __UpperCamelCase  		=						val

    return orig_state_dict




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




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_=None					,     lowercase_=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(lowercase_					,     lowercase_      )
    __UpperCamelCase  		=						XCLIPModel(lowercase_      )
    model.eval()

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

    __UpperCamelCase  		=						convert_state_dict(lowercase_					,     lowercase_      )

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

    __UpperCamelCase  		=						3_36 if model_name == '''xclip-large-patch14-16-frames''' else 2_24
    __UpperCamelCase  		=						VideoMAEImageProcessor(size=lowercase_      )
    __UpperCamelCase  		=						CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''      )
    __UpperCamelCase  		=						CLIPTokenizerFast.from_pretrained('''openai/clip-vit-base-patch32'''      )
    __UpperCamelCase  		=						XCLIPProcessor(image_processor=lowercase_					,     tokenizer=lowercase_      )

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

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

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

    # Verify outputs
    __UpperCamelCase  		=						outputs.logits_per_video
    __UpperCamelCase  		=						logits_per_video.softmax(dim=1      )
    print('''Probs:'''					,     lowercase_      )
    # 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.0_999E-04, 9.9_883E-01, 4.5_580E-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.6_937E-04, 9.9_728E-01, 1.9_473E-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.3_877E-04, 9.9_937E-01, 2.8_888E-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.8_554E-04, 9.9_929E-01, 3.2_754E-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.1_890E-06, 9.9_994E-01, 5.6_559E-05]]      )
    elif model_name == "xclip-base-patch16-hmdb-4-shot":
        __UpperCamelCase  		=						torch.tensor([[1.0_320E-05, 9.9_993E-01, 6.2_435E-05]]      )
    elif model_name == "xclip-base-patch16-hmdb-8-shot":
        __UpperCamelCase  		=						torch.tensor([[4.1_377E-06, 9.9_990E-01, 9.8_386E-05]]      )
    elif model_name == "xclip-base-patch16-hmdb-16-shot":
        __UpperCamelCase  		=						torch.tensor([[4.1_347E-05, 9.9_962E-01, 3.3_411E-04]]      )
    elif model_name == "xclip-base-patch16-ucf-2-shot":
        __UpperCamelCase  		=						torch.tensor([[8.5_857E-05, 9.9_928E-01, 6.3_291E-04]]      )
    elif model_name == "xclip-base-patch16-ucf-4-shot":
        __UpperCamelCase  		=						torch.tensor([[8.5_857E-05, 9.9_928E-01, 6.3_291E-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.8_219E-04, 9.9_593E-01, 3.0_863E-03]]      )
    # zero shot
    elif model_name == "xclip-base-patch16-zero-shot":
        __UpperCamelCase  		=						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(lowercase_					,     lowercase_					,     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(lowercase_      )

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


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

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











def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_ = False      )       ->	str:

    """simple docstring"""


    if not isinstance(lowercase_					,     lowercase_      ):
        __UpperCamelCase  		=						F"Expected string as input, found {type(lowercase_      )}"
        raise ValueError(lowercase_      )
    if not isinstance(lowercase_					,     lowercase_      ):
        __UpperCamelCase  		=						F"Expected boolean as use_pascal parameter, found {type(lowercase_      )}"
        raise ValueError(lowercase_      )

    __UpperCamelCase  		=						input_str.split('''_'''      )

    __UpperCamelCase  		=						0 if use_pascal else 1

    __UpperCamelCase  		=						words[start_index:]

    __UpperCamelCase  		=						[word[0].upper() + word[1:] for word in words_to_capitalize]

    __UpperCamelCase  		=						'''''' if use_pascal else words[0]

    return "".join([initial_word, *capitalized_words]      )


if __name__ == "__main__":
 from doctest import testmod

 testmod()






