



"""simple docstring"""







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      ):
		_lowercase				: Optional[Any]     =   {}
		for name in unet.attn_processors.keys():
				_lowercase				: str     =   None if name.endswith("""attn1.processor"""      ) else unet.config.cross_attention_dim
				if name.startswith("""mid_block"""      ):
						_lowercase				: Dict     =   unet.config.block_out_channels[-1]
				elif name.startswith("""up_blocks"""      ):
						_lowercase				: Any     =   int(name[len("""up_blocks."""      )]      )
						_lowercase				: Any     =   list(reversed(unet.config.block_out_channels      )      )[block_id]
				elif name.startswith("""down_blocks"""      ):
						_lowercase				: int     =   int(name[len("""down_blocks."""      )]      )
						_lowercase				: int     =   unet.config.block_out_channels[block_id]
				_lowercase				: Any     =   (
				    LoRAAttnProcessor2_0 if hasattr(snake_case_   ,				"""scaled_dot_product_attention"""      ) else LoRAAttnProcessor
				)
				_lowercase				: Tuple     =   lora_attn_processor_class(
				    hidden_size=snake_case_   ,				cross_attention_dim=snake_case_      )
		_lowercase				: Any     =   AttnProcsLayers(snake_case_      )
		return lora_attn_procs, unet_lora_layers


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase      ):
		_lowercase				: int     =   {}
		_lowercase				: str     =   (
		    LoRAAttnProcessor2_0 if hasattr(snake_case_   ,				"""scaled_dot_product_attention"""      ) else LoRAAttnProcessor
		)
		for name, module in text_encoder.named_modules():
				if name.endswith(snake_case_      ):
						_lowercase				: Optional[int]     =   lora_attn_processor_class(
						    hidden_size=module.out_proj.out_features   ,				cross_attention_dim=snake_case_      )
		return text_lora_attn_procs


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase      ):
		_lowercase				: List[str]     =   create_text_encoder_lora_attn_procs(snake_case_      )
		_lowercase				: List[Any]     =   AttnProcsLayers(snake_case_      )
		return text_encoder_lora_layers


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase   ,				__UpperCAmelCase=False      ):
		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"""      ):
								_lowercase				: Tuple     =   (
								    torch.randn_like(layer_module.up.weight      )
								    if randn_weight
								    else torch.zeros_like(layer_module.up.weight      )
								)
								_lowercase				: Any     =   torch.nn.Parameter(snake_case_      )





class   UpperCamelCase	(     unittest.TestCase       ):




	"""simple docstring"""






	def    lowerCamelCase__					(     self   ):
			torch.manual_seed(0   )
			_lowercase				: 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				,)
			_lowercase				: Optional[Any]     =   DDIMScheduler(
			    beta_start=0.00085				,beta_end=0.012				,beta_schedule="""scaled_linear"""				,clip_sample=UpperCamelCase_				,set_alpha_to_one=UpperCamelCase_				,steps_offset=1				,)
			torch.manual_seed(0   )
			_lowercase				: 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				,)
			_lowercase				: 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				,)
			_lowercase				: List[Any]     =   CLIPTextModel(UpperCamelCase_   )
			_lowercase				: Union[str, Any]     =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )

			_lowercase				: Optional[int]     =   create_unet_lora_layers(UpperCamelCase_   )
			_lowercase				: Optional[int]     =   create_text_encoder_lora_layers(UpperCamelCase_   )

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

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

			_lowercase				: 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    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[Any]     =   77

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

			_lowercase				: List[str]     =   {}
			_lowercase				: Optional[Any]     =   inputs
			return prepared_inputs






	def    lowerCamelCase__					(     self				,UpperCAmelCase_   ):
			_lowercase				: 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    lowerCamelCase__					(     self   ):
			_lowercase				: Any     =   self.get_dummy_components()
			_lowercase				: Tuple     =   StableDiffusionPipeline(**UpperCamelCase_   )
			_lowercase				: Optional[Any]     =   sd_pipe.to(UpperCamelCase_   )
			sd_pipe.set_progress_bar_config(disable=UpperCamelCase_   )

			_lowercase				: Dict     =   self.get_dummy_inputs()

			_lowercase				: Optional[int]     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: Optional[Any]     =   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_   )

			_lowercase				: Dict     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: Any     =   lora_images[0, -3:, -3:, -1]

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






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Tuple     =   self.get_dummy_components()
			_lowercase				: int     =   StableDiffusionPipeline(**UpperCamelCase_   )
			_lowercase				: Dict     =   sd_pipe.to(UpperCamelCase_   )
			sd_pipe.set_progress_bar_config(disable=UpperCamelCase_   )

			_lowercase				: str     =   self.get_dummy_inputs()

			_lowercase				: Optional[Any]     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: str     =   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_   )

			_lowercase				: int     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: int     =   lora_images[0, -3:, -3:, -1]

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






	def    lowerCamelCase__					(     self   ):
			_lowercase				: str     =   self.get_dummy_components()
			_lowercase				: Optional[Any]     =   lora_components['''unet_lora_attn_procs''']
			_lowercase				: Optional[int]     =   StableDiffusionPipeline(**UpperCamelCase_   )
			_lowercase				: List[str]     =   sd_pipe.to(UpperCamelCase_   )
			sd_pipe.set_progress_bar_config(disable=UpperCamelCase_   )

			_lowercase				: Any     =   self.get_dummy_inputs()

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

			with tempfile.TemporaryDirectory() as tmpdirname:
					_lowercase				: 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_   )

			_lowercase				: List[str]     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: Tuple     =   lora_images[0, -3:, -3:, -1]

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






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Union[str, Any]     =   self.get_dummy_components()
			_lowercase				: Any     =   StableDiffusionPipeline(**UpperCamelCase_   )

			_lowercase				: int     =   self.get_dummy_tokens()

			# inference without lora
			_lowercase				: Dict     =   pipe.text_encoder(**UpperCamelCase_   )[0]
			assert outputs_without_lora.shape == (1, 77, 32)

			# create lora_attn_procs with zeroed out up.weights
			_lowercase				: Any     =   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
			_lowercase				: Dict     =   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
			_lowercase				: 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
			_lowercase				: List[str]     =   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    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[Any]     =   self.get_dummy_components()
			_lowercase				: int     =   StableDiffusionPipeline(**UpperCamelCase_   )

			_lowercase				: Tuple     =   self.get_dummy_tokens()

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

			# create lora_attn_procs with randn up.weights
			_lowercase				: Tuple     =   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
			_lowercase				: Tuple     =   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
			_lowercase				: Dict     =   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    lowerCamelCase__					(     self   ):
			_lowercase				: str     =   self.get_dummy_components()
			_lowercase				: Optional[Any]     =   StableDiffusionPipeline(**UpperCamelCase_   )
			_lowercase				: str     =   sd_pipe.to(UpperCamelCase_   )
			sd_pipe.set_progress_bar_config(disable=UpperCamelCase_   )

			_lowercase				: str     =   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_   )

			_lowercase				: Dict     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: List[Any]     =   lora_images[0, -3:, -3:, -1]

			_lowercase				: List[str]     =   sd_pipe(**UpperCamelCase_				,cross_attention_kwargs={"""scale""": 0.5}   ).images
			_lowercase				: Tuple     =   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    lowerCamelCase__					(     self   ):





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

					_lowercase				: Union[str, Any]     =   self.get_dummy_components()
					_lowercase				: Tuple     =   StableDiffusionPipeline(**UpperCamelCase_   )
					_lowercase				: List[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_   ):
									_lowercase				: Optional[Any]     =   (
									    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    lowerCamelCase__					(     self   ):





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

					_lowercase				: Dict     =   self.get_dummy_components()
					_lowercase				: Dict     =   StableDiffusionPipeline(**UpperCamelCase_   )
					_lowercase				: str     =   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    lowerCamelCase__					(     self   ):
			_lowercase				: Dict     =   self.get_dummy_components()
			_lowercase				: Dict     =   StableDiffusionPipeline(**UpperCamelCase_   )
			_lowercase				: Union[str, Any]     =   sd_pipe.to(UpperCamelCase_   )
			sd_pipe.set_progress_bar_config(disable=UpperCamelCase_   )

			_lowercase				: Dict     =   self.get_dummy_inputs()

			# enable XFormers
			sd_pipe.enable_xformers_memory_efficient_attention()

			_lowercase				: Union[str, Any]     =   sd_pipe(**UpperCamelCase_   ).images
			_lowercase				: Dict     =   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_   )

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

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








'''simple docstring'''






import abc
import copy
import dataclasses
from dataclasses import dataclass
from typing import ClassVar, Dict, Type, TypeVar

from ..features import Features


lowerCamelCase__	  =		TypeVar('T', bound='TaskTemplate')

@dataclass(frozen=__A       )
class 	_lowerCAmelCase     (    abc.ABC       ):






		'''simple docstring'''
		# `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
		snake_case_							= 42
		snake_case_							= 42
		snake_case_							= 42
		def 						__lowercase      (						self							:       T      ,		UpperCamelCase_							:       Features		)      ->			T:







									'''simple docstring'''





									return copy.deepcopy(self		)
		@property
		def 						__lowercase      (						self							:       int		)      ->			Features:







									'''simple docstring'''





									return Features(**self.input_schema      ,		**self.label_schema		)
		@property
		@abc.abstractmethod
		def 						__lowercase      (						self							:       List[str]		)      ->			Dict[str, str]:







									'''simple docstring'''





									raise NotImplementedError





		@classmethod
		def 						__lowercase      (						cls							:       Type[T]      ,		UpperCamelCase_							:       dict		)      ->			T:







									'''simple docstring'''





									_lowercase							:		List[str]				=   {f.name for f in dataclasses.fields(cls		)}
									return cls(**{k: v for k, v in template_dict.items() if k in field_names}		)



