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


def       a__      (					a							)		->	Union[str, Any]:
       A_							:       List[str]	      =				{}
       A_							:       Tuple	      =				(
           LoRAAttnProcessor2_0 if hasattr(a  ,   '''scaled_dot_product_attention'''							) else LoRAAttnProcessor
       )
       for name, module in text_encoder.named_modules():
              if name.endswith(a							):
                     A_							:       Dict	      =				lora_attn_processor_class(
                         hidden_size=module.out_proj.out_features  ,   cross_attention_dim=a							)
       return text_lora_attn_procs


def       a__      (					a							)		->	Tuple:
       A_							:       Any	      =				create_text_encoder_lora_attn_procs(a							)
       A_							:       int	      =				AttnProcsLayers(a							)
       return text_encoder_lora_layers


def       a__      (					a  ,   a=False							)		->	List[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_							:       int	      =				(
                                torch.randn_like(layer_module.up.weight							)
                                if randn_weight
                                else torch.zeros_like(layer_module.up.weight							)
                            )
                            A_							:       List[Any]	      =				torch.nn.Parameter(a							)






class     __UpperCAmelCase(    unittest.TestCase				):






 """simple docstring"""




 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        torch.manual_seed(0      )
        A_							:       Optional[Any]	      =				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_							:       Any	      =				DDIMScheduler(
            beta_start=0.0_00_85 ,       beta_end=0.0_12 ,       beta_schedule='''scaled_linear''' ,       clip_sample=__magic_name__ ,       set_alpha_to_one=__magic_name__ ,       steps_offset=1 ,       )
        torch.manual_seed(0      )
        A_							:       List[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 ,       )
        A_							:       str	      =				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_							:       List[str]	      =				CLIPTextModel(__magic_name__      )
        A_							:       int	      =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

        A_      ,							A_							:       Union[str, Any]	      =				create_unet_lora_layers(__magic_name__      )
        A_							:       List[str]	      =				create_text_encoder_lora_layers(__magic_name__      )

        A_							:       Union[str, Any]	      =				{
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        A_							:       Optional[int]	      =				{
            '''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"""

        A_							:       Optional[int]	      =				1
        A_							:       List[Any]	      =				10
        A_							:       int	      =				4
        A_							:       Any	      =				(32, 32)

        A_							:       int	      =				torch.manual_seed(0      )
        A_							:       int	      =				floats_tensor((batch_size, num_channels) + sizes      )
        A_							:       List[Any]	      =				torch.randint(1 ,       __magic_name__ ,       size=(batch_size, sequence_length) ,       generator=__magic_name__      )

        A_							:       List[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  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_							:       Optional[int]	      =				77

        A_							:       str	      =				torch.randint(2 ,       56 ,       size=(1, max_seq_length) ,       generator=torch.manual_seed(0      )      )

        A_							:       Optional[int]	      =				{}
        A_							:       Optional[int]	      =				inputs
        return prepared_inputs
 def  UpperCAmelCase   (    self ,       __magic_name__      ):

        """simple docstring"""

        A_      ,							A_							:       List[Any]	      =				self.get_dummy_components()
        LoraLoaderMixin.save_lora_weights(
            save_directory=__magic_name__ ,       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(__magic_name__ ,       '''pytorch_lora_weights.bin'''      )      )      )
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       int	      =				self.get_dummy_components()
        A_							:       Any	      =				StableDiffusionPipeline(**__magic_name__      )
        A_							:       int	      =				sd_pipe.to(__magic_name__      )
        sd_pipe.set_progress_bar_config(disable=__magic_name__      )

        A_      ,							A_      ,							A_							:       Optional[Any]	      =				self.get_dummy_inputs()

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

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

        A_							:       Optional[Any]	      =				sd_pipe(**__magic_name__      ).images
        A_							:       Dict	      =				lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__magic_name__      ) ,       torch.from_numpy(__magic_name__      )      )      )
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       List[Any]	      =				self.get_dummy_components()
        A_							:       Union[str, Any]	      =				StableDiffusionPipeline(**__magic_name__      )
        A_							:       Optional[Any]	      =				sd_pipe.to(__magic_name__      )
        sd_pipe.set_progress_bar_config(disable=__magic_name__      )

        A_      ,							A_      ,							A_							:       Optional[int]	      =				self.get_dummy_inputs()

        A_							:       Any	      =				sd_pipe(**__magic_name__      ).images
        A_							:       List[Any]	      =				original_images[0, -3:, -3:, -1]

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

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

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__magic_name__      ) ,       torch.from_numpy(__magic_name__      )      )      )
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       Any	      =				self.get_dummy_components()
        A_							:       Any	      =				lora_components['''unet_lora_attn_procs''']
        A_							:       List[Any]	      =				StableDiffusionPipeline(**__magic_name__      )
        A_							:       Tuple	      =				sd_pipe.to(__magic_name__      )
        sd_pipe.set_progress_bar_config(disable=__magic_name__      )

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

        A_							:       int	      =				sd_pipe(**__magic_name__      ).images
        A_							:       Optional[int]	      =				original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
               A_							:       Tuple	      =				sd_pipe.unet
               unet.set_attn_processor(__magic_name__      )
               unet.save_attn_procs(__magic_name__      )
               self.assertTrue(os.path.isfile(os.path.join(__magic_name__ ,       '''pytorch_lora_weights.bin'''      )      )      )
               sd_pipe.load_lora_weights(__magic_name__      )

        A_							:       Optional[int]	      =				sd_pipe(**__magic_name__      ).images
        A_							:       Union[str, Any]	      =				lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__magic_name__      ) ,       torch.from_numpy(__magic_name__      )      )      )
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       Optional[Any]	      =				self.get_dummy_components()
        A_							:       Any	      =				StableDiffusionPipeline(**__magic_name__      )

        A_							:       Dict	      =				self.get_dummy_tokens()

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

        # create lora_attn_procs with zeroed out up.weights
        A_							:       int	      =				create_text_encoder_lora_attn_procs(pipe.text_encoder      )
        set_lora_up_weights(__magic_name__ ,       randn_weight=__magic_name__      )

        # monkey patch
        pipe._modify_text_encoder(__magic_name__      )

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

        assert torch.allclose(
            __magic_name__ ,       __magic_name__      ), "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_							:       Optional[int]	      =				create_text_encoder_lora_attn_procs(pipe.text_encoder      )
        set_lora_up_weights(__magic_name__ ,       randn_weight=__magic_name__      )

        # monkey patch
        pipe._modify_text_encoder(__magic_name__      )

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

        assert not torch.allclose(
            __magic_name__ ,       __magic_name__      ), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       Union[str, Any]	      =				self.get_dummy_components()
        A_							:       List[Any]	      =				StableDiffusionPipeline(**__magic_name__      )

        A_							:       Tuple	      =				self.get_dummy_tokens()

        # inference without lora
        A_							:       int	      =				pipe.text_encoder(**__magic_name__      )[0]
        assert outputs_without_lora.shape == (1, 77, 32)

        # create lora_attn_procs with randn up.weights
        A_							:       Dict	      =				create_text_encoder_lora_attn_procs(pipe.text_encoder      )
        set_lora_up_weights(__magic_name__ ,       randn_weight=__magic_name__      )

        # monkey patch
        pipe._modify_text_encoder(__magic_name__      )

        # 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_							:       List[str]	      =				pipe.text_encoder(**__magic_name__      )[0]
        assert outputs_with_lora.shape == (1, 77, 32)

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

        # remove monkey patch
        pipe._remove_text_encoder_monkey_patch()

        # inference with removed lora
        A_							:       List[Any]	      =				pipe.text_encoder(**__magic_name__      )[0]
        assert outputs_without_lora_removed.shape == (1, 77, 32)

        assert torch.allclose(
            __magic_name__ ,       __magic_name__      ), "remove lora monkey patch should restore the original outputs"
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       Optional[Any]	      =				self.get_dummy_components()
        A_							:       Optional[int]	      =				StableDiffusionPipeline(**__magic_name__      )
        A_							:       Dict	      =				sd_pipe.to(__magic_name__      )
        sd_pipe.set_progress_bar_config(disable=__magic_name__      )

        A_      ,							A_      ,							A_							:       Any	      =				self.get_dummy_inputs()

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

        A_							:       int	      =				sd_pipe(**__magic_name__      ).images
        A_							:       Optional[Any]	      =				lora_images[0, -3:, -3:, -1]

        A_							:       Dict	      =				sd_pipe(**__magic_name__ ,       cross_attention_kwargs={'''scale''': 0.5}      ).images
        A_							:       Dict	      =				lora_images_with_scale[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(
            torch.allclose(torch.from_numpy(__magic_name__      ) ,       torch.from_numpy(__magic_name__      )      )      )
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

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

               A_      ,							A_							:       Tuple	      =				self.get_dummy_components()
               A_							:       Tuple	      =				StableDiffusionPipeline(**__magic_name__      )
               A_							:       str	      =				sd_pipe.to(__magic_name__      )
               sd_pipe.set_progress_bar_config(disable=__magic_name__      )

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(__magic_name__      )

               # check if lora attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                      if isinstance(__magic_name__ ,       __magic_name__      ):
                             A_							:       str	      =				(
                                 LoRAAttnProcessor2_0 if hasattr(__magic_name__ ,       '''scaled_dot_product_attention'''      ) else LoRAAttnProcessor
                             )
                             self.assertIsInstance(module.processor ,       __magic_name__      )
 @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(__magic_name__      )

               A_      ,							A_							:       Tuple	      =				self.get_dummy_components()
               A_							:       Tuple	      =				StableDiffusionPipeline(**__magic_name__      )
               A_							:       Optional[Any]	      =				sd_pipe.to(__magic_name__      )
               sd_pipe.set_progress_bar_config(disable=__magic_name__      )

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(__magic_name__      )

               # check if lora attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                      if isinstance(__magic_name__ ,       __magic_name__      ):
                             self.assertIsInstance(module.processor ,       __magic_name__      )
 @unittest.skipIf(torch_device != '''cuda''' ,       '''This test is supposed to run on GPU'''      )
 def  UpperCAmelCase   (    self      ):

        """simple docstring"""

        A_      ,							A_							:       Any	      =				self.get_dummy_components()
        A_							:       Tuple	      =				StableDiffusionPipeline(**__magic_name__      )
        A_							:       int	      =				sd_pipe.to(__magic_name__      )
        sd_pipe.set_progress_bar_config(disable=__magic_name__      )

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

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

        A_							:       Any	      =				sd_pipe(**__magic_name__      ).images
        A_							:       Tuple	      =				original_images[0, -3:, -3:, -1]

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

        A_							:       List[Any]	      =				sd_pipe(**__magic_name__      ).images
        A_							:       Optional[int]	      =				lora_images[0, -3:, -3:, -1]

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



from transformers import ProcessorMixin






class     __UpperCAmelCase(    A__				):






 """simple docstring"""




 __magic_name__          =   """AutoFeatureExtractor"""
 __magic_name__          =   """AutoTokenizer"""


