
"""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       __UpperCAmelCase		(		a_      :		nn.Module					)  ->							List[str]:
   _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(a_     ,	"scaled_dot_product_attention"					) else LoRAAttnProcessor
      )
      _UpperCAmelCase 						=       lora_attn_processor_class(
          hidden_size=a_     ,	cross_attention_dim=a_					)
   _UpperCAmelCase 						=       AttnProcsLayers(a_					)
   return lora_attn_procs, unet_lora_layers
def       __UpperCAmelCase		(		a_      :		nn.Module					)  ->							List[str]:
   _UpperCAmelCase 						=       {}
   _UpperCAmelCase 						=       (
       LoRAAttnProcessor2_0 if hasattr(a_     ,	"scaled_dot_product_attention"					) else LoRAAttnProcessor
   )
   for name, module in text_encoder.named_modules():
      if name.endswith(a_					):
         _UpperCAmelCase 						=       lora_attn_processor_class(
             hidden_size=module.out_proj.out_features     ,	cross_attention_dim=a_					)
   return text_lora_attn_procs
def       __UpperCAmelCase		(		a_      :		nn.Module					)  ->							List[Any]:
   _UpperCAmelCase 						=       create_text_encoder_lora_attn_procs(a_					)
   _UpperCAmelCase 						=       AttnProcsLayers(a_					)
   return text_encoder_lora_layers
def       __UpperCAmelCase		(		a_      :		str     ,	a_      :		List[Any]=False					)  ->							Optional[Any]:
   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(a_					)






class      __a      (       unittest.TestCase      ):







       def  _snake_case      (       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_0_0_8_5				,beta_end=0.0_1_2				,beta_schedule="scaled_linear"				,clip_sample=lowerCamelCase_				,set_alpha_to_one=lowerCamelCase_				,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(lowerCamelCase_		)
          _UpperCAmelCase 						=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

          _UpperCAmelCase						,				_UpperCAmelCase 						=       create_unet_lora_layers(lowerCamelCase_		)
          _UpperCAmelCase 						=       create_text_encoder_lora_layers(lowerCamelCase_		)

          _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  _snake_case      (       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				,lowerCamelCase_				,size=(batch_size, sequence_length)				,generator=lowerCamelCase_		)

          _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  _snake_case      (       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  _snake_case      (       self				,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase						,				_UpperCAmelCase 						=       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  _snake_case      (       self		):



          '''simple docstring'''

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

          _UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_inputs()

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

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

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







       def  _snake_case      (       self		):



          '''simple docstring'''

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

          _UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_inputs()

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

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

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







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_components()
          _UpperCAmelCase 						=       lora_components["unet_lora_attn_procs"]
          _UpperCAmelCase 						=       StableDiffusionPipeline(**lowerCamelCase_		)
          _UpperCAmelCase 						=       sd_pipe.to(lowerCamelCase_		)
          sd_pipe.set_progress_bar_config(disable=lowerCamelCase_		)

          _UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_inputs()

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

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

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

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







       def  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_components()
          _UpperCAmelCase 						=       StableDiffusionPipeline(**lowerCamelCase_		)

          _UpperCAmelCase 						=       self.get_dummy_tokens()

          # inference without lora
          _UpperCAmelCase 						=       pipe.text_encoder(**lowerCamelCase_		)[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(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
          _UpperCAmelCase 						=       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
          _UpperCAmelCase 						=       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
          _UpperCAmelCase 						=       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  _snake_case      (       self		):



          '''simple docstring'''

          _UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_components()
          _UpperCAmelCase 						=       StableDiffusionPipeline(**lowerCamelCase_		)

          _UpperCAmelCase 						=       self.get_dummy_tokens()

          # inference without lora
          _UpperCAmelCase 						=       pipe.text_encoder(**lowerCamelCase_		)[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(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
          _UpperCAmelCase 						=       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
          _UpperCAmelCase 						=       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  _snake_case      (       self		):



          '''simple docstring'''

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

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

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

          _UpperCAmelCase 						=       sd_pipe(**lowerCamelCase_				,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(lowerCamelCase_		)				,torch.from_numpy(lowerCamelCase_		)		)		)







       def  _snake_case      (       self		):



          '''simple docstring'''

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

             _UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_components()
             _UpperCAmelCase 						=       StableDiffusionPipeline(**lowerCamelCase_		)
             _UpperCAmelCase 						=       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_		):
                   _UpperCAmelCase 						=       (
                       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  _snake_case      (       self		):



          '''simple docstring'''

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

             _UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_components()
             _UpperCAmelCase 						=       StableDiffusionPipeline(**lowerCamelCase_		)
             _UpperCAmelCase 						=       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  _snake_case      (       self		):



          '''simple docstring'''

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

          _UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase 						=       self.get_dummy_inputs()

          # enable XFormers
          sd_pipe.enable_xformers_memory_efficient_attention()

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

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

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



"""simple docstring"""







from typing import List

import PIL
import torch
from PIL import Image

from ...configuration_utils import ConfigMixin
from ...models.modeling_utils import ModelMixin
from ...utils import PIL_INTERPOLATION






class      __a      (       lowerCAmelCase_					,						lowerCAmelCase_      ):







       def __init__(       self		):



          '''simple docstring'''

          super().__init__()

          self.register_buffer("watermark_image"				,torch.zeros((62, 62, 4)		)		)
          _UpperCAmelCase 						=       None







       def  _snake_case      (       self				,lowerCamelCase_				,lowerCamelCase_=None		):



          '''simple docstring'''

          _UpperCAmelCase 						=       images[0].height
          _UpperCAmelCase 						=       images[0].width

          _UpperCAmelCase 						=       sample_size or h

          _UpperCAmelCase 						=       min(h / sample_size				,w / sample_size		)
          _UpperCAmelCase						,				_UpperCAmelCase 						=       (int(h / coef		), int(w / coef		)) if coef < 1 else (h, w)

          _UpperCAmelCase						,				_UpperCAmelCase 						=       1024**2, img_w * img_h
          _UpperCAmelCase 						=       (Sa / Sa) ** 0.5
          _UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase 						=       int(K * 62		), img_w - int(14 * K		), img_h - int(14 * K		)

          if self.watermark_image_as_pil is None:
             _UpperCAmelCase 						=       self.watermark_image.to(torch.uinta		).cpu().numpy()
             _UpperCAmelCase 						=       Image.fromarray(lowerCamelCase_				,mode="RGBA"		)
             _UpperCAmelCase 						=       watermark_image

          _UpperCAmelCase 						=       self.watermark_image_as_pil.resize(
              (wm_size, wm_size)				,PIL_INTERPOLATION["bicubic"]				,reducing_gap=lowerCamelCase_		)

          for pil_img in images:
             pil_img.paste(lowerCamelCase_				,box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y)				,mask=wm_img.split()[-1]		)

          return images

