



'''simple docstring'''





import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)




def        lowerCamelCase  (				_snake_case  :						Optional[Any]     ,_snake_case  :						int=False       ):





    '''simple docstring'''


    lowercase__		       =						[]
    for i in range(config.num_hidden_layers       ):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((f'''module.blocks.{i}.norm1.weight''', f'''vit.encoder.layer.{i}.layernorm_before.weight''')       )
        rename_keys.append((f'''module.blocks.{i}.norm1.bias''', f'''vit.encoder.layer.{i}.layernorm_before.bias''')       )
        rename_keys.append(
            (f'''module.blocks.{i}.attn.proj.weight''', f'''vit.encoder.layer.{i}.attention.output.dense.weight''')       )
        rename_keys.append((f'''module.blocks.{i}.attn.proj.bias''', f'''vit.encoder.layer.{i}.attention.output.dense.bias''')       )
        rename_keys.append((f'''module.blocks.{i}.norm2.weight''', f'''vit.encoder.layer.{i}.layernorm_after.weight''')       )
        rename_keys.append((f'''module.blocks.{i}.norm2.bias''', f'''vit.encoder.layer.{i}.layernorm_after.bias''')       )
        rename_keys.append((f'''module.blocks.{i}.mlp.fc1.weight''', f'''vit.encoder.layer.{i}.intermediate.dense.weight''')       )
        rename_keys.append((f'''module.blocks.{i}.mlp.fc1.bias''', f'''vit.encoder.layer.{i}.intermediate.dense.bias''')       )
        rename_keys.append((f'''module.blocks.{i}.mlp.fc2.weight''', f'''vit.encoder.layer.{i}.output.dense.weight''')       )
        rename_keys.append((f'''module.blocks.{i}.mlp.fc2.bias''', f'''vit.encoder.layer.{i}.output.dense.bias''')       )

    # projection layer + position embeddings
    rename_keys.extend(
        [
            ("module.cls_token", "vit.embeddings.cls_token"),
            ("module.patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
            ("module.patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
            ("module.pos_embed", "vit.embeddings.position_embeddings"),
        ]       )

    if base_model:
        # layernorm + pooler
        rename_keys.extend(
            [
                ("module.norm.weight", "layernorm.weight"),
                ("module.norm.bias", "layernorm.bias"),
            ]       )

        # if just the base model, we should remove "vit" from all keys that start with "vit"
        lowercase__		       =						[(pair[0], pair[1][4:]) if pair[1].startswith("vit"       ) else pair for pair in rename_keys]
    else:
        # layernorm + classification head
        rename_keys.extend(
            [
                ("norm.weight", "vit.layernorm.weight"),
                ("norm.bias", "vit.layernorm.bias"),
                ("head.weight", "classifier.weight"),
                ("head.bias", "classifier.bias"),
            ]       )

    return rename_keys




def        lowerCamelCase  (				_snake_case  :						int     ,_snake_case  :						Any     ,_snake_case  :						Union[str, Any]=False       ):





    '''simple docstring'''


    for i in range(config.num_hidden_layers       ):
        if base_model:
            lowercase__		       =						""
        else:
            lowercase__		       =						"vit."
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        lowercase__		       =						state_dict.pop(f'''module.blocks.{i}.attn.qkv.weight'''       )
        lowercase__		       =						state_dict.pop(f'''module.blocks.{i}.attn.qkv.bias'''       )
        # next, add query, keys and values (in that order) to the state dict
        lowercase__		       =						in_proj_weight[
            : config.hidden_size, :
        ]
        lowercase__		       =						in_proj_bias[: config.hidden_size]
        lowercase__		       =						in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        lowercase__		       =						in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        lowercase__		       =						in_proj_weight[
            -config.hidden_size :, :
        ]
        lowercase__		       =						in_proj_bias[-config.hidden_size :]




def        lowerCamelCase  (				_snake_case  :						Optional[int]       ):





    '''simple docstring'''


    lowercase__		       =						["head.weight", "head.bias"]
    for k in ignore_keys:
        state_dict.pop(_snake_case     ,_snake_case       )




def        lowerCamelCase  (				_snake_case  :						List[str]       ):





    '''simple docstring'''


    lowercase__		       =						[
        "module.fc.fc1.weight",
        "module.fc.fc1.bias",
        "module.fc.bn1.weight",
        "module.fc.bn1.bias",
        "module.fc.bn1.running_mean",
        "module.fc.bn1.running_var",
        "module.fc.bn1.num_batches_tracked",
        "module.fc.fc2.weight",
        "module.fc.fc2.bias",
        "module.fc.bn2.weight",
        "module.fc.bn2.bias",
        "module.fc.bn2.running_mean",
        "module.fc.bn2.running_var",
        "module.fc.bn2.num_batches_tracked",
        "module.fc.fc3.weight",
        "module.fc.fc3.bias",
    ]
    for k in ignore_keys:
        state_dict.pop(_snake_case     ,_snake_case       )




def        lowerCamelCase  (				_snake_case  :						Any     ,_snake_case  :						Union[str, Any]     ,_snake_case  :						int       ):





    '''simple docstring'''


    lowercase__		       =						dct.pop(_snake_case       )
    lowercase__		       =						val




def        lowerCamelCase  (				_snake_case  :						Optional[int]     ,_snake_case  :						int       ):





    '''simple docstring'''


    lowercase__		       =						ViTMSNConfig()
    lowercase__		       =						1_000

    lowercase__		       =						"datasets/huggingface/label-files"
    lowercase__		       =						"imagenet-1k-id2label.json"
    lowercase__		       =						json.load(open(hf_hub_download(_snake_case     ,_snake_case       )     ,"r"       )       )
    lowercase__		       =						{int(_snake_case       ): v for k, v in idalabel.items()}
    lowercase__		       =						idalabel
    lowercase__		       =						{v: k for k, v in idalabel.items()}

    if "s16" in checkpoint_url:
        lowercase__		       =						384
        lowercase__		       =						1_536
        lowercase__		       =						6
    elif "l16" in checkpoint_url:
        lowercase__		       =						1_024
        lowercase__		       =						4_096
        lowercase__		       =						24
        lowercase__		       =						16
        lowercase__		       =						0.1
    elif "b4" in checkpoint_url:
        lowercase__		       =						4
    elif "l7" in checkpoint_url:
        lowercase__		       =						7
        lowercase__		       =						1_024
        lowercase__		       =						4_096
        lowercase__		       =						24
        lowercase__		       =						16
        lowercase__		       =						0.1

    lowercase__		       =						ViTMSNModel(_snake_case       )

    lowercase__		       =						torch.hub.load_state_dict_from_url(_snake_case     ,map_location="cpu"       )["target_encoder"]

    lowercase__		       =						ViTImageProcessor(size=config.image_size       )

    remove_projection_head(_snake_case       )
    lowercase__		       =						create_rename_keys(_snake_case     ,base_model=_snake_case       )

    for src, dest in rename_keys:
        rename_key(_snake_case     ,_snake_case     ,_snake_case       )
    read_in_q_k_v(_snake_case     ,_snake_case     ,base_model=_snake_case       )

    model.load_state_dict(_snake_case       )
    model.eval()

    lowercase__		       =						"http://images.cocodataset.org/val2017/000000039769.jpg"

    lowercase__		       =						Image.open(requests.get(_snake_case     ,stream=_snake_case       ).raw       )
    lowercase__		       =						ViTImageProcessor(
        size=config.image_size     ,image_mean=_snake_case     ,image_std=_snake_case       )
    lowercase__		       =						image_processor(images=_snake_case     ,return_tensors="pt"       )

    # forward pass
    torch.manual_seed(2       )
    lowercase__		       =						model(**_snake_case       )
    lowercase__		       =						outputs.last_hidden_state

    # The following Colab Notebook was used to generate these outputs:
    # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
    if "s16" in checkpoint_url:
        lowercase__		       =						torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]       )
    elif "b16" in checkpoint_url:
        lowercase__		       =						torch.tensor([[14.2_889, -18.9_045, 11.7_281]]       )
    elif "l16" in checkpoint_url:
        lowercase__		       =						torch.tensor([[41.5_028, -22.8_681, 45.6_475]]       )
    elif "b4" in checkpoint_url:
        lowercase__		       =						torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]       )
    else:
        lowercase__		       =						torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]       )

    # verify logits
    assert torch.allclose(last_hidden_state[:, 0, :3]     ,_snake_case     ,atol=1e-4       )

    print(f'''Saving model to {pytorch_dump_folder_path}'''       )
    model.save_pretrained(_snake_case       )

    print(f'''Saving image processor to {pytorch_dump_folder_path}'''       )
    image_processor.save_pretrained(_snake_case       )


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE__      					=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--checkpoint_url",
     default="https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar",
     type=str,
     help="URL of the checkpoint you'd like to convert.",
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
 )

 SCREAMING_SNAKE_CASE__      					=     parser.parse_args()
 convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''





