code: 
  import argparse
  import json
  import os.path
  from collections import OrderedDict
  
  import numpy as np
  import requests
  import torch
  from flax.training.checkpoints import restore_checkpoint
  from huggingface_hub import hf_hub_download
  
  from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
  from transformers.image_utils import PILImageResampling
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      Optional[int]     )      -> Tuple:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  Optional[Any]    			=      "https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"
  
  					with open(_UpperCAmelCase   , "wb"     ) as f:
  										with requests.get(_UpperCAmelCase   , stream=_UpperCAmelCase     ) as req:
  															for chunk in req.iter_content(chunk_size=2048     ):
  																				f.write(_UpperCAmelCase     )
  
  
  
  def 		UpperCamelCase    (					)      -> VivitConfig:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  Tuple    			=      VivitConfig()
  
  					_lowercase					:  Dict    			=      400
  					_lowercase					:  Dict    			=      "huggingface/label-files"
  					_lowercase					:  str    			=      "kinetics400-id2label.json"
  
  					_lowercase					:  Union[str, Any]    			=      json.load(open(hf_hub_download(_UpperCAmelCase   , _UpperCAmelCase   , repo_type="dataset"     )   , "r"     )     )
  					_lowercase					:  str    			=      {int(_UpperCAmelCase     ): v for k, v in idalabel.items()}
  					_lowercase					:  Optional[Any]    			=      idalabel
  					_lowercase					:  int    			=      {v: k for k, v in idalabel.items()}
  					return config
  
  
  
  def 		UpperCamelCase    (					)      -> Any:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  Union[str, Any]    			=      hf_hub_download(
  					    repo_id="hf-internal-testing/spaghetti-video"   , filename="eating_spaghetti_32_frames.npy"   , repo_type="dataset"     )
  					_lowercase					:  Any    			=      np.load(_UpperCAmelCase     )
  					return list(_UpperCAmelCase     )
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Optional[int]:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					if np.ndim(_UpperCAmelCase     ) == 2:
  										return transform_attention_bias(_UpperCAmelCase     )
  
  					elif np.ndim(_UpperCAmelCase     ) == 3:
  										return transform_attention_kernel(_UpperCAmelCase     )
  
  					else:
  										raise Exception(f"""Invalid number of dimesions: {np.ndim(_UpperCAmelCase     )}"""     )
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Tuple:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					return current.flatten()
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Dict:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					return np.reshape(_UpperCAmelCase   , (current.shape[0], current.shape[1] * current.shape[2])     ).T
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      np.ndarray     )      -> Optional[Any]:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					return np.reshape(_UpperCAmelCase   , (current.shape[0] * current.shape[1], current.shape[2])     ).T
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      Union[str, Any]   , _UpperCAmelCase     :      str     )      -> Optional[Any]:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  int    			=      state_dict["optimizer"]["target"]["Transformer"][f"""encoderblock_{i}"""]
  
  					_lowercase					:  List[Any]    			=      f"""encoder.layer.{i}."""
  					_lowercase					:  Tuple    			=      {
  					    prefix + "intermediate.dense.bias": state["MlpBlock_0"]["Dense_0"]["bias"],
  					    prefix + "intermediate.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_0"]["kernel"]     ),
  					    prefix + "output.dense.bias": state["MlpBlock_0"]["Dense_1"]["bias"],
  					    prefix + "output.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_1"]["kernel"]     ),
  					    prefix + "layernorm_before.bias": state["LayerNorm_0"]["bias"],
  					    prefix + "layernorm_before.weight": state["LayerNorm_0"]["scale"],
  					    prefix + "layernorm_after.bias": state["LayerNorm_1"]["bias"],
  					    prefix + "layernorm_after.weight": state["LayerNorm_1"]["scale"],
  					    prefix
  					    + "attention.attention.query.bias": transform_attention(
  					        state["MultiHeadDotProductAttention_0"]["query"]["bias"]     ),
  					    prefix
  					    + "attention.attention.query.weight": transform_attention(
  					        state["MultiHeadDotProductAttention_0"]["query"]["kernel"]     ),
  					    prefix
  					    + "attention.attention.key.bias": transform_attention(state["MultiHeadDotProductAttention_0"]["key"]["bias"]     ),
  					    prefix
  					    + "attention.attention.key.weight": transform_attention(
  					        state["MultiHeadDotProductAttention_0"]["key"]["kernel"]     ),
  					    prefix
  					    + "attention.attention.value.bias": transform_attention(
  					        state["MultiHeadDotProductAttention_0"]["value"]["bias"]     ),
  					    prefix
  					    + "attention.attention.value.weight": transform_attention(
  					        state["MultiHeadDotProductAttention_0"]["value"]["kernel"]     ),
  					    prefix + "attention.output.dense.bias": state["MultiHeadDotProductAttention_0"]["out"]["bias"],
  					    prefix
  					    + "attention.output.dense.weight": transform_attention_output_weight(
  					        state["MultiHeadDotProductAttention_0"]["out"]["kernel"]     ),
  					}
  
  					return new_state
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      Optional[Any]     )      -> List[str]:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]     )
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      List[str]   , _UpperCAmelCase     :      Tuple=False     )      -> Tuple:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  Optional[Any]    			=      get_n_layers(_UpperCAmelCase     )
  
  					_lowercase					:  Dict    			=      OrderedDict()
  
  					_lowercase					:  Dict    			=      state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
  					_lowercase					:  Union[str, Any]    			=      state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]
  
  					_lowercase					:  Optional[Any]    			=      np.transpose(
  					    state_dict["optimizer"]["target"]["embedding"]["kernel"]   , (4, 3, 0, 1, 2)     )
  					_lowercase					:  Tuple    			=      state_dict["optimizer"]["target"]["embedding"]["bias"]
  
  					_lowercase					:  Tuple    			=      state_dict["optimizer"]["target"]["cls"]
  					_lowercase					:  int    			=      state_dict["optimizer"]["target"]["Transformer"]["posembed_input"][
  					    "pos_embedding"
  					]
  
  					for i in range(_UpperCAmelCase     ):
  										new_state.update(transform_state_encoder_block(_UpperCAmelCase   , _UpperCAmelCase     )     )
  
  					if classification_head:
  										_lowercase					:  Optional[Any]    			=      {"vivit." + k: v for k, v in new_state.items()}
  										_lowercase					:  Optional[Any]    			=      np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"]     )
  										_lowercase					:  Union[str, Any]    			=      np.transpose(state_dict["optimizer"]["target"]["output_projection"]["bias"]     )
  
  					return {k: torch.tensor(_UpperCAmelCase     ) for k, v in new_state.items()}
  
  
  
  def 		UpperCamelCase    (					)      -> VivitImageProcessor:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  Union[str, Any]    			=      VivitImageProcessor()
  
  					assert extractor.do_resize is True
  					assert extractor.size == {"shortest_edge": 256}
  					assert extractor.do_center_crop is True
  					assert extractor.crop_size == {"width": 224, "height": 224}
  					assert extractor.resample == PILImageResampling.BILINEAR
  
  					# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
  					# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
  					# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
  					assert extractor.do_normalize is False
  					assert extractor.do_rescale is True
  					assert extractor.rescale_factor == 1 / 255
  
  					# zero-centering = True in original implementation
  					assert extractor.do_zero_centering is True
  
  					return extractor
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      str     )      -> Tuple:
  
  
  
  
  
  
  					'''simple docstring'''
  
  					_lowercase					:  Any    			=      "checkpoint"
  
  					if not os.path.exists(_UpperCAmelCase     ):
  										download_checkpoint(_UpperCAmelCase     )
  
  					_lowercase					:  List[str]    			=      restore_checkpoint(_UpperCAmelCase   , _UpperCAmelCase     )
  					_lowercase					:  Union[str, Any]    			=      transform_state(_UpperCAmelCase   , classification_head=_UpperCAmelCase     )
  
  					_lowercase					:  Dict    			=      get_vivit_config()
  
  					assert config.image_size == 224
  					assert config.num_frames == 32
  
  					_lowercase					:  int    			=      VivitForVideoClassification(_UpperCAmelCase     )
  					model.load_state_dict(_UpperCAmelCase     )
  					model.eval()
  
  					_lowercase					:  List[Any]    			=      get_processor()
  
  					_lowercase					:  Union[str, Any]    			=      prepare_video()
  					_lowercase					:  Optional[int]    			=      extractor(_UpperCAmelCase   , return_tensors="pt"     )
  
  					_lowercase					:  Union[str, Any]    			=      model(**_UpperCAmelCase     )
  
  					_lowercase					:  List[str]    			=      torch.Size([1, 400]     )
  					_lowercase					:  Optional[int]    			=      torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]     )
  
  					assert outputs.logits.shape == expected_shape
  					assert torch.allclose(outputs.logits[0, :5]   , _UpperCAmelCase   , atol=1e-4     ), outputs.logits[0, :5]
  
  					model.save_pretrained(_UpperCAmelCase     )
  					extractor.save_pretrained(_UpperCAmelCase     )
  
  
  if __name__ == "__main__":
  							UpperCamelCase_	:    Tuple  					=       argparse.ArgumentParser()
  
  							parser.add_argument("""--output_model_name""", """-o""", type=str, help="""Output path for the converted HuggingFace model""")
  
  							UpperCamelCase_	:    int  					=       parser.parse_args()
  							convert(args.output_model_name)
  
  
  
  
  
