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__(    lowerCamelCase_	):



			"""simple docstring"""





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






def  UpperCAmelCase__(    lowerCamelCase_	):



			"""simple docstring"""





			_snake_case														=   {}
			_snake_case														=   (
			    LoRAAttnProcessor2_0 if hasattr(lowercase_				,						"scaled_dot_product_attention"	) else LoRAAttnProcessor
			)
			for name, module in text_encoder.named_modules():
						if name.endswith(lowercase_	):
									_snake_case														=   lora_attn_processor_class(
									    hidden_size=module.out_proj.out_features				,						cross_attention_dim=lowercase_	)
			return text_lora_attn_procs






def  UpperCAmelCase__(    lowerCamelCase_	):



			"""simple docstring"""





			_snake_case														=   create_text_encoder_lora_attn_procs(lowercase_	)
			_snake_case														=   AttnProcsLayers(lowercase_	)
			return text_encoder_lora_layers






def  UpperCAmelCase__(    lowerCamelCase_				,						lowerCamelCase_=False	):



			"""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"	):
												_snake_case														=   (
												    torch.randn_like(layer_module.up.weight	)
												    if randn_weight
												    else torch.zeros_like(layer_module.up.weight	)
												)
												_snake_case														=   torch.nn.Parameter(lowercase_	)







class 						_SCREAMING_SNAKE_CASE				(  unittest.TestCase				):






		"""simple docstring"""







		def 					_a					(      self     :					Dict ):






					"""simple docstring"""
					torch.manual_seed(0 )
					_snake_case														=   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	,       )
					_snake_case														=   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 )
					_snake_case														=   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	,       )
					_snake_case														=   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	,       )
					_snake_case														=   CLIPTextModel(lowerCAmelCase__ )
					_snake_case														=   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

					_snake_case							,			_snake_case														=   create_unet_lora_layers(lowerCAmelCase__ )
					_snake_case														=   create_text_encoder_lora_layers(lowerCAmelCase__ )

					_snake_case														=   {
					    "unet": unet,
					    "scheduler": scheduler,
					    "vae": vae,
					    "text_encoder": text_encoder,
					    "tokenizer": tokenizer,
					    "safety_checker": None,
					    "feature_extractor": None,
					}
					_snake_case														=   {
					    "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[str] ):






					"""simple docstring"""
					_snake_case														=   1
					_snake_case														=   10
					_snake_case														=   4
					_snake_case														=   (32, 32)

					_snake_case														=   torch.manual_seed(0 )
					_snake_case														=   floats_tensor((batch_size, num_channels) + sizes )
					_snake_case														=   torch.randint(1	,       lowerCAmelCase__	,       size=(batch_size, sequence_length)	,       generator=lowerCAmelCase__ )

					_snake_case														=   {
					    "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     :					List[str] ):






					"""simple docstring"""
					_snake_case														=   77

					_snake_case														=   torch.randint(2	,       56	,       size=(1, max_seq_length)	,       generator=torch.manual_seed(0 ) )

					_snake_case														=   {}
					_snake_case														=   inputs
					return prepared_inputs



		def 					_a					(      self     :					int	,       __SCREAMING_SNAKE_CASE     :					List[Any] ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   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 					_a					(      self     :					Any ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
					_snake_case														=   sd_pipe.to(lowerCAmelCase__ )
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__ )

					_snake_case							,			_snake_case							,			_snake_case														=   self.get_dummy_inputs()

					_snake_case														=   sd_pipe(**lowerCAmelCase__ ).images
					_snake_case														=   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__ )

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

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



		def 					_a					(      self     :					Tuple ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
					_snake_case														=   sd_pipe.to(lowerCAmelCase__ )
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__ )

					_snake_case							,			_snake_case							,			_snake_case														=   self.get_dummy_inputs()

					_snake_case														=   sd_pipe(**lowerCAmelCase__ ).images
					_snake_case														=   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__ )

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

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



		def 					_a					(      self     :					List[str] ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   lora_components["unet_lora_attn_procs"]
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
					_snake_case														=   sd_pipe.to(lowerCAmelCase__ )
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__ )

					_snake_case							,			_snake_case							,			_snake_case														=   self.get_dummy_inputs()

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

					with tempfile.TemporaryDirectory() as tmpdirname:
								_snake_case														=   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__ )

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

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



		def 					_a					(      self     :					Dict ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )

					_snake_case														=   self.get_dummy_tokens()

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

					# create lora_attn_procs with zeroed out up.weights
					_snake_case														=   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
					_snake_case														=   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
					_snake_case														=   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
					_snake_case														=   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 					_a					(      self     :					Optional[int] ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )

					_snake_case														=   self.get_dummy_tokens()

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

					# create lora_attn_procs with randn up.weights
					_snake_case														=   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
					_snake_case														=   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
					_snake_case														=   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 					_a					(      self     :					Dict ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
					_snake_case														=   sd_pipe.to(lowerCAmelCase__ )
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__ )

					_snake_case							,			_snake_case							,			_snake_case														=   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__ )

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

					_snake_case														=   sd_pipe(**lowerCAmelCase__	,       cross_attention_kwargs={"scale": 0.5} ).images
					_snake_case														=   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 					_a					(      self     :					str ):






					"""simple docstring"""
					with tempfile.TemporaryDirectory() as tmpdirname:
								self.create_lora_weight_file(lowerCAmelCase__ )

								_snake_case							,			_snake_case														=   self.get_dummy_components()
								_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
								_snake_case														=   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__ ):
														_snake_case														=   (
														    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 					_a					(      self     :					str ):






					"""simple docstring"""
					with tempfile.TemporaryDirectory() as tmpdirname:
								self.create_lora_weight_file(lowerCAmelCase__ )

								_snake_case							,			_snake_case														=   self.get_dummy_components()
								_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
								_snake_case														=   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 					_a					(      self     :					int ):






					"""simple docstring"""
					_snake_case							,			_snake_case														=   self.get_dummy_components()
					_snake_case														=   StableDiffusionPipeline(**lowerCAmelCase__ )
					_snake_case														=   sd_pipe.to(lowerCAmelCase__ )
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__ )

					_snake_case							,			_snake_case							,			_snake_case														=   self.get_dummy_inputs()

					# enable XFormers
					sd_pipe.enable_xformers_memory_efficient_attention()

					_snake_case														=   sd_pipe(**lowerCAmelCase__ ).images
					_snake_case														=   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__ )

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

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









# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
__UpperCAmelCase      =				abspath(join(dirname(dirname(dirname(__file__))), "src"))
sys.path.insert(1, git_repo_path)


# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action="ignore", category=FutureWarning)





def 						A_					(						lowercase_			)				->str:







 """simple docstring"""

 from transformers.testing_utils import pytest_addoption_shared

 pytest_addoption_shared(lowercase_			)





def 						A_					(						lowercase_			)				->Tuple:







 """simple docstring"""

 from transformers.testing_utils import pytest_terminal_summary_main

 SCREAMING_SNAKE_CASE				    =		terminalreporter.config.getoption('--make-reports'			)
 if make_reports:
  pytest_terminal_summary_main(lowercase_	,  id=lowercase_			)




