



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        A_  (      _lowercase     ):







						'''simple docstring'''


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







						'''simple docstring'''


						snake_case_			:List[Any]						=      {}
						snake_case_			:List[Any]						=      (
						    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_			:str						=      lora_attn_processor_class(
																		    hidden_size=module.out_proj.out_features, cross_attention_dim=_lowercase     )
						return text_lora_attn_procs
def        A_  (      _lowercase     ):







						'''simple docstring'''


						snake_case_			:List[str]						=      create_text_encoder_lora_attn_procs(_lowercase     )
						snake_case_			:List[Any]						=      AttnProcsLayers(_lowercase     )
						return text_encoder_lora_layers
def        A_  (      _lowercase, _lowercase=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_			:Dict						=      (
																								    torch.randn_like(layer_module.up.weight     )
																								    if randn_weight
																								    else torch.zeros_like(layer_module.up.weight     )
																								)
																								snake_case_			:Dict						=      torch.nn.Parameter(_lowercase     )

class     lowerCamelCase  (      unittest.TestCase    ):






					'''simple docstring'''





					def        lowerCAmelCase_    (	self:							Optional[int]						)       ->     List[str]:
											torch.manual_seed(0						)
											snake_case_			: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						,					)
											snake_case_			:Any						=      DDIMScheduler(
											    beta_start=0.0_0_0_8_5						,					beta_end=0.0_1_2						,					beta_schedule="""scaled_linear"""						,					clip_sample=snake_case						,					set_alpha_to_one=snake_case						,					steps_offset=1						,					)
											torch.manual_seed(0						)
											snake_case_			: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						,					)
											snake_case_			:Dict						=      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_			:List[Any]						=      CLIPTextModel(snake_case						)
											snake_case_			:Tuple						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

											snake_case_			:List[Any]						=      create_unet_lora_layers(snake_case						)
											snake_case_			:Tuple						=      create_text_encoder_lora_layers(snake_case						)

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

											snake_case_			:Optional[Any]						=      torch.manual_seed(0						)
											snake_case_			:Optional[Any]						=      floats_tensor((batch_size, num_channels) + sizes						)
											snake_case_			:Dict						=      torch.randint(1						,					snake_case						,					size=(batch_size, sequence_length)						,					generator=snake_case						)

											snake_case_			:Dict						=      {
											    """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:							List[Any]						)       ->     int:
											snake_case_			:Tuple						=      77

											snake_case_			:List[str]						=      torch.randint(2						,					56						,					size=(1, max_seq_length)						,					generator=torch.manual_seed(0						)						)

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






					def        lowerCAmelCase_    (	self:							int						,					snake_case:							List[str]						)       ->     int:
											snake_case_			:Union[str, Any]						=      self.get_dummy_components()
											LoraLoaderMixin.save_lora_weights(
											    save_directory=snake_case						,					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(snake_case						,					"""pytorch_lora_weights.bin"""						)						)						)






					def        lowerCAmelCase_    (	self:							Dict						)       ->     Tuple:
											snake_case_			:Optional[Any]						=      self.get_dummy_components()
											snake_case_			:Union[str, Any]						=      StableDiffusionPipeline(**snake_case						)
											snake_case_			:Tuple						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:Optional[int]						=      self.get_dummy_inputs()

											snake_case_			:Any						=      sd_pipe(**snake_case						).images
											snake_case_			:Optional[int]						=      original_images[0, -3:, -3:, -1]

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

											snake_case_			:Any						=      sd_pipe(**snake_case						).images
											snake_case_			:Union[str, Any]						=      lora_images[0, -3:, -3:, -1]

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






					def        lowerCAmelCase_    (	self:							Optional[Any]						)       ->     int:
											snake_case_			:int						=      self.get_dummy_components()
											snake_case_			:str						=      StableDiffusionPipeline(**snake_case						)
											snake_case_			:Optional[Any]						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:str						=      self.get_dummy_inputs()

											snake_case_			:int						=      sd_pipe(**snake_case						).images
											snake_case_			:Any						=      original_images[0, -3:, -3:, -1]

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

											snake_case_			:Tuple						=      sd_pipe(**snake_case						).images
											snake_case_			:List[Any]						=      lora_images[0, -3:, -3:, -1]

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






					def        lowerCAmelCase_    (	self:							List[Any]						)       ->     List[Any]:
											snake_case_			:Any						=      self.get_dummy_components()
											snake_case_			:List[Any]						=      lora_components["""unet_lora_attn_procs"""]
											snake_case_			:Dict						=      StableDiffusionPipeline(**snake_case						)
											snake_case_			:Tuple						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:Optional[int]						=      self.get_dummy_inputs()

											snake_case_			:Optional[Any]						=      sd_pipe(**snake_case						).images
											snake_case_			:List[str]						=      original_images[0, -3:, -3:, -1]

											with tempfile.TemporaryDirectory() as tmpdirname:
																	snake_case_			:Any						=      sd_pipe.unet
																	unet.set_attn_processor(snake_case						)
																	unet.save_attn_procs(snake_case						)
																	self.assertTrue(os.path.isfile(os.path.join(snake_case						,					"""pytorch_lora_weights.bin"""						)						)						)
																	sd_pipe.load_lora_weights(snake_case						)

											snake_case_			:Optional[int]						=      sd_pipe(**snake_case						).images
											snake_case_			:int						=      lora_images[0, -3:, -3:, -1]

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






					def        lowerCAmelCase_    (	self:							Dict						)       ->     Optional[int]:
											snake_case_			:str						=      self.get_dummy_components()
											snake_case_			:Tuple						=      StableDiffusionPipeline(**snake_case						)

											snake_case_			:Optional[int]						=      self.get_dummy_tokens()

											# inference without lora
											snake_case_			:Any						=      pipe.text_encoder(**snake_case						)[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(snake_case						,					randn_weight=snake_case						)

											# monkey patch
											pipe._modify_text_encoder(snake_case						)

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

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

											# monkey patch
											pipe._modify_text_encoder(snake_case						)

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

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






					def        lowerCAmelCase_    (	self:							List[str]						)       ->     Dict:
											snake_case_			:Dict						=      self.get_dummy_components()
											snake_case_			:str						=      StableDiffusionPipeline(**snake_case						)

											snake_case_			:str						=      self.get_dummy_tokens()

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

											# create lora_attn_procs with randn up.weights
											snake_case_			:Tuple						=      create_text_encoder_lora_attn_procs(pipe.text_encoder						)
											set_lora_up_weights(snake_case						,					randn_weight=snake_case						)

											# monkey patch
											pipe._modify_text_encoder(snake_case						)

											# 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(**snake_case						)[0]
											assert outputs_with_lora.shape == (1, 77, 32)

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

											# remove monkey patch
											pipe._remove_text_encoder_monkey_patch()

											# inference with removed lora
											snake_case_			:List[str]						=      pipe.text_encoder(**snake_case						)[0]
											assert outputs_without_lora_removed.shape == (1, 77, 32)

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






					def        lowerCAmelCase_    (	self:							Optional[int]						)       ->     Optional[int]:
											snake_case_			:Any						=      self.get_dummy_components()
											snake_case_			:Any						=      StableDiffusionPipeline(**snake_case						)
											snake_case_			:List[str]						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:Any						=      self.get_dummy_inputs()

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

											snake_case_			:Tuple						=      sd_pipe(**snake_case						).images
											snake_case_			:int						=      lora_images[0, -3:, -3:, -1]

											snake_case_			:Tuple						=      sd_pipe(**snake_case						,					cross_attention_kwargs={"""scale""": 0.5}						).images
											snake_case_			:List[Any]						=      lora_images_with_scale[0, -3:, -3:, -1]

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






					def        lowerCAmelCase_    (	self:							Optional[int]						)       ->     List[Any]:
											with tempfile.TemporaryDirectory() as tmpdirname:
																	self.create_lora_weight_file(snake_case						)

																	snake_case_			:List[Any]						=      self.get_dummy_components()
																	snake_case_			:int						=      StableDiffusionPipeline(**snake_case						)
																	snake_case_			:List[Any]						=      sd_pipe.to(snake_case						)
																	sd_pipe.set_progress_bar_config(disable=snake_case						)

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

            # load LoRA weight file
																	sd_pipe.load_lora_weights(snake_case						)

																	# check if lora attention processors are used
																	for _, module in sd_pipe.unet.named_modules():
																							if isinstance(snake_case						,					snake_case						):
																													snake_case_			:int						=      (
																													    LoRAAttnProcessor2_0 if hasattr(snake_case						,					"""scaled_dot_product_attention"""						) else LoRAAttnProcessor
																													)
																													self.assertIsInstance(module.processor						,					snake_case						)






					@unittest.skipIf(torch_device != """cuda"""						,					"""This test is supposed to run on GPU"""						)
					def        lowerCAmelCase_    (	self:							Optional[Any]						)       ->     int:
											with tempfile.TemporaryDirectory() as tmpdirname:
																	self.create_lora_weight_file(snake_case						)

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

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

            # load LoRA weight file
																	sd_pipe.load_lora_weights(snake_case						)

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






					@unittest.skipIf(torch_device != """cuda"""						,					"""This test is supposed to run on GPU"""						)
					def        lowerCAmelCase_    (	self:							List[Any]						)       ->     Union[str, Any]:
											snake_case_			:int						=      self.get_dummy_components()
											snake_case_			:List[Any]						=      StableDiffusionPipeline(**snake_case						)
											snake_case_			:Optional[int]						=      sd_pipe.to(snake_case						)
											sd_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_			:Dict						=      self.get_dummy_inputs()

											# enable XFormers
											sd_pipe.enable_xformers_memory_efficient_attention()

											snake_case_			:Union[str, Any]						=      sd_pipe(**snake_case						).images
											snake_case_			:List[str]						=      original_images[0, -3:, -3:, -1]

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

											snake_case_			:Dict						=      sd_pipe(**snake_case						).images
											snake_case_			:int						=      lora_images[0, -3:, -3:, -1]

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





"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


__a     =    {
    "configuration_xlm_prophetnet": ["XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMProphetNetConfig"],
}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a     =    ["XLMProphetNetTokenizer"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a     =    [
						    "XLM_PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "XLMProphetNetDecoder",
						    "XLMProphetNetEncoder",
						    "XLMProphetNetForCausalLM",
						    "XLMProphetNetForConditionalGeneration",
						    "XLMProphetNetModel",
						    "XLMProphetNetPreTrainedModel",
						]


if TYPE_CHECKING:
						from .configuration_xlm_prophetnet import XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMProphetNetConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_xlm_prophetnet import XLMProphetNetTokenizer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_xlm_prophetnet import (
												    XLM_PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
												    XLMProphetNetDecoder,
												    XLMProphetNetEncoder,
												    XLMProphetNetForCausalLM,
												    XLMProphetNetForConditionalGeneration,
												    XLMProphetNetModel,
												    XLMProphetNetPreTrainedModel,
												)

else:
						import sys

						__a     =    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
