
import copy

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging


UpperCamelCase_	:    List[Any]  					=       logging.get_logger(__name__)




class      __lowercase	(	__snake_case					):
	_A   			=      "hybrid-clip"
	_A   			=      True






	def __init__(self							:   Optional[int]			,  snake_case							:   List[str]=512			,  **snake_case							:   List[Any]						)   -> List[Any]:
							super().__init__(**snake_case						)

							if "text_config" not in kwargs:
													raise ValueError("`text_config` can not be `None`."						)

							if "vision_config" not in kwargs:
													raise ValueError("`vision_config` can not be `None`."						)

							_lowercase					:  int    			=      kwargs.pop("text_config"						)
							_lowercase					:  Optional[int]    			=      kwargs.pop("vision_config"						)

							_lowercase					:  Any    			=      text_config.pop("model_type"						)
							_lowercase					:  Any    			=      vision_config.pop("model_type"						)

							from transformers import AutoConfig

							_lowercase					:  Optional[Any]    			=      AutoConfig.for_model(snake_case			,  **snake_case						)

							if vision_model_type == "clip":
													_lowercase					:  str    			=      AutoConfig.for_model(snake_case			,  **snake_case						).vision_config
							elif vision_model_type == "clip_vision_model":
													from transformers import CLIPVisionConfig

													_lowercase					:  List[Any]    			=      CLIPVisionConfig(**snake_case						)
							else:
													_lowercase					:  List[Any]    			=      AutoConfig.for_model(snake_case			,  **snake_case						)

							_lowercase					:  Union[str, Any]    			=      projection_dim
							_lowercase					:  Optional[Any]    			=      1.0






	@classmethod
	def 			_a(cls							:   Union[str, Any]			,  snake_case							:   PretrainedConfig			,  snake_case							:   PretrainedConfig			,  **snake_case							:   Union[str, Any]						)   -> List[str]:
							return cls(text_config=text_config.to_dict()			,  vision_config=vision_config.to_dict()			,  **snake_case						)






	def 			_a(self							:   Union[str, Any]						)   -> Optional[Any]:
							_lowercase					:  List[str]    			=      copy.deepcopy(self.__dict__						)
							_lowercase					:  List[Any]    			=      self.text_config.to_dict()
							_lowercase					:  str    			=      self.vision_config.to_dict()
							_lowercase					:  int    			=      self.__class__.model_type
							return output






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						)						)						)




