
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 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_	)							->						Any:
			lowerCAmelCase__ : Dict   =						{}
			for name in unet.attn_processors.keys():
						lowerCAmelCase__ : int   =						None if name.endswith('attn1.processor'	) else unet.config.cross_attention_dim
						if name.startswith('mid_block'	):
									lowerCAmelCase__ : Any   =						unet.config.block_out_channels[-1]
						elif name.startswith('up_blocks'	):
									lowerCAmelCase__ : Union[str, Any]   =						int(name[len('up_blocks.'	)]	)
									lowerCAmelCase__ : Any   =						list(reversed(unet.config.block_out_channels	)	)[block_id]
						elif name.startswith('down_blocks'	):
									lowerCAmelCase__ : Dict   =						int(name[len('down_blocks.'	)]	)
									lowerCAmelCase__ : Optional[Any]   =						unet.config.block_out_channels[block_id]
						lowerCAmelCase__ : Union[str, Any]   =						(
						    LoRAAttnProcessor2_0 if hasattr(__lowerCAmelCase			,				'scaled_dot_product_attention'	) else LoRAAttnProcessor
						)
						lowerCAmelCase__ : Optional[Any]   =						lora_attn_processor_class(
						    hidden_size=__lowerCAmelCase			,				cross_attention_dim=__lowerCAmelCase	)
			lowerCAmelCase__ : Dict   =						AttnProcsLayers(__lowerCAmelCase	)
			return lora_attn_procs, unet_lora_layers
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_	)							->						str:
			lowerCAmelCase__ : List[str]   =						{}
			lowerCAmelCase__ : Any   =						(
			    LoRAAttnProcessor2_0 if hasattr(__lowerCAmelCase			,				'scaled_dot_product_attention'	) else LoRAAttnProcessor
			)
			for name, module in text_encoder.named_modules():
						if name.endswith(__lowerCAmelCase	):
									lowerCAmelCase__ : Dict   =						lora_attn_processor_class(
									    hidden_size=module.out_proj.out_features			,				cross_attention_dim=__lowerCAmelCase	)
			return text_lora_attn_procs
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_	)							->						Dict:
			lowerCAmelCase__ : int   =						create_text_encoder_lora_attn_procs(__lowerCAmelCase	)
			lowerCAmelCase__ : Optional[Any]   =						AttnProcsLayers(__lowerCAmelCase	)
			return text_encoder_lora_layers
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_=False	)							->						Optional[int]:
			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'	):
												lowerCAmelCase__ : Optional[int]   =						(
												    torch.randn_like(layer_module.up.weight	)
												    if randn_weight
												    else torch.zeros_like(layer_module.up.weight	)
												)
												lowerCAmelCase__ : Tuple   =						torch.nn.Parameter(__lowerCAmelCase	)







