




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  snake_case (							A      )							->							List[Any]:




	'''simple docstring'''


	lowerCamelCase_							:		List[Any]							       =							{}
	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_							:		List[Any]							       =							unet.config.block_out_channels[-1]
		elif name.startswith('up_blocks'      ):
			lowerCamelCase_							:		int							       =							int(name[len('up_blocks.'      )]      )
			lowerCamelCase_							:		Any							       =							list(reversed(unet.config.block_out_channels      )      )[block_id]
		elif name.startswith('down_blocks'      ):
			lowerCamelCase_							:		Any							       =							int(name[len('down_blocks.'      )]      )
			lowerCamelCase_							:		int							       =							unet.config.block_out_channels[block_id]
		lowerCamelCase_							:		Dict							       =							(
		    LoRAAttnProcessor2_0 if hasattr(A						,       'scaled_dot_product_attention'      ) else LoRAAttnProcessor
		)
		lowerCamelCase_							:		Optional[Any]							       =							lora_attn_processor_class(
		    hidden_size=A						,       cross_attention_dim=A      )
	lowerCamelCase_							:		str							       =							AttnProcsLayers(A      )
	return lora_attn_procs, unet_lora_layers






def  snake_case (							A      )							->							Dict:




	'''simple docstring'''


	lowerCamelCase_							:		Dict							       =							{}
	lowerCamelCase_							:		int							       =							(
	    LoRAAttnProcessor2_0 if hasattr(A						,       'scaled_dot_product_attention'      ) else LoRAAttnProcessor
	)
	for name, module in text_encoder.named_modules():
		if name.endswith(A      ):
			lowerCamelCase_							:		Optional[Any]							       =							lora_attn_processor_class(
			    hidden_size=module.out_proj.out_features						,       cross_attention_dim=A      )
	return text_lora_attn_procs






def  snake_case (							A      )							->							Optional[Any]:




	'''simple docstring'''


	lowerCamelCase_							:		Tuple							       =							create_text_encoder_lora_attn_procs(A      )
	lowerCamelCase_							:		List[str]							       =							AttnProcsLayers(A      )
	return text_encoder_lora_layers






