
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 					SCREAMING_SNAKE_CASE  (   __UpperCamelCase			:    nn.Module							)			->							Union[str, Any]:

				"""simple docstring"""







				A__	:						str								=				{}
				for name in unet.attn_processors.keys():
								A__	:						List[Any]								=				None if name.endswith('''attn1.processor'''							) else unet.config.cross_attention_dim
								if name.startswith('''mid_block'''							):
												A__	:						Optional[int]								=				unet.config.block_out_channels[-1]
								elif name.startswith('''up_blocks'''							):
												A__	:						Tuple								=				int(name[len('''up_blocks.'''							)]							)
												A__	:						int								=				list(reversed(unet.config.block_out_channels							)							)[block_id]
								elif name.startswith('''down_blocks'''							):
												A__	:						Any								=				int(name[len('''down_blocks.'''							)]							)
												A__	:						str								=				unet.config.block_out_channels[block_id]
								A__	:						List[str]								=				(
								    LoRAAttnProcessor2_0 if hasattr(__UpperCamelCase				,    '''scaled_dot_product_attention'''							) else LoRAAttnProcessor
								)
								A__	:						int								=				lora_attn_processor_class(
								    hidden_size=__UpperCamelCase				,    cross_attention_dim=__UpperCamelCase							)
				A__	:						Optional[Any]								=				AttnProcsLayers(__UpperCamelCase							)
				return lora_attn_procs, unet_lora_layers

def 					SCREAMING_SNAKE_CASE  (   __UpperCamelCase			:    nn.Module							)			->							str:

				"""simple docstring"""







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

def 					SCREAMING_SNAKE_CASE  (   __UpperCamelCase			:    nn.Module							)			->							Tuple:

				"""simple docstring"""







				A__	:						List[Any]								=				create_text_encoder_lora_attn_procs(__UpperCamelCase							)
				A__	:						List[str]								=				AttnProcsLayers(__UpperCamelCase							)
				return text_encoder_lora_layers

