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






def    __lowerCamelCase	(		lowerCAmelCase__			):
       lowerCAmelCase__  			=      {}
       lowerCAmelCase__  			=      (
           LoRAAttnProcessor2_0 if hasattr(lowerCAmelCase__			,						'scaled_dot_product_attention'			) else LoRAAttnProcessor
       )
       for name, module in text_encoder.named_modules():
              if name.endswith(lowerCAmelCase__			):
                     lowerCAmelCase__  			=      lora_attn_processor_class(
                         hidden_size=module.out_proj.out_features			,						cross_attention_dim=lowerCAmelCase__			)
       return text_lora_attn_procs






def    __lowerCamelCase	(		lowerCAmelCase__			):
       lowerCAmelCase__  			=      create_text_encoder_lora_attn_procs(lowerCAmelCase__			)
       lowerCAmelCase__  			=      AttnProcsLayers(lowerCAmelCase__			)
       return text_encoder_lora_layers






def    __lowerCamelCase	(		lowerCAmelCase__			,						lowerCAmelCase__=False			):
       for _, attn_proc in text_lora_attn_procs.items():
              # set up.weights
              for layer_name, layer_module in attn_proc.named_modules():
                     if layer_name.endswith('_lora'			):
                            lowerCAmelCase__  			=      (
                                torch.randn_like(layer_module.up.weight			)
                                if randn_weight
                                else torch.zeros_like(layer_module.up.weight			)
                            )
                            lowerCAmelCase__  			=      torch.nn.Parameter(lowerCAmelCase__			)



class        a_				(		unittest.TestCase		):






 '''simple docstring'''






 def 	__snake_case		(    self      :						int):







        '''simple docstring'''



        torch.manual_seed(0)
        lowerCAmelCase__  			=      UNetaDConditionModel(
            block_out_channels=(32, 64)    ,					layers_per_block=2    ,					sample_size=32    ,					in_channels=4    ,					out_channels=4    ,					down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')    ,					up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')    ,					cross_attention_dim=32    ,					)
        lowerCAmelCase__  			=      DDIMScheduler(
            beta_start=0.00_085    ,					beta_end=0.012    ,					beta_schedule='scaled_linear'    ,					clip_sample=lowercase__    ,					set_alpha_to_one=lowercase__    ,					steps_offset=1    ,					)
        torch.manual_seed(0)
        lowerCAmelCase__  			=      AutoencoderKL(
            block_out_channels=[32, 64]    ,					in_channels=3    ,					out_channels=3    ,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,					latent_channels=4    ,					)
        lowerCAmelCase__  			=      CLIPTextConfig(
            bos_token_id=0    ,					eos_token_id=2    ,					hidden_size=32    ,					intermediate_size=37    ,					layer_norm_eps=1e-05    ,					num_attention_heads=4    ,					num_hidden_layers=5    ,					pad_token_id=1    ,					vocab_size=1_000    ,					)
        lowerCAmelCase__  			=      CLIPTextModel(lowercase__)
        lowerCAmelCase__  			=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

        lowerCAmelCase__			,      lowerCAmelCase__  			=      create_unet_lora_layers(lowercase__)
        lowerCAmelCase__  			=      create_text_encoder_lora_layers(lowercase__)

        lowerCAmelCase__  			=      {
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }
        lowerCAmelCase__  			=      {
            'unet_lora_layers': unet_lora_layers,
            'text_encoder_lora_layers': text_encoder_lora_layers,
            'unet_lora_attn_procs': unet_lora_attn_procs,
        }
        return pipeline_components, lora_components




 def 	__snake_case		(    self      :						Tuple):







        '''simple docstring'''



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

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

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

        return noise, input_ids, pipeline_inputs

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




 def 	__snake_case		(    self      :						List[Any]):







        '''simple docstring'''



        lowerCAmelCase__  			=      77

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

        lowerCAmelCase__  			=      {}
        lowerCAmelCase__  			=      inputs
        return prepared_inputs




 def 	__snake_case		(    self      :						Tuple    ,					lowercase__      :						Optional[Any]):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        LoraLoaderMixin.save_lora_weights(
            save_directory=lowercase__    ,					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(lowercase__    ,					'pytorch_lora_weights.bin')))




 def 	__snake_case		(    self      :						int):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
        lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__			,      lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_inputs()

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

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

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

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




 def 	__snake_case		(    self      :						str):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
        lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__			,      lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_inputs()

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

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

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

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




 def 	__snake_case		(    self      :						Tuple):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      lora_components['unet_lora_attn_procs']
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
        lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__			,      lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_inputs()

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

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

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

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




 def 	__snake_case		(    self      :						Any):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)

        lowerCAmelCase__  			=      self.get_dummy_tokens()

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

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

        # monkey patch
        pipe._modify_text_encoder(lowercase__)

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

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

        assert torch.allclose(
            lowercase__    ,					lowercase__), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

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

        # monkey patch
        pipe._modify_text_encoder(lowercase__)

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

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

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




 def 	__snake_case		(    self      :						Any):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)

        lowerCAmelCase__  			=      self.get_dummy_tokens()

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

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

        # monkey patch
        pipe._modify_text_encoder(lowercase__)

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

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

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

        # remove monkey patch
        pipe._remove_text_encoder_monkey_patch()

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

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




 def 	__snake_case		(    self      :						str):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
        lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__			,      lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_inputs()

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

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

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

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




 def 	__snake_case		(    self      :						Any):







        '''simple docstring'''



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

               lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
               lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
               lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
               sd_pipe.set_progress_bar_config(disable=lowercase__)

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(lowercase__)

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




 @unittest.skipIf(torch_device != 'cuda'    ,					'This test is supposed to run on GPU')
 def 	__snake_case		(    self      :						int):







        '''simple docstring'''



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

               lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
               lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
               lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
               sd_pipe.set_progress_bar_config(disable=lowercase__)

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(lowercase__)

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




 @unittest.skipIf(torch_device != 'cuda'    ,					'This test is supposed to run on GPU')
 def 	__snake_case		(    self      :						int):







        '''simple docstring'''



        lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_components()
        lowerCAmelCase__  			=      StableDiffusionPipeline(**lowercase__)
        lowerCAmelCase__  			=      sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__			,      lowerCAmelCase__			,      lowerCAmelCase__  			=      self.get_dummy_inputs()

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

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

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

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

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

def    __lowerCamelCase	(		lowerCAmelCase__			):
       return " ".join([word[::-1] for word in input_str.split()]			)


if __name__ == "__main__":
    import doctest

    doctest.testmod()
