






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   (				__lowerCAmelCase     )->  Tuple:





						"""simple docstring"""



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


def __A   (				__lowerCAmelCase     )->  int:





						"""simple docstring"""



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


def __A   (				__lowerCAmelCase     )->  Optional[Any]:





						"""simple docstring"""



						_UpperCAmelCase     =      create_text_encoder_lora_attn_procs(_lowerCAmelCase     )
						_UpperCAmelCase     =      AttnProcsLayers(_lowerCAmelCase     )
						return text_encoder_lora_layers


def __A   (				__lowerCAmelCase			,		__lowerCAmelCase=False     )->  int:





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





class __lowerCamelCase			(				unittest.TestCase):

	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							torch.manual_seed(0   )
							_UpperCAmelCase     =      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	,   )
							_UpperCAmelCase     =      DDIMScheduler(
							    beta_start=0.0_00_85	,   beta_end=0.0_12	,   beta_schedule='scaled_linear'	,   clip_sample=__snake_case	,   set_alpha_to_one=__snake_case	,   steps_offset=1	,   )
							torch.manual_seed(0   )
							_UpperCAmelCase     =      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	,   )
							_UpperCAmelCase     =      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=1000	,   )
							_UpperCAmelCase     =      CLIPTextModel(__snake_case   )
							_UpperCAmelCase     =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

							_UpperCAmelCase     =      create_unet_lora_layers(__snake_case   )
							_UpperCAmelCase     =      create_text_encoder_lora_layers(__snake_case   )

							_UpperCAmelCase     =      {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							_UpperCAmelCase     =      {
							    '''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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      1
							_UpperCAmelCase     =      10
							_UpperCAmelCase     =      4
							_UpperCAmelCase     =      (32, 32)

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

							_UpperCAmelCase     =      {
							    '''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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      77

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

							_UpperCAmelCase     =      {}
							_UpperCAmelCase     =      inputs
							return prepared_inputs

	def 	UpperCamelCase   (							self	,   UpperCAmelCase   ):
							"""simple docstring"""







							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
							_UpperCAmelCase     =      sd_pipe.to(__snake_case   )
							sd_pipe.set_progress_bar_config(disable=__snake_case   )

							_UpperCAmelCase     =      self.get_dummy_inputs()

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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   )

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
							_UpperCAmelCase     =      sd_pipe.to(__snake_case   )
							sd_pipe.set_progress_bar_config(disable=__snake_case   )

							_UpperCAmelCase     =      self.get_dummy_inputs()

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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   )

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      lora_components['''unet_lora_attn_procs''']
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
							_UpperCAmelCase     =      sd_pipe.to(__snake_case   )
							sd_pipe.set_progress_bar_config(disable=__snake_case   )

							_UpperCAmelCase     =      self.get_dummy_inputs()

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

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

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )

							_UpperCAmelCase     =      self.get_dummy_tokens()

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

							# create lora_attn_procs with zeroed out up.weights
							_UpperCAmelCase     =      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
							_UpperCAmelCase     =      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
							_UpperCAmelCase     =      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
							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )

							_UpperCAmelCase     =      self.get_dummy_tokens()

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

							# create lora_attn_procs with randn up.weights
							_UpperCAmelCase     =      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
							_UpperCAmelCase     =      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
							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
							_UpperCAmelCase     =      sd_pipe.to(__snake_case   )
							sd_pipe.set_progress_bar_config(disable=__snake_case   )

							_UpperCAmelCase     =      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   )

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

							_UpperCAmelCase     =      sd_pipe(**__snake_case	,   cross_attention_kwargs={'scale': 0.5}   ).images
							_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







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

													_UpperCAmelCase     =      self.get_dummy_components()
													_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
													_UpperCAmelCase     =      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   ):
																									_UpperCAmelCase     =      (
																									    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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







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

													_UpperCAmelCase     =      self.get_dummy_components()
													_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
													_UpperCAmelCase     =      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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.get_dummy_components()
							_UpperCAmelCase     =      StableDiffusionPipeline(**__snake_case   )
							_UpperCAmelCase     =      sd_pipe.to(__snake_case   )
							sd_pipe.set_progress_bar_config(disable=__snake_case   )

							_UpperCAmelCase     =      self.get_dummy_inputs()

							# enable XFormers
							sd_pipe.enable_xformers_memory_efficient_attention()

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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   )

							_UpperCAmelCase     =      sd_pipe(**__snake_case   ).images
							_UpperCAmelCase     =      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 unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


_a				     =  '''
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
'''





class __lowerCamelCase			(				unittest.TestCase					,	snake_case__):





	"""simple docstring"""






	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      load_tool('summarization'   )
							self.tool.setup()
							_UpperCAmelCase     =      load_tool('summarization'	,   remote=UpperCAmelCase   )

	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.tool(UpperCAmelCase   )
							self.assertEqual(
							    UpperCAmelCase	,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'	,   )

	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.remote_tool(UpperCAmelCase   )
							self.assertEqual(
							    UpperCAmelCase	,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'	,   )

	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.tool(text=UpperCAmelCase   )
							self.assertEqual(
							    UpperCAmelCase	,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'	,   )

	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      self.remote_tool(text=UpperCAmelCase   )
							self.assertEqual(
							    UpperCAmelCase	,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'	,   )
