





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





def    _lowercase    (  _UpperCAmelCase )  ->     int:
				lowerCamelCase		       ={}
				lowerCamelCase		       =(
				    LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase	,					"""scaled_dot_product_attention""" ) else LoRAAttnProcessor
				)
				for name, module in text_encoder.named_modules():
								if name.endswith(_UpperCAmelCase ):
												lowerCamelCase		       =lora_attn_processor_class(
												    hidden_size=module.out_proj.out_features	,					cross_attention_dim=_UpperCAmelCase )
				return text_lora_attn_procs





def    _lowercase    (  _UpperCAmelCase )  ->     Any:
				lowerCamelCase		       =create_text_encoder_lora_attn_procs(_UpperCAmelCase )
				lowerCamelCase		       =AttnProcsLayers(_UpperCAmelCase )
				return text_encoder_lora_layers





def    _lowercase    (  _UpperCAmelCase	,					_UpperCAmelCase=False )  ->     int:
				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""" ):
																lowerCamelCase		       =(
																    torch.randn_like(layer_module.up.weight )
																    if randn_weight
																    else torch.zeros_like(layer_module.up.weight )
																)
																lowerCamelCase		       =torch.nn.Parameter(_UpperCAmelCase )




class 					__A   (  unittest.TestCase			):


				def _snake_case				(						self							):
								torch.manual_seed(0							)
								lowerCamelCase		       =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      ,       )
								lowerCamelCase		       =DDIMScheduler(
								    beta_start=0.0_0_0_8_5      ,       beta_end=0.0_1_2      ,       beta_schedule="""scaled_linear"""      ,       clip_sample=UpperCAmelCase_      ,       set_alpha_to_one=UpperCAmelCase_      ,       steps_offset=1      ,       )
								torch.manual_seed(0							)
								lowerCamelCase		       =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      ,       )
								lowerCamelCase		       =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      ,       )
								lowerCamelCase		       =CLIPTextModel(UpperCAmelCase_							)
								lowerCamelCase		       =CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

								lowerCamelCase       ,     lowerCamelCase		       =create_unet_lora_layers(UpperCAmelCase_							)
								lowerCamelCase		       =create_text_encoder_lora_layers(UpperCAmelCase_							)

								lowerCamelCase		       ={
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """vae""": vae,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """safety_checker""": None,
								    """feature_extractor""": None,
								}
								lowerCamelCase		       ={
								    """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 _snake_case				(						self							):
								lowerCamelCase		       =1
								lowerCamelCase		       =10
								lowerCamelCase		       =4
								lowerCamelCase		       =(32, 32)

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

								lowerCamelCase		       ={
								    """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 _snake_case				(						self							):
								lowerCamelCase		       =77

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

								lowerCamelCase		       ={}
								lowerCamelCase		       =inputs
								return prepared_inputs


				def _snake_case				(						self      ,       UpperCAmelCase_							):
								lowerCamelCase       ,     lowerCamelCase		       =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 _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
								lowerCamelCase		       =sd_pipe.to(UpperCAmelCase_							)
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

								lowerCamelCase       ,     lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_inputs()

								lowerCamelCase		       =sd_pipe(**UpperCAmelCase_							).images
								lowerCamelCase		       =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_							)

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

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


				def _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
								lowerCamelCase		       =sd_pipe.to(UpperCAmelCase_							)
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

								lowerCamelCase       ,     lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_inputs()

								lowerCamelCase		       =sd_pipe(**UpperCAmelCase_							).images
								lowerCamelCase		       =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_							)

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

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


				def _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =lora_components["""unet_lora_attn_procs"""]
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
								lowerCamelCase		       =sd_pipe.to(UpperCAmelCase_							)
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

								lowerCamelCase       ,     lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_inputs()

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

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

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

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


				def _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)

								lowerCamelCase		       =self.get_dummy_tokens()

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

								# create lora_attn_procs with zeroed out up.weights
								lowerCamelCase		       =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
								lowerCamelCase		       =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
								lowerCamelCase		       =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
								lowerCamelCase		       =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 _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)

								lowerCamelCase		       =self.get_dummy_tokens()

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

								# create lora_attn_procs with randn up.weights
								lowerCamelCase		       =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
								lowerCamelCase		       =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
								lowerCamelCase		       =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 _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
								lowerCamelCase		       =sd_pipe.to(UpperCAmelCase_							)
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

								lowerCamelCase       ,     lowerCamelCase       ,     lowerCamelCase		       =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_							)

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

								lowerCamelCase		       =sd_pipe(**UpperCAmelCase_      ,       cross_attention_kwargs={"""scale""": 0.5}							).images
								lowerCamelCase		       =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 _snake_case				(						self							):
								with tempfile.TemporaryDirectory() as tmpdirname:
												self.create_lora_weight_file(UpperCAmelCase_							)

												lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
												lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
												lowerCamelCase		       =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_							):
																				lowerCamelCase		       =(
																				    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 _snake_case				(						self							):
								with tempfile.TemporaryDirectory() as tmpdirname:
												self.create_lora_weight_file(UpperCAmelCase_							)

												lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
												lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
												lowerCamelCase		       =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 _snake_case				(						self							):
								lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_components()
								lowerCamelCase		       =StableDiffusionPipeline(**UpperCAmelCase_							)
								lowerCamelCase		       =sd_pipe.to(UpperCAmelCase_							)
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

								lowerCamelCase       ,     lowerCamelCase       ,     lowerCamelCase		       =self.get_dummy_inputs()

								# enable XFormers
								sd_pipe.enable_xformers_memory_efficient_attention()

								lowerCamelCase		       =sd_pipe(**UpperCAmelCase_							).images
								lowerCamelCase		       =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_							)

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

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











'''simple docstring'''
def 				snake_case_  (   a__		:	int			):



    """simple docstring"""




    if not isinstance(a__    ,a__			) or number < 0:
        raise ValueError("""Input must be a non-negative integer"""			)

    __lowercase							=						number & ~(number - 1)
    __lowercase							=						0
    while intermediate:
        intermediate >>= 1
        index += 1
    return index - 1


if __name__ == "__main__":
  import doctest

  doctest.testmod(verbose=True)




