



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





class 		__SCREAMING_SNAKE_CASE      (				unittest.TestCase      ):






				"""simple docstring"""

				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					torch.manual_seed(0      )
					_snake_case					:							Dict					     =   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					:							Any					     =   DDIMScheduler(
					    beta_start=0.0_0085						,	beta_end=0.012						,	beta_schedule="""scaled_linear"""						,	clip_sample=lowerCamelCase_						,	set_alpha_to_one=lowerCamelCase_						,	steps_offset=1						,	)
					torch.manual_seed(0      )
					_snake_case					:							Tuple					     =   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					:							List[Any]					     =   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=10_00						,	)
					_snake_case					:							List[Any]					     =   CLIPTextModel(lowerCamelCase_      )
					_snake_case					:							List[str]					     =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

					_snake_case     , _snake_case					:							Optional[int]					     =   create_unet_lora_layers(lowerCamelCase_      )
					_snake_case					:							Dict					     =   create_text_encoder_lora_layers(lowerCamelCase_      )

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






					"""simple docstring"""

					_snake_case					:							List[str]					     =   1
					_snake_case					:							str					     =   10
					_snake_case					:							Optional[int]					     =   4
					_snake_case					:							Optional[Any]					     =   (32, 32)

					_snake_case					:							Any					     =   torch.manual_seed(0      )
					_snake_case					:							Union[str, Any]					     =   floats_tensor((batch_size, num_channels) + sizes      )
					_snake_case					:							Any					     =   torch.randint(1						,	lowerCamelCase_						,	size=(batch_size, sequence_length)						,	generator=lowerCamelCase_      )

					_snake_case					:							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 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							Dict					     =   77

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

					_snake_case					:							int					     =   {}
					_snake_case					:							Optional[Any]					     =   inputs
					return prepared_inputs







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

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






					"""simple docstring"""

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

					_snake_case     , _snake_case     , _snake_case					:							str					     =   self.get_dummy_inputs()

					_snake_case					:							int					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							List[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"""]						,	)
						self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase_						,	"""pytorch_lora_weights.bin"""      )      )      )
						sd_pipe.load_lora_weights(lowerCamelCase_      )

					_snake_case					:							Optional[int]					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							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      ):






					"""simple docstring"""

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

					_snake_case     , _snake_case     , _snake_case					:							Any					     =   self.get_dummy_inputs()

					_snake_case					:							List[str]					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							List[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_      )

					_snake_case					:							int					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							List[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      ):






					"""simple docstring"""

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

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

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

					with tempfile.TemporaryDirectory() as tmpdirname:
						_snake_case					:							Optional[Any]					     =   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					:							Dict					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							Tuple					     =   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      ):






					"""simple docstring"""

					_snake_case     , _snake_case					:							str					     =   self.get_dummy_components()
					_snake_case					:							int					     =   StableDiffusionPipeline(**lowerCamelCase_      )

					_snake_case					:							str					     =   self.get_dummy_tokens()

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

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






					"""simple docstring"""

					_snake_case     , _snake_case					:							Optional[Any]					     =   self.get_dummy_components()
					_snake_case					:							Optional[Any]					     =   StableDiffusionPipeline(**lowerCamelCase_      )

					_snake_case					:							Optional[Any]					     =   self.get_dummy_tokens()

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

					# create lora_attn_procs with randn up.weights
					_snake_case					:							Union[str, 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
					_snake_case					:							List[str]					     =   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					:							Optional[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      ):






					"""simple docstring"""

					_snake_case     , _snake_case					:							Tuple					     =   self.get_dummy_components()
					_snake_case					:							str					     =   StableDiffusionPipeline(**lowerCamelCase_      )
					_snake_case					:							int					     =   sd_pipe.to(lowerCamelCase_      )
					sd_pipe.set_progress_bar_config(disable=lowerCamelCase_      )

					_snake_case     , _snake_case     , _snake_case					:							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_      )

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

					_snake_case					:							List[str]					     =   sd_pipe(**lowerCamelCase_						,	cross_attention_kwargs={"""scale""": 0.5}      ).images
					_snake_case					:							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      ):






					"""simple docstring"""

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

						_snake_case     , _snake_case					:							List[str]					     =   self.get_dummy_components()
						_snake_case					:							str					     =   StableDiffusionPipeline(**lowerCamelCase_      )
						_snake_case					:							Union[str, Any]					     =   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					:							Dict					     =   (
								    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      ):






					"""simple docstring"""

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

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






					"""simple docstring"""

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

					_snake_case     , _snake_case     , _snake_case					:							str					     =   self.get_dummy_inputs()

					# enable XFormers
					sd_pipe.enable_xformers_memory_efficient_attention()

					_snake_case					:							Any					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							int					     =   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					:							Dict					     =   sd_pipe(**lowerCamelCase_      ).images
					_snake_case					:							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"""







import warnings

from ...utils import logging
from .image_processing_mobilenet_va import MobileNetVaImageProcessor


SCREAMING_SNAKE_CASE_  					= logging.get_logger(__name__)






class      __a      (       lowerCAmelCase_      ):







       def __init__(       self				,*lowerCamelCase_				,**lowerCamelCase_		):



          '''simple docstring'''

          warnings.warn(
              "The class MobileNetV1FeatureExtractor is deprecated and will be removed in version 5 of Transformers."
              " Please use MobileNetV1ImageProcessor instead."				,lowerCamelCase_				,)
          super().__init__(*lowerCamelCase_				,**lowerCamelCase_		)