style_context: 
  import gc
  import os
  import tempfile
  import unittest
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
  
  from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
  from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
  from diffusers.models.attention_processor import (
      Attention,
      AttnProcessor,
      AttnProcessor2_0,
      LoRAAttnProcessor,
      LoRAAttnProcessor2_0,
      LoRAXFormersAttnProcessor,
      XFormersAttnProcessor,
  )
  from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      nn.Module     )      -> Any:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  Optional[int]    			=      {}
  						for name in unet.attn_processors.keys():
  												_lowercase					:  Any    			=      None if name.endswith("attn1.processor"     ) else unet.config.cross_attention_dim
  												if name.startswith("mid_block"     ):
  																		_lowercase					:  Any    			=      unet.config.block_out_channels[-1]
  												elif name.startswith("up_blocks"     ):
  																		_lowercase					:  List[Any]    			=      int(name[len("up_blocks."     )]     )
  																		_lowercase					:  Any    			=      list(reversed(unet.config.block_out_channels     )     )[block_id]
  												elif name.startswith("down_blocks"     ):
  																		_lowercase					:  Union[str, Any]    			=      int(name[len("down_blocks."     )]     )
  																		_lowercase					:  Dict    			=      unet.config.block_out_channels[block_id]
  												_lowercase					:  List[str]    			=      (
  												    LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase   , "scaled_dot_product_attention"     ) else LoRAAttnProcessor
  												)
  												_lowercase					:  Optional[Any]    			=      lora_attn_processor_class(
  												    hidden_size=_UpperCAmelCase   , cross_attention_dim=_UpperCAmelCase     )
  						_lowercase					:  Dict    			=      AttnProcsLayers(_UpperCAmelCase     )
  						return lora_attn_procs, unet_lora_layers
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      nn.Module     )      -> Any:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  List[str]    			=      {}
  						_lowercase					:  Dict    			=      (
  						    LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase   , "scaled_dot_product_attention"     ) else LoRAAttnProcessor
  						)
  						for name, module in text_encoder.named_modules():
  												if name.endswith(_UpperCAmelCase     ):
  																		_lowercase					:  Tuple    			=      lora_attn_processor_class(
  																		    hidden_size=module.out_proj.out_features   , cross_attention_dim=_UpperCAmelCase     )
  						return text_lora_attn_procs
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      nn.Module     )      -> Union[str, Any]:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						_lowercase					:  Union[str, Any]    			=      create_text_encoder_lora_attn_procs(_UpperCAmelCase     )
  						_lowercase					:  str    			=      AttnProcsLayers(_UpperCAmelCase     )
  						return text_encoder_lora_layers
  
  
  
  def 		UpperCamelCase    (					_UpperCAmelCase     :      str   , _UpperCAmelCase     :      int=False     )      -> List[Any]:
  
  
  
  
  
  
  						'''simple docstring'''
  
  						for _, attn_proc in text_lora_attn_procs.items():
  												# set up.weights
  												for layer_name, layer_module in attn_proc.named_modules():
  																		if layer_name.endswith("_lora"     ):
  																								_lowercase					:  Tuple    			=      (
  																								    torch.randn_like(layer_module.up.weight     )
  																								    if randn_weight
  																								    else torch.zeros_like(layer_module.up.weight     )
  																								)
  																								_lowercase					:  Dict    			=      torch.nn.Parameter(_UpperCAmelCase     )
  
  
  
  
  class      __lowercase	(	unittest.TestCase					):
  
  
  
  
  
  
  	def 			_a(self							:   str						)   -> Tuple:
  							torch.manual_seed(0						)
  							_lowercase					:  int    			=      UNetaDConditionModel(
  							    block_out_channels=(32, 64)			,  layers_per_block=2			,  sample_size=32			,  in_channels=4			,  out_channels=4			,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,  up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,  cross_attention_dim=32			,  )
  							_lowercase					:  Union[str, Any]    			=      DDIMScheduler(
  							    beta_start=0.0_00_85			,  beta_end=0.0_12			,  beta_schedule="scaled_linear"			,  clip_sample=snake_case			,  set_alpha_to_one=snake_case			,  steps_offset=1			,  )
  							torch.manual_seed(0						)
  							_lowercase					:  int    			=      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			,  )
  							_lowercase					:  Optional[int]    			=      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			,  )
  							_lowercase					:  List[Any]    			=      CLIPTextModel(snake_case						)
  							_lowercase					:  Dict    			=      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)
  
  							_lowercase			,						_lowercase					:  List[Any]    			=      create_unet_lora_layers(snake_case						)
  							_lowercase					:  Union[str, Any]    			=      create_text_encoder_lora_layers(snake_case						)
  
  							_lowercase					:  List[Any]    			=      {
  							    "unet": unet,
  							    "scheduler": scheduler,
  							    "vae": vae,
  							    "text_encoder": text_encoder,
  							    "tokenizer": tokenizer,
  							    "safety_checker": None,
  							    "feature_extractor": None,
  							}
  							_lowercase					:  Union[str, Any]    			=      {
  							    "unet_lora_layers": unet_lora_layers,
  							    "text_encoder_lora_layers": text_encoder_lora_layers,
  							    "unet_lora_attn_procs": unet_lora_attn_procs,
  							}
  							return pipeline_components, lora_components
  
  
  
  
  
  
  	def 			_a(self							:   List[Any]						)   -> Optional[Any]:
  							_lowercase					:  str    			=      1
  							_lowercase					:  str    			=      10
  							_lowercase					:  Any    			=      4
  							_lowercase					:  List[str]    			=      (32, 32)
  
  							_lowercase					:  Union[str, Any]    			=      torch.manual_seed(0						)
  							_lowercase					:  List[str]    			=      floats_tensor((batch_size, num_channels) + sizes						)
  							_lowercase					:  Dict    			=      torch.randint(1			,  snake_case			,  size=(batch_size, sequence_length)			,  generator=snake_case						)
  
  							_lowercase					:  Tuple    			=      {
  							    "prompt": "A painting of a squirrel eating a burger",
  							    "generator": generator,
  							    "num_inference_steps": 2,
  							    "guidance_scale": 6.0,
  							    "output_type": "numpy",
  							}
  
  							return noise, input_ids, pipeline_inputs
  
  							# copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb
  
  
  
  
  
  
  	def 			_a(self							:   Union[str, Any]						)   -> int:
  							_lowercase					:  Any    			=      77
  
  							_lowercase					:  Any    			=      torch.randint(2			,  56			,  size=(1, max_seq_length)			,  generator=torch.manual_seed(0						)						)
  
  							_lowercase					:  int    			=      {}
  							_lowercase					:  int    			=      inputs
  							return prepared_inputs
  
  
  
  
  
  
  	def 			_a(self							:   Union[str, Any]			,  snake_case							:   Dict						)   -> List[Any]:
  							_lowercase			,						_lowercase					:  List[Any]    			=      self.get_dummy_components()
  							LoraLoaderMixin.save_lora_weights(
  							    save_directory=snake_case			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
  							self.assertTrue(os.path.isfile(os.path.join(snake_case			,  "pytorch_lora_weights.bin"						)						)						)
  
  
  
  
  
  
  	def 			_a(self							:   Any						)   -> Union[str, Any]:
  							_lowercase			,						_lowercase					:  int    			=      self.get_dummy_components()
  							_lowercase					:  List[str]    			=      StableDiffusionPipeline(**snake_case						)
  							_lowercase					:  Dict    			=      sd_pipe.to(snake_case						)
  							sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  							_lowercase			,						_lowercase			,						_lowercase					:  Optional[int]    			=      self.get_dummy_inputs()
  
  							_lowercase					:  Any    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  int    			=      original_images[0, -3:, -3:, -1]
  
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													LoraLoaderMixin.save_lora_weights(
  													    save_directory=snake_case			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
  													self.assertTrue(os.path.isfile(os.path.join(snake_case			,  "pytorch_lora_weights.bin"						)						)						)
  													sd_pipe.load_lora_weights(snake_case						)
  
  							_lowercase					:  Dict    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  List[Any]    			=      lora_images[0, -3:, -3:, -1]
  
  							# Outputs shouldn't match.
  							self.assertFalse(torch.allclose(torch.from_numpy(snake_case						)			,  torch.from_numpy(snake_case						)						)						)
  
  
  
  
  
  
  	def 			_a(self							:   Optional[int]						)   -> List[Any]:
  							_lowercase			,						_lowercase					:  str    			=      self.get_dummy_components()
  							_lowercase					:  Any    			=      StableDiffusionPipeline(**snake_case						)
  							_lowercase					:  Any    			=      sd_pipe.to(snake_case						)
  							sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  							_lowercase			,						_lowercase			,						_lowercase					:  Optional[int]    			=      self.get_dummy_inputs()
  
  							_lowercase					:  Optional[int]    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  Any    			=      original_images[0, -3:, -3:, -1]
  
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													LoraLoaderMixin.save_lora_weights(
  													    save_directory=snake_case			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  safe_serialization=snake_case			,  )
  													self.assertTrue(os.path.isfile(os.path.join(snake_case			,  "pytorch_lora_weights.safetensors"						)						)						)
  													sd_pipe.load_lora_weights(snake_case						)
  
  							_lowercase					:  Any    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  Optional[int]    			=      lora_images[0, -3:, -3:, -1]
  
  							# Outputs shouldn't match.
  							self.assertFalse(torch.allclose(torch.from_numpy(snake_case						)			,  torch.from_numpy(snake_case						)						)						)
  
  
  
  
  
  
  	def 			_a(self							:   str						)   -> Union[str, Any]:
  							_lowercase			,						_lowercase					:  Optional[Any]    			=      self.get_dummy_components()
  							_lowercase					:  int    			=      lora_components["unet_lora_attn_procs"]
  							_lowercase					:  str    			=      StableDiffusionPipeline(**snake_case						)
  							_lowercase					:  List[str]    			=      sd_pipe.to(snake_case						)
  							sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  							_lowercase			,						_lowercase			,						_lowercase					:  Any    			=      self.get_dummy_inputs()
  
  							_lowercase					:  int    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  List[str]    			=      original_images[0, -3:, -3:, -1]
  
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													_lowercase					:  Optional[Any]    			=      sd_pipe.unet
  													unet.set_attn_processor(snake_case						)
  													unet.save_attn_procs(snake_case						)
  													self.assertTrue(os.path.isfile(os.path.join(snake_case			,  "pytorch_lora_weights.bin"						)						)						)
  													sd_pipe.load_lora_weights(snake_case						)
  
  							_lowercase					:  int    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  int    			=      lora_images[0, -3:, -3:, -1]
  
  							# Outputs shouldn't match.
  							self.assertFalse(torch.allclose(torch.from_numpy(snake_case						)			,  torch.from_numpy(snake_case						)						)						)
  
  
  
  
  
  
  	def 			_a(self							:   Dict						)   -> List[str]:
  							_lowercase			,						_lowercase					:  Optional[Any]    			=      self.get_dummy_components()
  							_lowercase					:  List[Any]    			=      StableDiffusionPipeline(**snake_case						)
  
  							_lowercase					:  Tuple    			=      self.get_dummy_tokens()
  
  							# inference without lora
  							_lowercase					:  int    			=      pipe.text_encoder(**snake_case						)[0]
  							assert outputs_without_lora.shape == (1, 77, 32)
  
  							# create lora_attn_procs with zeroed out up.weights
  							_lowercase					:  str    			=      create_text_encoder_lora_attn_procs(pipe.text_encoder						)
  							set_lora_up_weights(snake_case			,  randn_weight=snake_case						)
  
  							# monkey patch
  							pipe._modify_text_encoder(snake_case						)
  
  							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
  							del text_attn_procs
  							gc.collect()
  
  							# inference with lora
  							_lowercase					:  Dict    			=      pipe.text_encoder(**snake_case						)[0]
  							assert outputs_with_lora.shape == (1, 77, 32)
  
  							assert torch.allclose(
  							    snake_case			,  snake_case						), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"
  
  							# create lora_attn_procs with randn up.weights
  							_lowercase					:  Any    			=      create_text_encoder_lora_attn_procs(pipe.text_encoder						)
  							set_lora_up_weights(snake_case			,  randn_weight=snake_case						)
  
  							# monkey patch
  							pipe._modify_text_encoder(snake_case						)
  
  							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
  							del text_attn_procs
  							gc.collect()
  
  							# inference with lora
  							_lowercase					:  Any    			=      pipe.text_encoder(**snake_case						)[0]
  							assert outputs_with_lora.shape == (1, 77, 32)
  
  							assert not torch.allclose(
  							    snake_case			,  snake_case						), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"
  
  
  
  
  
  
  	def 			_a(self							:   str						)   -> Tuple:
  							_lowercase			,						_lowercase					:  List[Any]    			=      self.get_dummy_components()
  							_lowercase					:  str    			=      StableDiffusionPipeline(**snake_case						)
  
  							_lowercase					:  Dict    			=      self.get_dummy_tokens()
  
  							# inference without lora
  							_lowercase					:  Tuple    			=      pipe.text_encoder(**snake_case						)[0]
  							assert outputs_without_lora.shape == (1, 77, 32)
  
  							# create lora_attn_procs with randn up.weights
  							_lowercase					:  Optional[int]    			=      create_text_encoder_lora_attn_procs(pipe.text_encoder						)
  							set_lora_up_weights(snake_case			,  randn_weight=snake_case						)
  
  							# monkey patch
  							pipe._modify_text_encoder(snake_case						)
  
  							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
  							del text_attn_procs
  							gc.collect()
  
  							# inference with lora
  							_lowercase					:  Union[str, Any]    			=      pipe.text_encoder(**snake_case						)[0]
  							assert outputs_with_lora.shape == (1, 77, 32)
  
  							assert not torch.allclose(
  							    snake_case			,  snake_case						), "lora outputs should be different to without lora outputs"
  
  							# remove monkey patch
  							pipe._remove_text_encoder_monkey_patch()
  
  							# inference with removed lora
  							_lowercase					:  Union[str, Any]    			=      pipe.text_encoder(**snake_case						)[0]
  							assert outputs_without_lora_removed.shape == (1, 77, 32)
  
  							assert torch.allclose(
  							    snake_case			,  snake_case						), "remove lora monkey patch should restore the original outputs"
  
  
  
  
  
  
  	def 			_a(self							:   Optional[int]						)   -> Dict:
  							_lowercase			,						_lowercase					:  Dict    			=      self.get_dummy_components()
  							_lowercase					:  Optional[Any]    			=      StableDiffusionPipeline(**snake_case						)
  							_lowercase					:  Union[str, Any]    			=      sd_pipe.to(snake_case						)
  							sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  							_lowercase			,						_lowercase			,						_lowercase					:  str    			=      self.get_dummy_inputs()
  
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													LoraLoaderMixin.save_lora_weights(
  													    save_directory=snake_case			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
  													self.assertTrue(os.path.isfile(os.path.join(snake_case			,  "pytorch_lora_weights.bin"						)						)						)
  													sd_pipe.load_lora_weights(snake_case						)
  
  							_lowercase					:  Dict    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  int    			=      lora_images[0, -3:, -3:, -1]
  
  							_lowercase					:  List[str]    			=      sd_pipe(**snake_case			,  cross_attention_kwargs={"scale": 0.5}						).images
  							_lowercase					:  int    			=      lora_images_with_scale[0, -3:, -3:, -1]
  
  							# Outputs shouldn't match.
  							self.assertFalse(
  							    torch.allclose(torch.from_numpy(snake_case						)			,  torch.from_numpy(snake_case						)						)						)
  
  
  
  
  
  
  	def 			_a(self							:   Tuple						)   -> str:
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													self.create_lora_weight_file(snake_case						)
  
  													_lowercase			,						_lowercase					:  str    			=      self.get_dummy_components()
  													_lowercase					:  str    			=      StableDiffusionPipeline(**snake_case						)
  													_lowercase					:  List[Any]    			=      sd_pipe.to(snake_case						)
  													sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  													# check if vanilla attention processors are used
  													for _, module in sd_pipe.unet.named_modules():
  																			if isinstance(snake_case			,  snake_case						):
  																									self.assertIsInstance(module.processor			,  (AttnProcessor, AttnProcessor2_0)						)
  
              # load LoRA weight file
  													sd_pipe.load_lora_weights(snake_case						)
  
  													# check if lora attention processors are used
  													for _, module in sd_pipe.unet.named_modules():
  																			if isinstance(snake_case			,  snake_case						):
  																									_lowercase					:  Tuple    			=      (
  																									    LoRAAttnProcessor2_0 if hasattr(snake_case			,  "scaled_dot_product_attention"						) else LoRAAttnProcessor
  																									)
  																									self.assertIsInstance(module.processor			,  snake_case						)
  
  
  
  
  
  
  	@unittest.skipIf(torch_device != "cuda"			,  "This test is supposed to run on GPU"						)
  	def 			_a(self							:   List[str]						)   -> Union[str, Any]:
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													self.create_lora_weight_file(snake_case						)
  
  													_lowercase			,						_lowercase					:  List[str]    			=      self.get_dummy_components()
  													_lowercase					:  str    			=      StableDiffusionPipeline(**snake_case						)
  													_lowercase					:  Tuple    			=      sd_pipe.to(snake_case						)
  													sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  													# enable XFormers
  													sd_pipe.enable_xformers_memory_efficient_attention()
  
  													# check if xFormers attention processors are used
  													for _, module in sd_pipe.unet.named_modules():
  																			if isinstance(snake_case			,  snake_case						):
  																									self.assertIsInstance(module.processor			,  snake_case						)
  
              # load LoRA weight file
  													sd_pipe.load_lora_weights(snake_case						)
  
  													# check if lora attention processors are used
  													for _, module in sd_pipe.unet.named_modules():
  																			if isinstance(snake_case			,  snake_case						):
  																									self.assertIsInstance(module.processor			,  snake_case						)
  
  
  
  
  
  
  	@unittest.skipIf(torch_device != "cuda"			,  "This test is supposed to run on GPU"						)
  	def 			_a(self							:   List[str]						)   -> Any:
  							_lowercase			,						_lowercase					:  Tuple    			=      self.get_dummy_components()
  							_lowercase					:  List[Any]    			=      StableDiffusionPipeline(**snake_case						)
  							_lowercase					:  Union[str, Any]    			=      sd_pipe.to(snake_case						)
  							sd_pipe.set_progress_bar_config(disable=snake_case						)
  
  							_lowercase			,						_lowercase			,						_lowercase					:  List[Any]    			=      self.get_dummy_inputs()
  
  							# enable XFormers
  							sd_pipe.enable_xformers_memory_efficient_attention()
  
  							_lowercase					:  Tuple    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  Any    			=      original_images[0, -3:, -3:, -1]
  
  							with tempfile.TemporaryDirectory() as tmpdirname:
  													LoraLoaderMixin.save_lora_weights(
  													    save_directory=snake_case			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
  													self.assertTrue(os.path.isfile(os.path.join(snake_case			,  "pytorch_lora_weights.bin"						)						)						)
  													sd_pipe.load_lora_weights(snake_case						)
  
  							_lowercase					:  Optional[int]    			=      sd_pipe(**snake_case						).images
  							_lowercase					:  Tuple    			=      lora_images[0, -3:, -3:, -1]
  
  							# Outputs shouldn't match.
  							self.assertFalse(torch.allclose(torch.from_numpy(snake_case						)			,  torch.from_numpy(snake_case						)						)						)
  
  
  
  
  