import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

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


enable_full_determinism()


@skip_mps
class 				snake_case (UpperCamelCase		,      unittest.TestCase   ):
       lowerCAmelCase__  :int						   =					TextToVideoSDPipeline
       lowerCAmelCase__  :Union[str, Any]						   =					TEXT_TO_IMAGE_PARAMS
       lowerCAmelCase__  :List[str]						   =					TEXT_TO_IMAGE_BATCH_PARAMS
       # No `output_type`.
       lowerCAmelCase__  :Optional[Any]						   =					frozenset(
           [
               "num_inference_steps",
               "generator",
               "latents",
               "return_dict",
               "callback",
               "callback_steps",
           ]   )



       def 							_a      (				self      )  ->							Optional[int]:
           torch.manual_seed(0      )
           lowercase__		       =						UNetaDConditionModel(
               block_out_channels=(32, 64, 64, 64)     ,layers_per_block=2     ,sample_size=32     ,in_channels=4     ,out_channels=4     ,down_block_types=("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")     ,up_block_types=("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")     ,cross_attention_dim=32     ,attention_head_dim=4     ,)
           lowercase__		       =						DDIMScheduler(
               beta_start=0.0_00_85     ,beta_end=0.0_12     ,beta_schedule="scaled_linear"     ,clip_sample=UpperCAmelCase_     ,set_alpha_to_one=UpperCAmelCase_     ,)
           torch.manual_seed(0      )
           lowercase__		       =						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     ,sample_size=128     ,)
           torch.manual_seed(0      )
           lowercase__		       =						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=1_000     ,hidden_act="gelu"     ,projection_dim=512     ,)
           lowercase__		       =						CLIPTextModel(UpperCAmelCase_      )
           lowercase__		       =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

           lowercase__		       =						{
               "unet": unet,
               "scheduler": scheduler,
               "vae": vae,
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
           }
           return components



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_=0      )  ->							int:
           if str(UpperCAmelCase_      ).startswith("mps"      ):
               lowercase__		       =						torch.manual_seed(UpperCAmelCase_      )
           else:
               lowercase__		       =						torch.Generator(device=UpperCAmelCase_      ).manual_seed(UpperCAmelCase_      )
           lowercase__		       =						{
               "prompt": "A painting of a squirrel eating a burger",
               "generator": generator,
               "num_inference_steps": 2,
               "guidance_scale": 6.0,
               "output_type": "pt",
           }
           return inputs



       def 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowercase__		       =						self.get_dummy_components()
           lowercase__		       =						TextToVideoSDPipeline(**UpperCAmelCase_      )
           lowercase__		       =						sd_pipe.to(UpperCAmelCase_      )
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

           lowercase__		       =						self.get_dummy_inputs(UpperCAmelCase_      )
           lowercase__		       =						"np"
           lowercase__		       =						sd_pipe(**UpperCAmelCase_      ).frames
           lowercase__		       =						frames[0][-3:, -3:, -1]

           assert frames[0].shape == (64, 64, 3)
           lowercase__		       =						np.array([1_58.0, 1_60.0, 1_53.0, 1_25.0, 1_00.0, 1_21.0, 1_11.0, 93.0, 1_13.0]      )

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



       def 							_a      (				self      )  ->							List[Any]:
           self._test_attention_slicing_forward_pass(test_mean_pixel_difference=UpperCAmelCase_     ,expected_max_diff=3E-3      )



       @unittest.skipIf(
           torch_device != "cuda" or not is_xformers_available()     ,reason="XFormers attention is only available with CUDA and `xformers` installed"     ,)
       def 							_a      (				self      )  ->							Dict:
           self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=UpperCAmelCase_     ,expected_max_diff=1E-2      )



       @unittest.skip(reason="Batching needs to be properly figured out first for this pipeline."      )
       def 							_a      (				self      )  ->							Union[str, Any]:
           pass



       @unittest.skip(reason="Batching needs to be properly figured out first for this pipeline."      )
       def 							_a      (				self      )  ->							Union[str, Any]:
           pass



       @unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline."      )
       def 							_a      (				self      )  ->							Union[str, Any]:
           pass







       def 							_a      (				self      )  ->							str:
           return super().test_progress_bar()





@slow
@skip_mps
class 				snake_case (unittest.TestCase   ):



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy"      )

           lowercase__		       =						TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b"      )
           lowercase__		       =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
           lowercase__		       =						pipe.to("cuda"      )

           lowercase__		       =						"Spiderman is surfing"
           lowercase__		       =						torch.Generator(device="cpu"      ).manual_seed(0      )

           lowercase__		       =						pipe(UpperCAmelCase_     ,generator=UpperCAmelCase_     ,num_inference_steps=25     ,output_type="pt"      ).frames
           lowercase__		       =						video_frames.cpu().numpy()

           assert np.abs(expected_video - video      ).mean() < 5E-2







       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy"      )

           lowercase__		       =						TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b"      )
           lowercase__		       =						pipe.to("cuda"      )

           lowercase__		       =						"Spiderman is surfing"
           lowercase__		       =						torch.Generator(device="cpu"      ).manual_seed(0      )

           lowercase__		       =						pipe(UpperCAmelCase_     ,generator=UpperCAmelCase_     ,num_inference_steps=2     ,output_type="pt"      ).frames
           lowercase__		       =						video_frames.cpu().numpy()

           assert np.abs(expected_video - video      ).mean() < 5E-2




