




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

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase)    ->  Optional[Any]:
     a											=   {}
     a											=   (
         LoRAAttnProcessor2_0 if hasattr(__UpperCamelCase	,      "scaled_dot_product_attention") else LoRAAttnProcessor
     )
     for name, module in text_encoder.named_modules():
          if name.endswith(__UpperCamelCase):
               a											=   lora_attn_processor_class(
                   hidden_size=module.out_proj.out_features	,      cross_attention_dim=__UpperCamelCase)
     return text_lora_attn_procs

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase)    ->  str:
     a											=   create_text_encoder_lora_attn_procs(__UpperCamelCase)
     a											=   AttnProcsLayers(__UpperCamelCase)
     return text_encoder_lora_layers

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





class     a__ (      unittest.TestCase   ):







     def     lowerCAmelCase_       (			self				)						->					Union[str, Any]:



          '''simple docstring'''



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

          a						,   a											=   create_unet_lora_layers(A				)
          a											=   create_text_encoder_lora_layers(A				)

          a											=   {
              "unet": unet,
              "scheduler": scheduler,
              "vae": vae,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "safety_checker": None,
              "feature_extractor": None,
          }
          a											=   {
              "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]:



          '''simple docstring'''



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

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

          a											=   {
              "prompt": "A painting of a squirrel eating a burger",
              "generator": generator,
              "num_inference_steps": 2,
              "guidance_scale": 6.0,
              "output_type": "numpy",
          }

          return noise, input_ids, pipeline_inputs

          # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb







     def     lowerCAmelCase_       (			self				)						->					Optional[int]:



          '''simple docstring'''



          a											=   77

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

          a											=   {}
          a											=   inputs
          return prepared_inputs







     def     lowerCAmelCase_       (			self						,	A				)						->					Optional[int]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          LoraLoaderMixin.save_lora_weights(
              save_directory=A						,	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(A						,	"pytorch_lora_weights.bin"				)				)				)







     def     lowerCAmelCase_       (			self				)						->					Any:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

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

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

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

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







     def     lowerCAmelCase_       (			self				)						->					Optional[int]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

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

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

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

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







     def     lowerCAmelCase_       (			self				)						->					Union[str, Any]:



          '''simple docstring'''



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

          a						,   a						,   a											=   self.get_dummy_inputs()

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

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

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

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







     def     lowerCAmelCase_       (			self				)						->					Optional[Any]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)

          a											=   self.get_dummy_tokens()

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

          # create lora_attn_procs with zeroed out up.weights
          a											=   create_text_encoder_lora_attn_procs(pipe.text_encoder				)
          set_lora_up_weights(A						,	randn_weight=A				)

          # monkey patch
          pipe._modify_text_encoder(A				)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert torch.allclose(
              A						,	A				), "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
          a											=   create_text_encoder_lora_attn_procs(pipe.text_encoder				)
          set_lora_up_weights(A						,	randn_weight=A				)

          # monkey patch
          pipe._modify_text_encoder(A				)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

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







     def     lowerCAmelCase_       (			self				)						->					Any:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)

          a											=   self.get_dummy_tokens()

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

          # create lora_attn_procs with randn up.weights
          a											=   create_text_encoder_lora_attn_procs(pipe.text_encoder				)
          set_lora_up_weights(A						,	randn_weight=A				)

          # monkey patch
          pipe._modify_text_encoder(A				)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert not torch.allclose(
              A						,	A				), "lora outputs should be different to without lora outputs"

          # remove monkey patch
          pipe._remove_text_encoder_monkey_patch()

          # inference with removed lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_without_lora_removed.shape == (1, 77, 32)

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







     def     lowerCAmelCase_       (			self				)						->					str:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

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

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

          a											=   sd_pipe(**A						,	cross_attention_kwargs={"scale": 0.5}				).images
          a											=   lora_images_with_scale[0, -3:, -3:, -1]

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







     def     lowerCAmelCase_       (			self				)						->					Optional[Any]:



          '''simple docstring'''



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

               a						,   a											=   self.get_dummy_components()
               a											=   StableDiffusionPipeline(**A				)
               a											=   sd_pipe.to(A				)
               sd_pipe.set_progress_bar_config(disable=A				)

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(A				)

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







     @unittest.skipIf(torch_device != "cuda"						,	"This test is supposed to run on GPU"				)
     def     lowerCAmelCase_       (			self				)						->					Union[str, Any]:



          '''simple docstring'''



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

               a						,   a											=   self.get_dummy_components()
               a											=   StableDiffusionPipeline(**A				)
               a											=   sd_pipe.to(A				)
               sd_pipe.set_progress_bar_config(disable=A				)

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(A				)

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







     @unittest.skipIf(torch_device != "cuda"						,	"This test is supposed to run on GPU"				)
     def     lowerCAmelCase_       (			self				)						->					int:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

          # enable XFormers
          sd_pipe.enable_xformers_memory_efficient_attention()

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

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

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

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







from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL

from ...utils import BaseOutput, is_invisible_watermark_available, is_torch_available, is_transformers_available




@dataclass
class     a__ (      UpperCamelCase__   ):
     a	:    Union[List[PIL.Image.Image], np.ndarray]


if is_transformers_available() and is_torch_available() and is_invisible_watermark_available():
   from .pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
   from .pipeline_stable_diffusion_xl_imgaimg import StableDiffusionXLImgaImgPipeline

