






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




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




def       _a   (							lowerCamelCase__				)  ->      List[str]:
 lowerCamelCase_							:    Tuple							=     create_text_encoder_lora_attn_procs(lowerCamelCase__				)
 lowerCamelCase_							:    Optional[Any]							=     AttnProcsLayers(lowerCamelCase__				)
 return text_encoder_lora_layers




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





class 	lowerCamelCase__				(							unittest.TestCase ):






      def      UpperCAmelCase_     (self					:    str		)       ->   Dict:


       """simple docstring"""






       torch.manual_seed(0		)
       lowerCamelCase_							:    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		,		)
       lowerCamelCase_							:    List[str]							=     DDIMScheduler(
           beta_start=0.0_0085		,		beta_end=0.012		,		beta_schedule='scaled_linear'		,		clip_sample=_snake_case		,		set_alpha_to_one=_snake_case		,		steps_offset=1		,		)
       torch.manual_seed(0		)
       lowerCamelCase_							:    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		,		)
       lowerCamelCase_							:    List[Any]							=     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		,		)
       lowerCamelCase_							:    List[Any]							=     CLIPTextModel(_snake_case		)
       lowerCamelCase_							:    Union[str, Any]							=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

       lowerCamelCase_		,      lowerCamelCase_							:    Union[str, Any]							=     create_unet_lora_layers(_snake_case		)
       lowerCamelCase_							:    Any							=     create_text_encoder_lora_layers(_snake_case		)

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


       """simple docstring"""






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

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

       lowerCamelCase_							:    Union[str, 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					:    Tuple		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    str							=     77

       lowerCamelCase_							:    Optional[Any]							=     torch.randint(2		,		56		,		size=(1, max_seq_length)		,		generator=torch.manual_seed(0		)		)

       lowerCamelCase_							:    Optional[Any]							=     {}
       lowerCamelCase_							:    int							=     inputs
       return prepared_inputs






      def      UpperCAmelCase_     (self					:    List[Any]		,		_snake_case					:    Union[str, Any]		)       ->   Tuple:


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    Dict							=     self.get_dummy_components()
       LoraLoaderMixin.save_lora_weights(
           save_directory=_snake_case		,		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(_snake_case		,		'pytorch_lora_weights.bin'		)		)		)






      def      UpperCAmelCase_     (self					:    List[Any]		)       ->   str:


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    Dict							=     self.get_dummy_components()
       lowerCamelCase_							:    List[str]							=     StableDiffusionPipeline(**_snake_case		)
       lowerCamelCase_							:    List[str]							=     sd_pipe.to(_snake_case		)
       sd_pipe.set_progress_bar_config(disable=_snake_case		)

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

       lowerCamelCase_							:    Dict							=     sd_pipe(**_snake_case		).images
       lowerCamelCase_							:    Union[str, Any]							=     original_images[0, -3:, -3:, -1]

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

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

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






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


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    str							=     self.get_dummy_components()
       lowerCamelCase_							:    List[Any]							=     StableDiffusionPipeline(**_snake_case		)
       lowerCamelCase_							:    str							=     sd_pipe.to(_snake_case		)
       sd_pipe.set_progress_bar_config(disable=_snake_case		)

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

       lowerCamelCase_							:    List[Any]							=     sd_pipe(**_snake_case		).images
       lowerCamelCase_							:    Union[str, Any]							=     original_images[0, -3:, -3:, -1]

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

       lowerCamelCase_							:    Union[str, Any]							=     sd_pipe(**_snake_case		).images
       lowerCamelCase_							:    Tuple							=     lora_images[0, -3:, -3:, -1]

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






      def      UpperCAmelCase_     (self					:    Optional[Any]		)       ->   int:


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    Optional[Any]							=     self.get_dummy_components()
       lowerCamelCase_							:    Union[str, Any]							=     lora_components['unet_lora_attn_procs']
       lowerCamelCase_							:    Optional[Any]							=     StableDiffusionPipeline(**_snake_case		)
       lowerCamelCase_							:    int							=     sd_pipe.to(_snake_case		)
       sd_pipe.set_progress_bar_config(disable=_snake_case		)

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

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

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

       lowerCamelCase_							:    Dict							=     sd_pipe(**_snake_case		).images
       lowerCamelCase_							:    Union[str, Any]							=     lora_images[0, -3:, -3:, -1]

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






      def      UpperCAmelCase_     (self					:    Optional[int]		)       ->   Union[str, Any]:


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    Optional[int]							=     self.get_dummy_components()
       lowerCamelCase_							:    int							=     StableDiffusionPipeline(**_snake_case		)

       lowerCamelCase_							:    Dict							=     self.get_dummy_tokens()

       # inference without lora
       lowerCamelCase_							:    str							=     pipe.text_encoder(**_snake_case		)[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with zeroed out up.weights
       lowerCamelCase_							:    str							=     create_text_encoder_lora_attn_procs(pipe.text_encoder		)
       set_lora_up_weights(_snake_case		,		randn_weight=_snake_case		)

       # monkey patch
       pipe._modify_text_encoder(_snake_case		)

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

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

       # monkey patch
       pipe._modify_text_encoder(_snake_case		)

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

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






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


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    List[Any]							=     self.get_dummy_components()
       lowerCamelCase_							:    Dict							=     StableDiffusionPipeline(**_snake_case		)

       lowerCamelCase_							:    Dict							=     self.get_dummy_tokens()

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

       # create lora_attn_procs with randn up.weights
       lowerCamelCase_							:    List[Any]							=     create_text_encoder_lora_attn_procs(pipe.text_encoder		)
       set_lora_up_weights(_snake_case		,		randn_weight=_snake_case		)

       # monkey patch
       pipe._modify_text_encoder(_snake_case		)

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

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

       # remove monkey patch
       pipe._remove_text_encoder_monkey_patch()

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

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






      def      UpperCAmelCase_     (self					:    str		)       ->   str:


       """simple docstring"""






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

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

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

       lowerCamelCase_							:    Optional[Any]							=     sd_pipe(**_snake_case		).images
       lowerCamelCase_							:    List[str]							=     lora_images[0, -3:, -3:, -1]

       lowerCamelCase_							:    Tuple							=     sd_pipe(**_snake_case		,		cross_attention_kwargs={'scale': 0.5}		).images
       lowerCamelCase_							:    List[Any]							=     lora_images_with_scale[0, -3:, -3:, -1]

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






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


       """simple docstring"""






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

        lowerCamelCase_		,      lowerCamelCase_							:    Optional[int]							=     self.get_dummy_components()
        lowerCamelCase_							:    Any							=     StableDiffusionPipeline(**_snake_case		)
        lowerCamelCase_							:    List[Any]							=     sd_pipe.to(_snake_case		)
        sd_pipe.set_progress_bar_config(disable=_snake_case		)

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

            # load LoRA weight file
        sd_pipe.load_lora_weights(_snake_case		)

        # check if lora attention processors are used
        for _, module in sd_pipe.unet.named_modules():
         if isinstance(_snake_case		,		_snake_case		):
          lowerCamelCase_							:    Optional[int]							=     (
              LoRAAttnProcessor2_0 if hasattr(_snake_case		,		'scaled_dot_product_attention'		) else LoRAAttnProcessor
          )
          self.assertIsInstance(module.processor		,		_snake_case		)






      @unittest.skipIf(torch_device != 'cuda'		,		'This test is supposed to run on GPU'		)
      def      UpperCAmelCase_     (self					:    List[str]		)       ->   str:


       """simple docstring"""






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

        lowerCamelCase_		,      lowerCamelCase_							:    int							=     self.get_dummy_components()
        lowerCamelCase_							:    Tuple							=     StableDiffusionPipeline(**_snake_case		)
        lowerCamelCase_							:    Optional[int]							=     sd_pipe.to(_snake_case		)
        sd_pipe.set_progress_bar_config(disable=_snake_case		)

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

            # load LoRA weight file
        sd_pipe.load_lora_weights(_snake_case		)

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






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


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    List[str]							=     self.get_dummy_components()
       lowerCamelCase_							:    Any							=     StableDiffusionPipeline(**_snake_case		)
       lowerCamelCase_							:    str							=     sd_pipe.to(_snake_case		)
       sd_pipe.set_progress_bar_config(disable=_snake_case		)

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

       # enable XFormers
       sd_pipe.enable_xformers_memory_efficient_attention()

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

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

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

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








def       _a   (							lowerCamelCase__				)  ->      int:
 return int(lowerCamelCase__				) if x - int(lowerCamelCase__				) >= 0 else int(lowerCamelCase__				) - 1


if __name__ == "__main__":
       import doctest

       doctest.testmod()