def 					SCREAMING_SNAKE_CASE  (   __UpperCamelCase			:    Optional[Any]				,    __UpperCamelCase			:    Optional[int]=False							)			->							Optional[int]:

				"""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'''							):
																A__	:						Union[str, Any]								=				(
																    torch.randn_like(layer_module.up.weight							)
																    if randn_weight
																    else torch.zeros_like(layer_module.up.weight							)
																)
																A__	:						Dict								=				torch.nn.Parameter(__UpperCamelCase							)







class 			UpperCamelCase__		(     unittest.TestCase ):







	'''simple docstring'''





	def 	__snake_case					(      self		):
					torch.manual_seed(0		)
					A__	:						List[Any]								=				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					,							)
					A__	:						Dict								=				DDIMScheduler(
					    beta_start=0.0_0_0_8_5					,							beta_end=0.0_1_2					,							beta_schedule='''scaled_linear'''					,							clip_sample=UpperCamelCase__					,							set_alpha_to_one=UpperCamelCase__					,							steps_offset=1					,							)
					torch.manual_seed(0		)
					A__	:						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					,							)
					A__	:						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=1000					,							)
					A__	:						int								=				CLIPTextModel(UpperCamelCase__		)
					A__	:						Dict								=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

					A__							,				A__	:						Optional[Any]								=				create_unet_lora_layers(UpperCamelCase__		)
					A__	:						Union[str, Any]								=				create_text_encoder_lora_layers(UpperCamelCase__		)

					A__	:						int								=				{
					    '''unet''': unet,
					    '''scheduler''': scheduler,
					    '''vae''': vae,
					    '''text_encoder''': text_encoder,
					    '''tokenizer''': tokenizer,
					    '''safety_checker''': None,
					    '''feature_extractor''': None,
					}
					A__	:						Optional[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 	__snake_case					(      self		):
					A__	:						Union[str, Any]								=				1
					A__	:						int								=				10
					A__	:						int								=				4
					A__	:						List[Any]								=				(32, 32)

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

					A__	:						Optional[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 	__snake_case					(      self		):
					A__	:						Optional[int]								=				77

					A__	:						Optional[Any]								=				torch.randint(2					,							56					,							size=(1, max_seq_length)					,							generator=torch.manual_seed(0		)		)

					A__	:						int								=				{}
					A__	:						List[Any]								=				inputs
					return prepared_inputs





	def 	__snake_case					(      self					,							UpperCamelCase__		):
					A__							,				A__	:						Tuple								=				self.get_dummy_components()
					LoraLoaderMixin.save_lora_weights(
					    save_directory=UpperCamelCase__					,							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(UpperCamelCase__					,							'''pytorch_lora_weights.bin'''		)		)		)





	def 	__snake_case					(      self		):
					A__							,				A__	:						Tuple								=				self.get_dummy_components()
					A__	:						Optional[int]								=				StableDiffusionPipeline(**UpperCamelCase__		)
					A__	:						Union[str, Any]								=				sd_pipe.to(UpperCamelCase__		)
					sd_pipe.set_progress_bar_config(disable=UpperCamelCase__		)

					A__							,				A__							,				A__	:						List[Any]								=				self.get_dummy_inputs()

					A__	:						Tuple								=				sd_pipe(**UpperCamelCase__		).images
					A__	:						Tuple								=				original_images[0, -3:, -3:, -1]

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

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

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





	def 	__snake_case					(      self		):
					A__							,				A__	:						str								=				self.get_dummy_components()
					A__	:						List[str]								=				StableDiffusionPipeline(**UpperCamelCase__		)
					A__	:						Dict								=				sd_pipe.to(UpperCamelCase__		)
					sd_pipe.set_progress_bar_config(disable=UpperCamelCase__		)

					A__							,				A__							,				A__	:						Optional[Any]								=				self.get_dummy_inputs()

					A__	:						str								=				sd_pipe(**UpperCamelCase__		).images
					A__	:						Tuple								=				original_images[0, -3:, -3:, -1]

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

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

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





	def 	__snake_case					(      self		):
					A__							,				A__	:						Union[str, Any]								=				self.get_dummy_components()
					A__	:						Optional[Any]								=				lora_components['''unet_lora_attn_procs''']
					A__	:						Optional[Any]								=				StableDiffusionPipeline(**UpperCamelCase__		)
					A__	:						Any								=				sd_pipe.to(UpperCamelCase__		)
					sd_pipe.set_progress_bar_config(disable=UpperCamelCase__		)

					A__							,				A__							,				A__	:						List[Any]								=				self.get_dummy_inputs()

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

					with tempfile.TemporaryDirectory() as tmpdirname:
									A__	:						Union[str, Any]								=				sd_pipe.unet
									unet.set_attn_processor(UpperCamelCase__		)
									unet.save_attn_procs(UpperCamelCase__		)
									self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__					,							'''pytorch_lora_weights.bin'''		)		)		)
									sd_pipe.load_lora_weights(UpperCamelCase__		)

					A__	:						List[Any]								=				sd_pipe(**UpperCamelCase__		).images
					A__	:						List[str]								=				lora_images[0, -3:, -3:, -1]

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





	def 	__snake_case					(      self		):
					A__							,				A__	:						Union[str, Any]								=				self.get_dummy_components()
					A__	:						Union[str, Any]								=				StableDiffusionPipeline(**UpperCamelCase__		)

					A__	:						Union[str, Any]								=				self.get_dummy_tokens()

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

					# create lora_attn_procs with zeroed out up.weights
					A__	:						Optional[int]								=				create_text_encoder_lora_attn_procs(pipe.text_encoder		)
					set_lora_up_weights(UpperCamelCase__					,							randn_weight=UpperCamelCase__		)

					# monkey patch
					pipe._modify_text_encoder(UpperCamelCase__		)

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

					# inference with lora
					A__	:						List[Any]								=				pipe.text_encoder(**UpperCamelCase__		)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

					assert torch.allclose(
					    UpperCamelCase__					,							UpperCamelCase__		), "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
					A__	:						str								=				create_text_encoder_lora_attn_procs(pipe.text_encoder		)
					set_lora_up_weights(UpperCamelCase__					,							randn_weight=UpperCamelCase__		)

					# monkey patch
					pipe._modify_text_encoder(UpperCamelCase__		)

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

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

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





	def 	__snake_case					(      self		):
					A__							,				A__	:						str								=				self.get_dummy_components()
					A__	:						Any								=				StableDiffusionPipeline(**UpperCamelCase__		)

					A__	:						Union[str, Any]								=				self.get_dummy_tokens()

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

					# create lora_attn_procs with randn up.weights
					A__	:						int								=				create_text_encoder_lora_attn_procs(pipe.text_encoder		)
					set_lora_up_weights(UpperCamelCase__					,							randn_weight=UpperCamelCase__		)

					# monkey patch
					pipe._modify_text_encoder(UpperCamelCase__		)

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

					# inference with lora
					A__	:						List[str]								=				pipe.text_encoder(**UpperCamelCase__		)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

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

					# remove monkey patch
					pipe._remove_text_encoder_monkey_patch()

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

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





	def 	__snake_case					(      self		):
					A__							,				A__	:						int								=				self.get_dummy_components()
					A__	:						Optional[int]								=				StableDiffusionPipeline(**UpperCamelCase__		)
					A__	:						Any								=				sd_pipe.to(UpperCamelCase__		)
					sd_pipe.set_progress_bar_config(disable=UpperCamelCase__		)

					A__							,				A__							,				A__	:						List[Any]								=				self.get_dummy_inputs()

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

					A__	:						Optional[Any]								=				sd_pipe(**UpperCamelCase__		).images
					A__	:						str								=				lora_images[0, -3:, -3:, -1]

					A__	:						Dict								=				sd_pipe(**UpperCamelCase__					,							cross_attention_kwargs={'''scale''': 0.5}		).images
					A__	:						Optional[Any]								=				lora_images_with_scale[0, -3:, -3:, -1]

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





	def 	__snake_case					(      self		):
					with tempfile.TemporaryDirectory() as tmpdirname:
									self.create_lora_weight_file(UpperCamelCase__		)

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

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

            # load LoRA weight file
									sd_pipe.load_lora_weights(UpperCamelCase__		)

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





	@unittest.skipIf(torch_device != '''cuda'''					,							'''This test is supposed to run on GPU'''		)
	def 	__snake_case					(      self		):
					with tempfile.TemporaryDirectory() as tmpdirname:
									self.create_lora_weight_file(UpperCamelCase__		)

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

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

            # load LoRA weight file
									sd_pipe.load_lora_weights(UpperCamelCase__		)

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





	@unittest.skipIf(torch_device != '''cuda'''					,							'''This test is supposed to run on GPU'''		)
	def 	__snake_case					(      self		):
					A__							,				A__	:						Union[str, Any]								=				self.get_dummy_components()
					A__	:						Optional[int]								=				StableDiffusionPipeline(**UpperCamelCase__		)
					A__	:						Dict								=				sd_pipe.to(UpperCamelCase__		)
					sd_pipe.set_progress_bar_config(disable=UpperCamelCase__		)

					A__							,				A__							,				A__	:						Union[str, Any]								=				self.get_dummy_inputs()

					# enable XFormers
					sd_pipe.enable_xformers_memory_efficient_attention()

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

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

					A__	:						List[str]								=				sd_pipe(**UpperCamelCase__		).images
					A__	:						Any								=				lora_images[0, -3:, -3:, -1]

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

_SCREAMING_SNAKE_CASE     :   dict[str, float]							 =      {
    "kilogram": 1,
    "gram": pow(1_0, 3),
    "milligram": pow(1_0, 6),
    "metric-ton": pow(1_0, -3),
    "long-ton": 0.00_0984_2073,
    "short-ton": 0.00_1102_3122,
    "pound": 2.20_4624_4202,
    "stone": 0.15_7473_1728,
    "ounce": 35.2_7399_0723,
    "carrat": 5_0_0_0,
    "atomic-mass-unit": 6.022136652E26,
}

_SCREAMING_SNAKE_CASE     :   dict[str, float]							 =      {
    "kilogram": 1,
    "gram": pow(1_0, -3),
    "milligram": pow(1_0, -6),
    "metric-ton": pow(1_0, 3),
    "long-ton": 1016.0_4608,
    "short-ton": 907.184,
    "pound": 0.45_3592,
    "stone": 6.3_5029,
    "ounce": 0.028_3495,
    "carrat": 0.0002,
    "atomic-mass-unit": 1.660540199E-27,
}

def 					SCREAMING_SNAKE_CASE  (   __UpperCamelCase			:    str				,    __UpperCamelCase			:    str				,    __UpperCamelCase			:    float							)			->							float:

				"""simple docstring"""







				if to_type not in KILOGRAM_CHART or from_type not in WEIGHT_TYPE_CHART:
								A__	:						int								=				(
								    F"Invalid 'from_type' or 'to_type' value: {from_type!r}, {to_type!r}\n"
								    F"Supported values are: {', '.join(__UpperCamelCase							)}"
								)
								raise ValueError(__UpperCamelCase							)
				return value * KILOGRAM_CHART[to_type] * WEIGHT_TYPE_CHART[from_type]


if __name__ == "__main__":
					import doctest

					doctest.testmod()