





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      lowerCamelCase(						__snake_case			, unittest.TestCase	):
     '''simple docstring'''


     __magic_name__   =			TextToVideoSDPipeline
     __magic_name__   =			TEXT_TO_IMAGE_PARAMS
     __magic_name__   =			TEXT_TO_IMAGE_BATCH_PARAMS
     # No `output_type`.
     __magic_name__   =			frozenset(
         [
             'num_inference_steps',
             'generator',
             'latents',
             'return_dict',
             'callback',
             'callback_steps',
         ]	)




     def        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)
         _A     =							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				,     )
         _A     =							DDIMScheduler(
             beta_start=0.0_0085				,     beta_end=0.012				,     beta_schedule='scaled_linear'				,     clip_sample=snake_case_				,     set_alpha_to_one=snake_case_				,     )
         torch.manual_seed(0							)
         _A     =							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							)
         _A     =							CLIPTextConfig(
             bos_token_id=0				,     eos_token_id=2				,     hidden_size=32				,     intermediate_size=37				,     layer_norm_eps=1E-05				,     num_attention_heads=4				,     num_hidden_layers=5				,     pad_token_id=1				,     vocab_size=1000				,     hidden_act='gelu'				,     projection_dim=512				,     )
         _A     =							CLIPTextModel(snake_case_							)
         _A     =							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

         _A     =							{
             'unet': unet,
             'scheduler': scheduler,
             'vae': vae,
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
         }
         return components




     def        lowerCAmelCase__		(    self				,     snake_case_				,     snake_case_=0							):
         if str(snake_case_							).startswith('mps'							):
             _A     =							torch.manual_seed(snake_case_							)
         else:
             _A     =							torch.Generator(device=snake_case_							).manual_seed(snake_case_							)
         _A     =							{
             '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        lowerCAmelCase__		(    self							):
         _A     =							'cpu'  # ensure determinism for the device-dependent torch.Generator
         _A     =							self.get_dummy_components()
         _A     =							TextToVideoSDPipeline(**snake_case_							)
         _A     =							sd_pipe.to(snake_case_							)
         sd_pipe.set_progress_bar_config(disable=snake_case_							)

         _A     =							self.get_dummy_inputs(snake_case_							)
         _A     =							'np'
         _A     =							sd_pipe(**snake_case_							).frames
         _A     =							frames[0][-3:, -3:, -1]

         assert frames[0].shape == (64, 64, 3)
         _A     =							np.array([158.0, 160.0, 153.0, 125.0, 100.0, 121.0, 111.0, 93.0, 113.0]							)

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




     def        lowerCAmelCase__		(    self							):
         self._test_attention_slicing_forward_pass(test_mean_pixel_difference=snake_case_				,     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        lowerCAmelCase__		(    self							):
         self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=snake_case_				,     expected_max_diff=1E-2							)




     @unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'							)
     def        lowerCAmelCase__		(    self							):
         pass




     @unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'							)
     def        lowerCAmelCase__		(    self							):
         pass




     @unittest.skip(reason='`num_images_per_prompt` argument is not supported for this pipeline.'							)
     def        lowerCAmelCase__		(    self							):
         pass




     def        lowerCAmelCase__		(    self							):
         return super().test_progress_bar()







@slow
@skip_mps
class      lowerCamelCase(						unittest.TestCase	):
     '''simple docstring'''


     def        lowerCAmelCase__		(    self							):
         _A     =							load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy'							)

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

         _A     =							'Spiderman is surfing'
         _A     =							torch.Generator(device='cpu'							).manual_seed(0							)

         _A     =							pipe(snake_case_				,     generator=snake_case_				,     num_inference_steps=25				,     output_type='pt'							).frames
         _A     =							video_frames.cpu().numpy()

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




     def        lowerCAmelCase__		(    self							):
         _A     =							load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy'							)

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

         _A     =							'Spiderman is surfing'
         _A     =							torch.Generator(device='cpu'							).manual_seed(0							)

         _A     =							pipe(snake_case_				,     generator=snake_case_				,     num_inference_steps=2				,     output_type='pt'							).frames
         _A     =							video_frames.cpu().numpy()

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













import argparse
import json
from pathlib import Path

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

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase__							  =						logging.get_logger(__name__)
def 		_a       ( a					:List[str]      ,				a					:Tuple=False					)				->     List[Any]:
 a						    = []
 for i in range(config.num_hidden_layers					):
  # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
  rename_keys.append((F"""blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""")					)
  rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")					)
  rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")					)
  rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")					)
  rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")					)
  rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")					)
  rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")					)
  rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")					)
  rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")					)
  rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")					)

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

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

  # if just the base model, we should remove "vit" from all keys that start with "vit"
  a						    = [(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 		_a       ( a					:List[str]      ,				a					:int      ,				a					:Tuple=False					)				->     Any:
 for i in range(config.num_hidden_layers					):
  if base_model:
   a						    = ''''''
  else:
   a						    = '''vit.'''
  # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
  a						    = state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""					)
  a						    = state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""					)
  # next, add query, keys and values (in that order) to the state dict
  a						    = in_proj_weight[
      : config.hidden_size, :
  ]
  a						    = in_proj_bias[: config.hidden_size]
  a						    = in_proj_weight[
      config.hidden_size : config.hidden_size * 2, :
  ]
  a						    = in_proj_bias[
      config.hidden_size : config.hidden_size * 2
  ]
  a						    = in_proj_weight[
      -config.hidden_size :, :
  ]
  a						    = in_proj_bias[-config.hidden_size :]
def 		_a       ( a					:List[Any]					)				->     Dict:
 a						    = ['''head.weight''', '''head.bias''']
 for k in ignore_keys:
  state_dict.pop(a      ,				a					)
def 		_a       ( a					:Union[str, Any]      ,				a					:str      ,				a					:List[Any]					)				->     str:
 a						    = dct.pop(a					)
 a						    = val
def 		_a       ( )				->     Optional[Any]:
 a						    = '''http://images.cocodataset.org/val2017/000000039769.jpg'''
 a						    = Image.open(requests.get(a      ,				stream=a					).raw					)
 return im
@torch.no_grad()
def 		_a       ( a					:Tuple      ,				a					:str      ,				a					:Optional[int]=True					)				->     Dict:
 a						    = ViTConfig()
 # patch_size
 if model_name[-1] == "8":
  a						    = 8
 # set labels if required
 if not base_model:
  a						    = 1_000
  a						    = '''huggingface/label-files'''
  a						    = '''imagenet-1k-id2label.json'''
  a						    = json.load(open(hf_hub_download(a      ,				a      ,				repo_type='''dataset'''					)      ,				'''r'''					)					)
  a						    = {int(a					): v for k, v in idalabel.items()}
  a						    = idalabel
  a						    = {v: k for k, v in idalabel.items()}
 # size of the architecture
 if model_name in ["dino_vits8", "dino_vits16"]:
  a						    = 384
  a						    = 1_536
  a						    = 12
  a						    = 6

 # load original model from torch hub
 a						    = torch.hub.load('''facebookresearch/dino:main'''      ,				a					)
 original_model.eval()

 # load state_dict of original model, remove and rename some keys
 a						    = original_model.state_dict()
 if base_model:
  remove_classification_head_(a					)
 a						    = create_rename_keys(a      ,				base_model=a					)
 for src, dest in rename_keys:
  rename_key(a      ,				a      ,				a					)
 read_in_q_k_v(a      ,				a      ,				a					)

 # load HuggingFace model
 if base_model:
  a						    = ViTModel(a      ,				add_pooling_layer=a					).eval()
 else:
  a						    = ViTForImageClassification(a					).eval()
 model.load_state_dict(a					)

 # Check outputs on an image, prepared by ViTImageProcessor
 a						    = ViTImageProcessor()
 a						    = image_processor(images=prepare_img()      ,				return_tensors='''pt'''					)
 a						    = encoding['''pixel_values''']
 a						    = model(a					)

 if base_model:
  a						    = original_model(a					)
  assert torch.allclose(a      ,				outputs.last_hidden_state[:, 0, :]      ,				atol=1e-1					)
 else:
  a						    = original_model(a					)
  assert logits.shape == outputs.logits.shape
  assert torch.allclose(a      ,				outputs.logits      ,				atol=1e-3					)

 Path(a					).mkdir(exist_ok=a					)
 print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""					)
 model.save_pretrained(a					)
 print(F"""Saving image processor to {pytorch_dump_folder_path}"""					)
 image_processor.save_pretrained(a					)


if __name__ == "__main__":
     UpperCAmelCase__							  =						argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         "--model_name",
         default="dino_vitb16",
         type=str,
         help="Name of the model trained with DINO you'd like to convert.",
     )
     parser.add_argument(
         "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
     )
     parser.add_argument(
         "--base_model",
         action="store_true",
         help="Whether to only convert the base model (no projection head weights).",
     )

     parser.set_defaults(base_model=True)
     UpperCAmelCase__							  =						parser.parse_args()
     convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)