def  snake_case (							A						,       A=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'      ):
				lowerCamelCase_							:		List[Any]							       =							(
				    torch.randn_like(layer_module.up.weight      )
				    if randn_weight
				    else torch.zeros_like(layer_module.up.weight      )
				)
				lowerCamelCase_							:		str							       =							torch.nn.Parameter(A      )



class     __SCREAMING_SNAKE_CASE							(		unittest.TestCase  ):

					def      lowercase_							(    self ):







						'''simple docstring'''




						torch.manual_seed(0 )
						lowerCamelCase_							:		str							       =							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_							:		str							       =							DDIMScheduler(
						    beta_start=0.0_00_85			,	beta_end=0.0_12			,	beta_schedule='scaled_linear'			,	clip_sample=__lowercase			,	set_alpha_to_one=__lowercase			,	steps_offset=1			,	)
						torch.manual_seed(0 )
						lowerCamelCase_							:		Dict							       =							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_							:		List[str]							       =							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_							:		List[str]							       =							CLIPTextModel(__lowercase )
						lowerCamelCase_							:		Optional[int]							       =							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

						lowerCamelCase_       ,      lowerCamelCase_							:		Optional[int]							       =							create_unet_lora_layers(__lowercase )
						lowerCamelCase_							:		List[Any]							       =							create_text_encoder_lora_layers(__lowercase )

						lowerCamelCase_							:		Optional[Any]							       =							{
						    'unet': unet,
						    'scheduler': scheduler,
						    'vae': vae,
						    'text_encoder': text_encoder,
						    'tokenizer': tokenizer,
						    'safety_checker': None,
						    'feature_extractor': None,
						}
						lowerCamelCase_							:		Optional[int]							       =							{
						    '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      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Any							       =							1
						lowerCamelCase_							:		Optional[int]							       =							10
						lowerCamelCase_							:		int							       =							4
						lowerCamelCase_							:		int							       =							(32, 32)

						lowerCamelCase_							:		Any							       =							torch.manual_seed(0 )
						lowerCamelCase_							:		List[Any]							       =							floats_tensor((batch_size, num_channels) + sizes )
						lowerCamelCase_							:		Tuple							       =							torch.randint(1			,	__lowercase			,	size=(batch_size, sequence_length)			,	generator=__lowercase )

						lowerCamelCase_							:		Union[str, Any]							       =							{
						    '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      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							       =							77

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

						lowerCamelCase_							:		Any							       =							{}
						lowerCamelCase_							:		Any							       =							inputs
						return prepared_inputs

					def      lowercase_							(    self			,	__lowercase ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		Dict							       =							self.get_dummy_components()
						LoraLoaderMixin.save_lora_weights(
						    save_directory=__lowercase			,	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(__lowercase			,	'pytorch_lora_weights.bin' ) ) )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		Tuple							       =							self.get_dummy_components()
						lowerCamelCase_							:		Any							       =							StableDiffusionPipeline(**__lowercase )
						lowerCamelCase_							:		Optional[Any]							       =							sd_pipe.to(__lowercase )
						sd_pipe.set_progress_bar_config(disable=__lowercase )

						lowerCamelCase_       ,      lowerCamelCase_       ,      lowerCamelCase_							:		List[str]							       =							self.get_dummy_inputs()

						lowerCamelCase_							:		Optional[Any]							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Any							       =							original_images[0, -3:, -3:, -1]

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

						lowerCamelCase_							:		Tuple							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Union[str, Any]							       =							lora_images[0, -3:, -3:, -1]

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

					def      lowercase_							(    self ):







						'''simple docstring'''




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

						lowerCamelCase_       ,      lowerCamelCase_       ,      lowerCamelCase_							:		int							       =							self.get_dummy_inputs()

						lowerCamelCase_							:		List[str]							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Optional[int]							       =							original_images[0, -3:, -3:, -1]

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

						lowerCamelCase_							:		str							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Tuple							       =							lora_images[0, -3:, -3:, -1]

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

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		List[Any]							       =							self.get_dummy_components()
						lowerCamelCase_							:		Optional[Any]							       =							lora_components['unet_lora_attn_procs']
						lowerCamelCase_							:		List[str]							       =							StableDiffusionPipeline(**__lowercase )
						lowerCamelCase_							:		List[Any]							       =							sd_pipe.to(__lowercase )
						sd_pipe.set_progress_bar_config(disable=__lowercase )

						lowerCamelCase_       ,      lowerCamelCase_       ,      lowerCamelCase_							:		Tuple							       =							self.get_dummy_inputs()

						lowerCamelCase_							:		Any							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		List[Any]							       =							original_images[0, -3:, -3:, -1]

						with tempfile.TemporaryDirectory() as tmpdirname:
							lowerCamelCase_							:		Optional[int]							       =							sd_pipe.unet
							unet.set_attn_processor(__lowercase )
							unet.save_attn_procs(__lowercase )
							self.assertTrue(os.path.isfile(os.path.join(__lowercase			,	'pytorch_lora_weights.bin' ) ) )
							sd_pipe.load_lora_weights(__lowercase )

						lowerCamelCase_							:		int							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Any							       =							lora_images[0, -3:, -3:, -1]

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

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		List[str]							       =							self.get_dummy_components()
						lowerCamelCase_							:		str							       =							StableDiffusionPipeline(**__lowercase )

						lowerCamelCase_							:		Dict							       =							self.get_dummy_tokens()

						# inference without lora
						lowerCamelCase_							:		int							       =							pipe.text_encoder(**__lowercase )[0]
						assert outputs_without_lora.shape == (1, 77, 32)

						# create lora_attn_procs with zeroed out up.weights
						lowerCamelCase_							:		str							       =							create_text_encoder_lora_attn_procs(pipe.text_encoder )
						set_lora_up_weights(__lowercase			,	randn_weight=__lowercase )

						# monkey patch
						pipe._modify_text_encoder(__lowercase )

						# 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_							:		Any							       =							pipe.text_encoder(**__lowercase )[0]
						assert outputs_with_lora.shape == (1, 77, 32)

						assert torch.allclose(
						    __lowercase			,	__lowercase ), "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_							:		int							       =							create_text_encoder_lora_attn_procs(pipe.text_encoder )
						set_lora_up_weights(__lowercase			,	randn_weight=__lowercase )

						# monkey patch
						pipe._modify_text_encoder(__lowercase )

						# 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_							:		str							       =							pipe.text_encoder(**__lowercase )[0]
						assert outputs_with_lora.shape == (1, 77, 32)

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

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		List[str]							       =							self.get_dummy_components()
						lowerCamelCase_							:		Tuple							       =							StableDiffusionPipeline(**__lowercase )

						lowerCamelCase_							:		Dict							       =							self.get_dummy_tokens()

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

						# create lora_attn_procs with randn up.weights
						lowerCamelCase_							:		Dict							       =							create_text_encoder_lora_attn_procs(pipe.text_encoder )
						set_lora_up_weights(__lowercase			,	randn_weight=__lowercase )

						# monkey patch
						pipe._modify_text_encoder(__lowercase )

						# 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_							:		int							       =							pipe.text_encoder(**__lowercase )[0]
						assert outputs_with_lora.shape == (1, 77, 32)

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

						# remove monkey patch
						pipe._remove_text_encoder_monkey_patch()

						# inference with removed lora
						lowerCamelCase_							:		Any							       =							pipe.text_encoder(**__lowercase )[0]
						assert outputs_without_lora_removed.shape == (1, 77, 32)

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

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		Union[str, Any]							       =							self.get_dummy_components()
						lowerCamelCase_							:		int							       =							StableDiffusionPipeline(**__lowercase )
						lowerCamelCase_							:		List[str]							       =							sd_pipe.to(__lowercase )
						sd_pipe.set_progress_bar_config(disable=__lowercase )

						lowerCamelCase_       ,      lowerCamelCase_       ,      lowerCamelCase_							:		List[str]							       =							self.get_dummy_inputs()

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

						lowerCamelCase_							:		List[str]							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Optional[int]							       =							lora_images[0, -3:, -3:, -1]

						lowerCamelCase_							:		Dict							       =							sd_pipe(**__lowercase			,	cross_attention_kwargs={'scale': 0.5} ).images
						lowerCamelCase_							:		int							       =							lora_images_with_scale[0, -3:, -3:, -1]

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

					def      lowercase_							(    self ):







						'''simple docstring'''




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

							lowerCamelCase_       ,      lowerCamelCase_							:		Optional[int]							       =							self.get_dummy_components()
							lowerCamelCase_							:		str							       =							StableDiffusionPipeline(**__lowercase )
							lowerCamelCase_							:		Dict							       =							sd_pipe.to(__lowercase )
							sd_pipe.set_progress_bar_config(disable=__lowercase )

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

            # load LoRA weight file
							sd_pipe.load_lora_weights(__lowercase )

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

					@unittest.skipIf(torch_device != 'cuda'			,	'This test is supposed to run on GPU' )
					def      lowercase_							(    self ):







						'''simple docstring'''




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

							lowerCamelCase_       ,      lowerCamelCase_							:		str							       =							self.get_dummy_components()
							lowerCamelCase_							:		Optional[int]							       =							StableDiffusionPipeline(**__lowercase )
							lowerCamelCase_							:		List[str]							       =							sd_pipe.to(__lowercase )
							sd_pipe.set_progress_bar_config(disable=__lowercase )

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

            # load LoRA weight file
							sd_pipe.load_lora_weights(__lowercase )

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





					@unittest.skipIf(torch_device != 'cuda'			,	'This test is supposed to run on GPU' )
					def      lowercase_							(    self ):







						'''simple docstring'''




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

						lowerCamelCase_       ,      lowerCamelCase_       ,      lowerCamelCase_							:		int							       =							self.get_dummy_inputs()

						# enable XFormers
						sd_pipe.enable_xformers_memory_efficient_attention()

						lowerCamelCase_							:		Tuple							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		Optional[Any]							       =							original_images[0, -3:, -3:, -1]

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

						lowerCamelCase_							:		int							       =							sd_pipe(**__lowercase ).images
						lowerCamelCase_							:		List[Any]							       =							lora_images[0, -3:, -3:, -1]

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






from typing import Callable






def  snake_case (							A      )							->							List[str]:




	'''simple docstring'''


	def wrapper(A      ):
		lowerCamelCase_							:		Dict							       =							function_with_docstring.__doc__
		return target_function

	return wrapper
