



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







						'''simple docstring'''


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







						'''simple docstring'''


						snake_case_					:str		=     {}
						snake_case_					:int		=     (
						    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_					:Optional[Any]		=     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_					:Optional[int]		=     create_text_encoder_lora_attn_procs(_lowercase      )
						snake_case_					:int		=     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_					:int		=     (
																								    torch.randn_like(layer_module.up.weight      )
																								    if randn_weight
																								    else torch.zeros_like(layer_module.up.weight      )
																								)
																								snake_case_					:str		=     torch.nn.Parameter(_lowercase      )

class 							lowerCamelCase    (     unittest.TestCase							):






					'''simple docstring'''





					def   lowerCAmelCase_							(			self:				Any		)  ->      Tuple:
											torch.manual_seed(0		)
											snake_case_					: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		,	)
											snake_case_					:Optional[int]		=     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_					:Union[str, 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=1_000		,	)
											snake_case_					:Tuple		=     CLIPTextModel(snake_case		)
											snake_case_					:Optional[int]		=     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

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

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

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

											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:				Optional[int]		)  ->      Optional[int]:
											snake_case_					:Tuple		=     77

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

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






					def   lowerCAmelCase_							(			self:				Optional[Any]		,	snake_case:				Any		)  ->      List[str]:
											snake_case_,			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:				Optional[Any]		)  ->      Optional[Any]:
											snake_case_,			snake_case_					:Optional[int]		=     self.get_dummy_components()
											snake_case_					:Optional[Any]		=     StableDiffusionPipeline(**snake_case		)
											snake_case_					:Tuple		=     sd_pipe.to(snake_case		)
											sd_pipe.set_progress_bar_config(disable=snake_case		)

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

											snake_case_					:Dict		=     sd_pipe(**snake_case		).images
											snake_case_					:Union[str, 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"""]		,	)
																	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_					: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:				List[Any]		)  ->      Tuple:
											snake_case_,			snake_case_					:Optional[int]		=     self.get_dummy_components()
											snake_case_					:Tuple		=     StableDiffusionPipeline(**snake_case		)
											snake_case_					:Dict		=     sd_pipe.to(snake_case		)
											sd_pipe.set_progress_bar_config(disable=snake_case		)

											snake_case_,			snake_case_,			snake_case_					:Dict		=     self.get_dummy_inputs()

											snake_case_					:Optional[Any]		=     sd_pipe(**snake_case		).images
											snake_case_					: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"""]		,	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_					:Any		=     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:				Optional[int]		)  ->      str:
											snake_case_,			snake_case_					:Union[str, Any]		=     self.get_dummy_components()
											snake_case_					:Optional[int]		=     lora_components["""unet_lora_attn_procs"""]
											snake_case_					:Dict		=     StableDiffusionPipeline(**snake_case		)
											snake_case_					:Optional[int]		=     sd_pipe.to(snake_case		)
											sd_pipe.set_progress_bar_config(disable=snake_case		)

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

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

											with tempfile.TemporaryDirectory() as tmpdirname:
																	snake_case_					:List[str]		=     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_					:str		=     sd_pipe(**snake_case		).images
											snake_case_					:Optional[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:				Any		)  ->      Optional[Any]:
											snake_case_,			snake_case_					:int		=     self.get_dummy_components()
											snake_case_					:Union[str, Any]		=     StableDiffusionPipeline(**snake_case		)

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

											# inference without lora
											snake_case_					:Optional[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_					:Union[str, 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_					:Any		=     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_					:Dict		=     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 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:				Optional[int]		)  ->      Optional[Any]:
											snake_case_,			snake_case_					:Optional[Any]		=     self.get_dummy_components()
											snake_case_					:List[Any]		=     StableDiffusionPipeline(**snake_case		)

											snake_case_					:Tuple		=     self.get_dummy_tokens()

											# inference without lora
											snake_case_					:Union[str, 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_					:List[str]		=     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 outputs should be different to without lora outputs"

											# remove monkey patch
											pipe._remove_text_encoder_monkey_patch()

											# inference with removed lora
											snake_case_					: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[Any]		)  ->      List[str]:
											snake_case_,			snake_case_					:List[Any]		=     self.get_dummy_components()
											snake_case_					:str		=     StableDiffusionPipeline(**snake_case		)
											snake_case_					:Tuple		=     sd_pipe.to(snake_case		)
											sd_pipe.set_progress_bar_config(disable=snake_case		)

											snake_case_,			snake_case_,			snake_case_					:List[str]		=     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_					:Optional[Any]		=     sd_pipe(**snake_case		).images
											snake_case_					:List[str]		=     lora_images[0, -3:, -3:, -1]

											snake_case_					:Optional[Any]		=     sd_pipe(**snake_case		,	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(snake_case		)		,	torch.from_numpy(snake_case		)		)		)






					def   lowerCAmelCase_							(			self:				Tuple		)  ->      Any:
											with tempfile.TemporaryDirectory() as tmpdirname:
																	self.create_lora_weight_file(snake_case		)

																	snake_case_,			snake_case_					:int		=     self.get_dummy_components()
																	snake_case_					:Tuple		=     StableDiffusionPipeline(**snake_case		)
																	snake_case_					:List[str]		=     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_					:Dict		=     (
																													    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]		)  ->      List[str]:
											with tempfile.TemporaryDirectory() as tmpdirname:
																	self.create_lora_weight_file(snake_case		)

																	snake_case_,			snake_case_					:Union[str, Any]		=     self.get_dummy_components()
																	snake_case_					:Optional[Any]		=     StableDiffusionPipeline(**snake_case		)
																	snake_case_					:int		=     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:				Union[str, Any]		)  ->      int:
											snake_case_,			snake_case_					:Dict		=     self.get_dummy_components()
											snake_case_					:Optional[int]		=     StableDiffusionPipeline(**snake_case		)
											snake_case_					:List[Any]		=     sd_pipe.to(snake_case		)
											sd_pipe.set_progress_bar_config(disable=snake_case		)

											snake_case_,			snake_case_,			snake_case_					:List[str]		=     self.get_dummy_inputs()

											# enable XFormers
											sd_pipe.enable_xformers_memory_efficient_attention()

											snake_case_					:str		=     sd_pipe(**snake_case		).images
											snake_case_					:Dict		=     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_					:List[str]		=     sd_pipe(**snake_case		).images
											snake_case_					:Dict		=     lora_images[0, -3:, -3:, -1]

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