class 	A__				(	unittest.TestCase      ):






				def 				_lowerCamelCase						(		self   : Union[str, Any]  ):






							'''simple docstring'''







							torch.manual_seed(0  )
							lowerCAmelCase__ : Tuple   =						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       ,	)
							lowerCAmelCase__ : Union[str, Any]   =						DDIMScheduler(
							    beta_start=0.0_0_0_8_5       ,	beta_end=0.0_1_2       ,	beta_schedule='scaled_linear'       ,	clip_sample=lowerCamelCase__       ,	set_alpha_to_one=lowerCamelCase__       ,	steps_offset=1       ,	)
							torch.manual_seed(0  )
							lowerCAmelCase__ : Optional[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       ,	)
							lowerCAmelCase__ : 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=1_000       ,	)
							lowerCAmelCase__ : Dict   =						CLIPTextModel(lowerCamelCase__  )
							lowerCAmelCase__ : List[str]   =						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

							lowerCAmelCase__ : str   =						create_unet_lora_layers(lowerCamelCase__  )
							lowerCAmelCase__ : str   =						create_text_encoder_lora_layers(lowerCamelCase__  )

							lowerCAmelCase__ : str   =						{
							    "unet": unet,
							    "scheduler": scheduler,
							    "vae": vae,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "safety_checker": None,
							    "feature_extractor": None,
							}
							lowerCAmelCase__ : 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 				_lowerCamelCase						(		self   : str  ):






							'''simple docstring'''







							lowerCAmelCase__ : List[str]   =						1
							lowerCAmelCase__ : Any   =						10
							lowerCAmelCase__ : Union[str, Any]   =						4
							lowerCAmelCase__ : List[str]   =						(32, 32)

							lowerCAmelCase__ : Optional[Any]   =						torch.manual_seed(0  )
							lowerCAmelCase__ : List[Any]   =						floats_tensor((batch_size, num_channels) + sizes  )
							lowerCAmelCase__ : str   =						torch.randint(1       ,	lowerCamelCase__       ,	size=(batch_size, sequence_length)       ,	generator=lowerCamelCase__  )

							lowerCAmelCase__ : Optional[int]   =						{
							    "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 				_lowerCamelCase						(		self   : List[str]  ):






							'''simple docstring'''







							lowerCAmelCase__ : Any   =						77

							lowerCAmelCase__ : str   =						torch.randint(2       ,	56       ,	size=(1, max_seq_length)       ,	generator=torch.manual_seed(0  )  )

							lowerCAmelCase__ : Dict   =						{}
							lowerCAmelCase__ : Union[str, Any]   =						inputs
							return prepared_inputs






				def 				_lowerCamelCase						(		self   : Any       ,	a   : Tuple  ):






							'''simple docstring'''







							lowerCAmelCase__ : Any   =						self.get_dummy_components()
							LoraLoaderMixin.save_lora_weights(
							    save_directory=lowerCamelCase__       ,	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(lowerCamelCase__       ,	'pytorch_lora_weights.bin'  )  )  )






				def 				_lowerCamelCase						(		self   : Optional[int]  ):






							'''simple docstring'''







							lowerCAmelCase__ : Union[str, Any]   =						self.get_dummy_components()
							lowerCAmelCase__ : Union[str, Any]   =						StableDiffusionPipeline(**lowerCamelCase__  )
							lowerCAmelCase__ : Optional[int]   =						sd_pipe.to(lowerCamelCase__  )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

							lowerCAmelCase__ : Dict   =						self.get_dummy_inputs()

							lowerCAmelCase__ : List[str]   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : List[str]   =						original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
										LoraLoaderMixin.save_lora_weights(
										    save_directory=lowerCamelCase__       ,	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(lowerCamelCase__       ,	'pytorch_lora_weights.bin'  )  )  )
										sd_pipe.load_lora_weights(lowerCamelCase__  )

							lowerCAmelCase__ : Optional[Any]   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : Optional[Any]   =						lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__  )       ,	torch.from_numpy(lowerCamelCase__  )  )  )






				def 				_lowerCamelCase						(		self   : Any  ):






							'''simple docstring'''







							lowerCAmelCase__ : str   =						self.get_dummy_components()
							lowerCAmelCase__ : List[str]   =						StableDiffusionPipeline(**lowerCamelCase__  )
							lowerCAmelCase__ : Any   =						sd_pipe.to(lowerCamelCase__  )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

							lowerCAmelCase__ : List[Any]   =						self.get_dummy_inputs()

							lowerCAmelCase__ : int   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : Union[str, Any]   =						original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
										LoraLoaderMixin.save_lora_weights(
										    save_directory=lowerCamelCase__       ,	unet_lora_layers=lora_components['unet_lora_layers']       ,	text_encoder_lora_layers=lora_components['text_encoder_lora_layers']       ,	safe_serialization=lowerCamelCase__       ,	)
										self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__       ,	'pytorch_lora_weights.safetensors'  )  )  )
										sd_pipe.load_lora_weights(lowerCamelCase__  )

							lowerCAmelCase__ : Optional[Any]   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : Optional[Any]   =						lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__  )       ,	torch.from_numpy(lowerCamelCase__  )  )  )






				def 				_lowerCamelCase						(		self   : Optional[Any]  ):






							'''simple docstring'''







							lowerCAmelCase__ : Tuple   =						self.get_dummy_components()
							lowerCAmelCase__ : Tuple   =						lora_components["unet_lora_attn_procs"]
							lowerCAmelCase__ : Dict   =						StableDiffusionPipeline(**lowerCamelCase__  )
							lowerCAmelCase__ : List[Any]   =						sd_pipe.to(lowerCamelCase__  )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

							lowerCAmelCase__ : Dict   =						self.get_dummy_inputs()

							lowerCAmelCase__ : Any   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : Tuple   =						original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
										lowerCAmelCase__ : Dict   =						sd_pipe.unet
										unet.set_attn_processor(lowerCamelCase__  )
										unet.save_attn_procs(lowerCamelCase__  )
										self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase__       ,	'pytorch_lora_weights.bin'  )  )  )
										sd_pipe.load_lora_weights(lowerCamelCase__  )

							lowerCAmelCase__ : Dict   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : Optional[int]   =						lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__  )       ,	torch.from_numpy(lowerCamelCase__  )  )  )






				def 				_lowerCamelCase						(		self   : Dict  ):






							'''simple docstring'''







							lowerCAmelCase__ : Any   =						self.get_dummy_components()
							lowerCAmelCase__ : int   =						StableDiffusionPipeline(**lowerCamelCase__  )

							lowerCAmelCase__ : List[Any]   =						self.get_dummy_tokens()

							# inference without lora
							lowerCAmelCase__ : Any   =						pipe.text_encoder(**lowerCamelCase__  )[0]
							assert outputs_without_lora.shape == (1, 77, 32)

							# create lora_attn_procs with zeroed out up.weights
							lowerCAmelCase__ : Dict   =						create_text_encoder_lora_attn_procs(pipe.text_encoder  )
							set_lora_up_weights(lowerCamelCase__       ,	randn_weight=lowerCamelCase__  )

							# monkey patch
							pipe._modify_text_encoder(lowerCamelCase__  )

							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
							del text_attn_procs
							gc.collect()

							# inference with lora
							lowerCAmelCase__ : Union[str, Any]   =						pipe.text_encoder(**lowerCamelCase__  )[0]
							assert outputs_with_lora.shape == (1, 77, 32)

							assert torch.allclose(
							    lowerCamelCase__       ,	lowerCamelCase__  ), "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
							lowerCAmelCase__ : Any   =						create_text_encoder_lora_attn_procs(pipe.text_encoder  )
							set_lora_up_weights(lowerCamelCase__       ,	randn_weight=lowerCamelCase__  )

							# monkey patch
							pipe._modify_text_encoder(lowerCamelCase__  )

							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
							del text_attn_procs
							gc.collect()

							# inference with lora
							lowerCAmelCase__ : Dict   =						pipe.text_encoder(**lowerCamelCase__  )[0]
							assert outputs_with_lora.shape == (1, 77, 32)

							assert not torch.allclose(
							    lowerCamelCase__       ,	lowerCamelCase__  ), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"






				def 				_lowerCamelCase						(		self   : str  ):






							'''simple docstring'''







							lowerCAmelCase__ : Tuple   =						self.get_dummy_components()
							lowerCAmelCase__ : Dict   =						StableDiffusionPipeline(**lowerCamelCase__  )

							lowerCAmelCase__ : str   =						self.get_dummy_tokens()

							# inference without lora
							lowerCAmelCase__ : List[str]   =						pipe.text_encoder(**lowerCamelCase__  )[0]
							assert outputs_without_lora.shape == (1, 77, 32)

							# create lora_attn_procs with randn up.weights
							lowerCAmelCase__ : Optional[Any]   =						create_text_encoder_lora_attn_procs(pipe.text_encoder  )
							set_lora_up_weights(lowerCamelCase__       ,	randn_weight=lowerCamelCase__  )

							# monkey patch
							pipe._modify_text_encoder(lowerCamelCase__  )

							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
							del text_attn_procs
							gc.collect()

							# inference with lora
							lowerCAmelCase__ : Dict   =						pipe.text_encoder(**lowerCamelCase__  )[0]
							assert outputs_with_lora.shape == (1, 77, 32)

							assert not torch.allclose(
							    lowerCamelCase__       ,	lowerCamelCase__  ), "lora outputs should be different to without lora outputs"

							# remove monkey patch
							pipe._remove_text_encoder_monkey_patch()

							# inference with removed lora
							lowerCAmelCase__ : int   =						pipe.text_encoder(**lowerCamelCase__  )[0]
							assert outputs_without_lora_removed.shape == (1, 77, 32)

							assert torch.allclose(
							    lowerCamelCase__       ,	lowerCamelCase__  ), "remove lora monkey patch should restore the original outputs"






				def 				_lowerCamelCase						(		self   : Tuple  ):






							'''simple docstring'''







							lowerCAmelCase__ : Any   =						self.get_dummy_components()
							lowerCAmelCase__ : List[str]   =						StableDiffusionPipeline(**lowerCamelCase__  )
							lowerCAmelCase__ : Dict   =						sd_pipe.to(lowerCamelCase__  )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

							lowerCAmelCase__ : Optional[int]   =						self.get_dummy_inputs()

							with tempfile.TemporaryDirectory() as tmpdirname:
										LoraLoaderMixin.save_lora_weights(
										    save_directory=lowerCamelCase__       ,	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(lowerCamelCase__       ,	'pytorch_lora_weights.bin'  )  )  )
										sd_pipe.load_lora_weights(lowerCamelCase__  )

							lowerCAmelCase__ : Any   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : List[Any]   =						lora_images[0, -3:, -3:, -1]

							lowerCAmelCase__ : Optional[int]   =						sd_pipe(**lowerCamelCase__       ,	cross_attention_kwargs={'scale': 0.5}  ).images
							lowerCAmelCase__ : Optional[int]   =						lora_images_with_scale[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(
							    torch.allclose(torch.from_numpy(lowerCamelCase__  )       ,	torch.from_numpy(lowerCamelCase__  )  )  )






				def 				_lowerCamelCase						(		self   : Dict  ):






							'''simple docstring'''







							with tempfile.TemporaryDirectory() as tmpdirname:
										self.create_lora_weight_file(lowerCamelCase__  )

										lowerCAmelCase__ : Optional[int]   =						self.get_dummy_components()
										lowerCAmelCase__ : List[str]   =						StableDiffusionPipeline(**lowerCamelCase__  )
										lowerCAmelCase__ : List[str]   =						sd_pipe.to(lowerCamelCase__  )
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

										# check if vanilla attention processors are used
										for _, module in sd_pipe.unet.named_modules():
													if isinstance(lowerCamelCase__       ,	lowerCamelCase__  ):
																self.assertIsInstance(module.processor       ,	(AttnProcessor, AttnProcessor2_0)  )

            # load LoRA weight file
										sd_pipe.load_lora_weights(lowerCamelCase__  )

										# check if lora attention processors are used
										for _, module in sd_pipe.unet.named_modules():
													if isinstance(lowerCamelCase__       ,	lowerCamelCase__  ):
																lowerCAmelCase__ : Any   =						(
																    LoRAAttnProcessor2_0 if hasattr(lowerCamelCase__       ,	'scaled_dot_product_attention'  ) else LoRAAttnProcessor
																)
																self.assertIsInstance(module.processor       ,	lowerCamelCase__  )






				@unittest.skipIf(torch_device != 'cuda'       ,	'This test is supposed to run on GPU'  )
				def 				_lowerCamelCase						(		self   : Optional[Any]  ):






							'''simple docstring'''







							with tempfile.TemporaryDirectory() as tmpdirname:
										self.create_lora_weight_file(lowerCamelCase__  )

										lowerCAmelCase__ : str   =						self.get_dummy_components()
										lowerCAmelCase__ : Tuple   =						StableDiffusionPipeline(**lowerCamelCase__  )
										lowerCAmelCase__ : Union[str, Any]   =						sd_pipe.to(lowerCamelCase__  )
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

										# 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(lowerCamelCase__       ,	lowerCamelCase__  ):
																self.assertIsInstance(module.processor       ,	lowerCamelCase__  )

            # load LoRA weight file
										sd_pipe.load_lora_weights(lowerCamelCase__  )

										# check if lora attention processors are used
										for _, module in sd_pipe.unet.named_modules():
													if isinstance(lowerCamelCase__       ,	lowerCamelCase__  ):
																self.assertIsInstance(module.processor       ,	lowerCamelCase__  )






				@unittest.skipIf(torch_device != 'cuda'       ,	'This test is supposed to run on GPU'  )
				def 				_lowerCamelCase						(		self   : Optional[Any]  ):






							'''simple docstring'''







							lowerCAmelCase__ : List[str]   =						self.get_dummy_components()
							lowerCAmelCase__ : int   =						StableDiffusionPipeline(**lowerCamelCase__  )
							lowerCAmelCase__ : Optional[Any]   =						sd_pipe.to(lowerCamelCase__  )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__  )

							lowerCAmelCase__ : List[str]   =						self.get_dummy_inputs()

							# enable XFormers
							sd_pipe.enable_xformers_memory_efficient_attention()

							lowerCAmelCase__ : Union[str, Any]   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : str   =						original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
										LoraLoaderMixin.save_lora_weights(
										    save_directory=lowerCamelCase__       ,	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(lowerCamelCase__       ,	'pytorch_lora_weights.bin'  )  )  )
										sd_pipe.load_lora_weights(lowerCamelCase__  )

							lowerCAmelCase__ : Dict   =						sd_pipe(**lowerCamelCase__  ).images
							lowerCAmelCase__ : Optional[int]   =						lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase__  )       ,	torch.from_numpy(lowerCamelCase__  )  )  )



'''simple docstring'''


def 			__lowerCAmelCase			(__lowerCAmelCase ):
		return "".join(chr(ord(__lowerCAmelCase ) + 32 ) if "A" <= char <= "Z" else char for char in word )


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

						testmod()



