


'''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 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





       snake_case__					:	Tuple          =					{}
       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__					:	Optional[Any]          =					unet.config.block_out_channels[-1]
              elif name.startswith("""up_blocks"""     ):
                     snake_case__					:	Optional[int]          =					int(name[len("""up_blocks."""     )]     )
                     snake_case__					:	List[str]          =					list(reversed(unet.config.block_out_channels     )     )[block_id]
              elif name.startswith("""down_blocks"""     ):
                     snake_case__					:	List[str]          =					int(name[len("""down_blocks."""     )]     )
                     snake_case__					:	Dict          =					unet.config.block_out_channels[block_id]
              snake_case__					:	Dict          =					(
                  LoRAAttnProcessor2_0 if hasattr(UpperCAmelCase    ,       """scaled_dot_product_attention"""     ) else LoRAAttnProcessor
              )
              snake_case__					:	Optional[Any]          =					lora_attn_processor_class(
                  hidden_size=UpperCAmelCase    ,       cross_attention_dim=UpperCAmelCase     )
       snake_case__					:	str          =					AttnProcsLayers(UpperCAmelCase     )
       return lora_attn_procs, unet_lora_layers



def 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





       snake_case__					:	Optional[int]          =					{}
       snake_case__					:	Tuple          =					(
           LoRAAttnProcessor2_0 if hasattr(UpperCAmelCase    ,       """scaled_dot_product_attention"""     ) else LoRAAttnProcessor
       )
       for name, module in text_encoder.named_modules():
              if name.endswith(UpperCAmelCase     ):
                     snake_case__					:	int          =					lora_attn_processor_class(
                         hidden_size=module.out_proj.out_features    ,       cross_attention_dim=UpperCAmelCase     )
       return text_lora_attn_procs



def 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





       snake_case__					:	Any          =					create_text_encoder_lora_attn_procs(UpperCAmelCase     )
       snake_case__					:	int          =					AttnProcsLayers(UpperCAmelCase     )
       return text_encoder_lora_layers



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase=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__					:	Tuple          =					(
                                torch.randn_like(layer_module.up.weight     )
                                if randn_weight
                                else torch.zeros_like(layer_module.up.weight     )
                            )
                            snake_case__					:	List[str]          =					torch.nn.Parameter(UpperCAmelCase     )







class 					_A						(						unittest.TestCase							):



   '''simple docstring'''






   def   __lowerCAmelCase						(	self	:					List[str]						)->	int:
          torch.manual_seed(0						)
          snake_case__					:	str          =					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__					:	Tuple          =					DDIMScheduler(
              beta_start=0.00_085					, beta_end=0.012					, beta_schedule="""scaled_linear"""					, clip_sample=lowerCamelCase					, set_alpha_to_one=lowerCamelCase					, steps_offset=1					, )
          torch.manual_seed(0						)
          snake_case__					:	Optional[Any]          =					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__					:	Tuple          =					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__					:	Dict          =					CLIPTextModel(lowerCamelCase						)
          snake_case__					:	Any          =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

          snake_case__       ,   snake_case__					:	Dict          =					create_unet_lora_layers(lowerCamelCase						)
          snake_case__					:	Dict          =					create_text_encoder_lora_layers(lowerCamelCase						)

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

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

          snake_case__					:	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	:					List[str]						)->	Optional[int]:
          snake_case__					:	Optional[int]          =					77

          snake_case__					:	Union[str, Any]          =					torch.randint(2					, 56					, size=(1, max_seq_length)					, generator=torch.manual_seed(0						)						)

          snake_case__					:	Optional[Any]          =					{}
          snake_case__					:	Union[str, Any]          =					inputs
          return prepared_inputs




   def   __lowerCAmelCase						(	self	:					Optional[Any]					, lowerCamelCase	:					str						)->	int:
          snake_case__       ,   snake_case__					:	Tuple          =					self.get_dummy_components()
          LoraLoaderMixin.save_lora_weights(
              save_directory=lowerCamelCase					, 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(lowerCamelCase					, """pytorch_lora_weights.bin"""						)						)						)




   def   __lowerCAmelCase						(	self	:					Tuple						)->	Dict:
          snake_case__       ,   snake_case__					:	Tuple          =					self.get_dummy_components()
          snake_case__					:	Tuple          =					StableDiffusionPipeline(**lowerCamelCase						)
          snake_case__					:	List[Any]          =					sd_pipe.to(lowerCamelCase						)
          sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

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

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

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

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




   def   __lowerCAmelCase						(	self	:					int						)->	int:
          snake_case__       ,   snake_case__					:	Any          =					self.get_dummy_components()
          snake_case__					:	str          =					StableDiffusionPipeline(**lowerCamelCase						)
          snake_case__					:	str          =					sd_pipe.to(lowerCamelCase						)
          sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

          snake_case__       ,   snake_case__       ,   snake_case__					:	Tuple          =					self.get_dummy_inputs()

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

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

          snake_case__					:	Dict          =					sd_pipe(**lowerCamelCase						).images
          snake_case__					:	List[str]          =					lora_images[0, -3:, -3:, -1]

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




   def   __lowerCAmelCase						(	self	:					List[Any]						)->	Optional[Any]:
          snake_case__       ,   snake_case__					:	Union[str, Any]          =					self.get_dummy_components()
          snake_case__					:	List[str]          =					lora_components["""unet_lora_attn_procs"""]
          snake_case__					:	Union[str, Any]          =					StableDiffusionPipeline(**lowerCamelCase						)
          snake_case__					:	Tuple          =					sd_pipe.to(lowerCamelCase						)
          sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

          snake_case__					:	Optional[int]          =					sd_pipe(**lowerCamelCase						).images
          snake_case__					:	List[Any]          =					original_images[0, -3:, -3:, -1]

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

          snake_case__					:	List[Any]          =					sd_pipe(**lowerCamelCase						).images
          snake_case__					:	Union[str, Any]          =					lora_images[0, -3:, -3:, -1]

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




   def   __lowerCAmelCase						(	self	:					Any						)->	str:
          snake_case__       ,   snake_case__					:	str          =					self.get_dummy_components()
          snake_case__					:	Union[str, Any]          =					StableDiffusionPipeline(**lowerCamelCase						)

          snake_case__					:	Dict          =					self.get_dummy_tokens()

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

          # create lora_attn_procs with zeroed out up.weights
          snake_case__					:	Tuple          =					create_text_encoder_lora_attn_procs(pipe.text_encoder						)
          set_lora_up_weights(lowerCamelCase					, randn_weight=lowerCamelCase						)

          # monkey patch
          pipe._modify_text_encoder(lowerCamelCase						)

          # 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__					:	Optional[int]          =					pipe.text_encoder(**lowerCamelCase						)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert torch.allclose(
              lowerCamelCase					, lowerCamelCase						), "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(lowerCamelCase					, randn_weight=lowerCamelCase						)

          # monkey patch
          pipe._modify_text_encoder(lowerCamelCase						)

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

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




   def   __lowerCAmelCase						(	self	:					int						)->	Any:
          snake_case__       ,   snake_case__					:	Dict          =					self.get_dummy_components()
          snake_case__					:	Optional[Any]          =					StableDiffusionPipeline(**lowerCamelCase						)

          snake_case__					:	Optional[Any]          =					self.get_dummy_tokens()

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

          # create lora_attn_procs with randn up.weights
          snake_case__					:	str          =					create_text_encoder_lora_attn_procs(pipe.text_encoder						)
          set_lora_up_weights(lowerCamelCase					, randn_weight=lowerCamelCase						)

          # monkey patch
          pipe._modify_text_encoder(lowerCamelCase						)

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

          assert not torch.allclose(
              lowerCamelCase					, lowerCamelCase						), "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(**lowerCamelCase						)[0]
          assert outputs_without_lora_removed.shape == (1, 77, 32)

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




   def   __lowerCAmelCase						(	self	:					Optional[Any]						)->	Any:
          snake_case__       ,   snake_case__					:	Union[str, Any]          =					self.get_dummy_components()
          snake_case__					:	Dict          =					StableDiffusionPipeline(**lowerCamelCase						)
          snake_case__					:	str          =					sd_pipe.to(lowerCamelCase						)
          sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

          snake_case__       ,   snake_case__       ,   snake_case__					:	int          =					self.get_dummy_inputs()

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

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

          snake_case__					:	List[str]          =					sd_pipe(**lowerCamelCase					, cross_attention_kwargs={"""scale""": 0.5}						).images
          snake_case__					:	Optional[Any]          =					lora_images_with_scale[0, -3:, -3:, -1]

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




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

                 snake_case__       ,   snake_case__					:	Dict          =					self.get_dummy_components()
                 snake_case__					:	Union[str, Any]          =					StableDiffusionPipeline(**lowerCamelCase						)
                 snake_case__					:	List[str]          =					sd_pipe.to(lowerCamelCase						)
                 sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

            # load LoRA weight file
                 sd_pipe.load_lora_weights(lowerCamelCase						)

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




   @unittest.skipIf(torch_device != """cuda"""					, """This test is supposed to run on GPU"""						)
   def   __lowerCAmelCase						(	self	:					int						)->	Tuple:
          with tempfile.TemporaryDirectory() as tmpdirname:
                 self.create_lora_weight_file(lowerCamelCase						)

                 snake_case__       ,   snake_case__					:	Optional[Any]          =					self.get_dummy_components()
                 snake_case__					:	List[Any]          =					StableDiffusionPipeline(**lowerCamelCase						)
                 snake_case__					:	Dict          =					sd_pipe.to(lowerCamelCase						)
                 sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

            # load LoRA weight file
                 sd_pipe.load_lora_weights(lowerCamelCase						)

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




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

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

          # enable XFormers
          sd_pipe.enable_xformers_memory_efficient_attention()

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

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

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

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








'''simple docstring'''
def 				lowerCAmelCase__	(       UpperCAmelCase = 30     ):






       """simple docstring"""





       snake_case__					:	Tuple          =					exponent + 2
       snake_case__					:	List[Any]          =					(1 + 5**0.5) / 2
       snake_case__					:	Tuple          =					(phi**fibonacci_index - (phi - 1) ** fibonacci_index) / 5**0.5

       return int(UpperCAmelCase     )


if __name__ == "__main__":
     print(f"""{solution() = }""")