import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
      import torch

      from transformers import GPTaLMHeadModel

@require_torch
@require_sentencepiece
@require_tokenizers
class 	__snake_case					(      unittest.TestCase ):






    '''simple docstring'''





    @slow
    @unittest.skip('''Model is not available.'''				)
    def 						A		( self    :		str				):

      """simple docstring"""





      __lowerCAmelCase								=					'''nvidia/megatron-gpt2-345m/'''
      if "MYDIR" in os.environ:
        __lowerCAmelCase								=					os.path.join(os.environ['''MYDIR''']      ,    snake_case				)
      __lowerCAmelCase								=					GPTaLMHeadModel.from_pretrained(snake_case				)
      model.to(snake_case				)
      model.half()

      __lowerCAmelCase								=					torch.tensor(
          [[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]      ,    device=snake_case      ,    dtype=torch.long      ,    )

      with torch.no_grad():
        __lowerCAmelCase								=					model(snake_case				).logits

      __lowerCAmelCase								=					torch.Size((1, 9, 50_257)				)
      self.assertEqual(output.shape      ,    snake_case				)

      __lowerCAmelCase								=					torch.tensor(
          [
              4.9414,
              -0.2920,
              -1.2148,
              -4.0273,
              -0.5161,
              -5.2109,
              -1.2412,
              -1.8301,
              -1.7734,
              -4.7148,
              -0.2317,
              -1.0811,
              -2.1777,
              0.4141,
              -3.7969,
              -4.0586,
              -2.5332,
              -3.3809,
              4.3867,
          ]      ,    device=snake_case      ,    dtype=torch.half      ,    )

      for i in range(19				):
        __lowerCAmelCase						,    __lowerCAmelCase								=					8 * i // 17, 2_792 * i  # along the diagonal
        __lowerCAmelCase						,    __lowerCAmelCase								=					output[0, r, c], expected_diag[i]
        __lowerCAmelCase								=					f'''row={r} col={c} computed={computed} expected={expected}'''
        self.assertAlmostEqual(snake_case      ,    snake_case      ,    delta=1e-4      ,    msg=snake_case				)
