




import argparse

import torch

from diffusers import UNetaDConditionModel

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase	,      __UpperCamelCase	,      __UpperCamelCase=None	,      __UpperCamelCase=None	,      __UpperCamelCase=None)    ->  Union[str, Any]:
     assert isinstance(_A	,      _A), "Paths should be a list of dicts containing 'old' and 'new' keys."

     # Splits the attention layers into three variables.
     if attention_paths_to_split is not None:
          for path, path_map in attention_paths_to_split.items():
               a											=   old_checkpoint[path]
               a											=   old_tensor.shape[0] // 3

               a											=   (-1, channels) if len(old_tensor.shape) == 3 else (-1)

               a											=   old_tensor.shape[0] // config["num_head_channels"] // 3

               a											=   old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
               a						,   a						,   a											=   old_tensor.split(channels // num_heads	,      dim=1)

               a											=   query.reshape(_A)
               a											=   key.reshape(_A)
               a											=   value.reshape(_A)

     for path in paths:
          a											=   path["new"]

          # These have already been assigned
          if attention_paths_to_split is not None and new_path in attention_paths_to_split:
               continue

          if additional_replacements is not None:
               for replacement in additional_replacements:
                    a											=   new_path.replace(replacement["old"]	,      replacement["new"])

                    # proj_attn.weight has to be converted from conv 1D to linear
          a											=   old_checkpoint[path["old"]]
          a											=   ["proj_attn.weight"]
          a											=   ["proj_out.weight", "proj_in.weight"]
          if any(k in new_path for k in names):
               a											=   weight[:, :, 0]
          elif any(k in new_path for k in names_a) and len(weight.shape) > 2 and ".attentions." not in new_path:
               a											=   weight[:, :, 0]
          else:
               a											=   weight

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase=0)    ->  str:
     a											=   []
     for old_item in old_list:
          a											=   old_item

          #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
          #         new_item = new_item.replace('norm.bias', 'group_norm.bias')

          #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
          #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')

          #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)

          mapping.append({"old": old_item, "new": new_item})

     return mapping

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase=1)    ->  Optional[int]:







     if n_shave_prefix_segments >= 0:
          return ".".join(path.split(".")[n_shave_prefix_segments:])
     else:
          return ".".join(path.split(".")[:n_shave_prefix_segments])

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase=0)    ->  Union[str, Any]:
     a											=   []
     for old_item in old_list:
          mapping.append({"old": old_item, "new": old_item})

     return mapping

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase=0)    ->  List[Any]:
     a											=   []
     for old_item in old_list:
          a											=   old_item.replace("in_layers.0"	,      "norm1")
          a											=   new_item.replace("in_layers.2"	,      "conv1")

          a											=   new_item.replace("out_layers.0"	,      "norm2")
          a											=   new_item.replace("out_layers.3"	,      "conv2")

          a											=   new_item.replace("emb_layers.1"	,      "time_emb_proj")
          a											=   new_item.replace("skip_connection"	,      "conv_shortcut")

          a											=   shave_segments(_A	,      n_shave_prefix_segments=_A)

          if "temopral_conv" not in old_item:
               mapping.append({"old": old_item, "new": new_item})

     return mapping

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase	,      __UpperCamelCase=None	,      __UpperCamelCase=False)    ->  int:
     a											=   {}
     a											=   list(checkpoint.keys())

     a											=   "model.diffusion_model."

     # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
     if sum(k.startswith("model_ema") for k in keys) > 1_00 and extract_ema:
          print(f'''Checkpoint {path} has both EMA and non-EMA weights.''')
          print(
              "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"
              " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag.")
          for key in keys:
               if key.startswith("model.diffusion_model"):
                    a											=   "model_ema." + "".join(key.split(".")[1:])
                    a											=   checkpoint.pop(_A)
     else:
          if sum(k.startswith("model_ema") for k in keys) > 1_00:
               print(
                   "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"
                   " weights (usually better for inference), please make sure to add the `--extract_ema` flag.")

          for key in keys:
               a											=   checkpoint.pop(_A)

     a											=   {}

     a											=   unet_state_dict["time_embed.0.weight"]
     a											=   unet_state_dict["time_embed.0.bias"]
     a											=   unet_state_dict["time_embed.2.weight"]
     a											=   unet_state_dict["time_embed.2.bias"]

     if config["class_embed_type"] is None:
          # No parameters to port
          ...
     elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
          a											=   unet_state_dict["label_emb.0.0.weight"]
          a											=   unet_state_dict["label_emb.0.0.bias"]
          a											=   unet_state_dict["label_emb.0.2.weight"]
          a											=   unet_state_dict["label_emb.0.2.bias"]
     else:
          raise NotImplementedError(f'''Not implemented `class_embed_type`: {config["class_embed_type"]}''')

     a											=   unet_state_dict["input_blocks.0.0.weight"]
     a											=   unet_state_dict["input_blocks.0.0.bias"]

     a											=   [v for v in unet_state_dict if v.startswith("input_blocks.0.1")]
     a											=   renew_attention_paths(_A)
     a											=   {"old": "input_blocks.0.1", "new": "transformer_in"}
     assign_to_checkpoint(_A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

     a											=   unet_state_dict["out.0.weight"]
     a											=   unet_state_dict["out.0.bias"]
     a											=   unet_state_dict["out.2.weight"]
     a											=   unet_state_dict["out.2.bias"]

     # Retrieves the keys for the input blocks only
     a											=   len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
     a											=   {
         layer_id: [key for key in unet_state_dict if f'''input_blocks.{layer_id}''' in key]
         for layer_id in range(_A)
     }

     # Retrieves the keys for the middle blocks only
     a											=   len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
     a											=   {
         layer_id: [key for key in unet_state_dict if f'''middle_block.{layer_id}''' in key]
         for layer_id in range(_A)
     }

     # Retrieves the keys for the output blocks only
     a											=   len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
     a											=   {
         layer_id: [key for key in unet_state_dict if f'''output_blocks.{layer_id}''' in key]
         for layer_id in range(_A)
     }

     for i in range(1	,      _A):
          a											=   (i - 1) // (config["layers_per_block"] + 1)
          a											=   (i - 1) % (config["layers_per_block"] + 1)

          a											=   [
              key for key in input_blocks[i] if f'''input_blocks.{i}.0''' in key and f'''input_blocks.{i}.0.op''' not in key
          ]
          a											=   [key for key in input_blocks[i] if f'''input_blocks.{i}.1''' in key]
          a											=   [key for key in input_blocks[i] if f'''input_blocks.{i}.2''' in key]

          if f'''input_blocks.{i}.op.weight''' in unet_state_dict:
               a											=   unet_state_dict.pop(
                   f'''input_blocks.{i}.op.weight''')
               a											=   unet_state_dict.pop(
                   f'''input_blocks.{i}.op.bias''')

          a											=   renew_resnet_paths(_A)
          a											=   {"old": f'''input_blocks.{i}.0''', "new": f'''down_blocks.{block_id}.resnets.{layer_in_block_id}'''}
          assign_to_checkpoint(
              _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

          a											=   [key for key in resnets if "temopral_conv" in key]
          a											=   renew_temp_conv_paths(_A)
          a											=   {
              "old": f'''input_blocks.{i}.0.temopral_conv''',
              "new": f'''down_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
          }
          assign_to_checkpoint(
              _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

          if len(_A):
               a											=   renew_attention_paths(_A)
               a											=   {"old": f'''input_blocks.{i}.1''', "new": f'''down_blocks.{block_id}.attentions.{layer_in_block_id}'''}
               assign_to_checkpoint(
                   _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

          if len(_A):
               a											=   renew_attention_paths(_A)
               a											=   {
                   "old": f'''input_blocks.{i}.2''',
                   "new": f'''down_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
               }
               assign_to_checkpoint(
                   _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

     a											=   middle_blocks[0]
     a											=   [key for key in resnet_a if "temopral_conv" in key]
     a											=   middle_blocks[1]
     a											=   middle_blocks[2]
     a											=   middle_blocks[3]
     a											=   [key for key in resnet_a if "temopral_conv" in key]

     a											=   renew_resnet_paths(_A)
     a											=   {"old": "middle_block.0", "new": "mid_block.resnets.0"}
     assign_to_checkpoint(
         _A	,      _A	,      _A	,      config=_A	,      additional_replacements=[meta_path])

     a											=   renew_temp_conv_paths(_A)
     a											=   {"old": "middle_block.0.temopral_conv", "new": "mid_block.temp_convs.0"}
     assign_to_checkpoint(
         _A	,      _A	,      _A	,      config=_A	,      additional_replacements=[meta_path])

     a											=   renew_resnet_paths(_A)
     a											=   {"old": "middle_block.3", "new": "mid_block.resnets.1"}
     assign_to_checkpoint(
         _A	,      _A	,      _A	,      config=_A	,      additional_replacements=[meta_path])

     a											=   renew_temp_conv_paths(_A)
     a											=   {"old": "middle_block.3.temopral_conv", "new": "mid_block.temp_convs.1"}
     assign_to_checkpoint(
         _A	,      _A	,      _A	,      config=_A	,      additional_replacements=[meta_path])

     a											=   renew_attention_paths(_A)
     a											=   {"old": "middle_block.1", "new": "mid_block.attentions.0"}
     assign_to_checkpoint(
         _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

     a											=   renew_attention_paths(_A)
     a											=   {"old": "middle_block.2", "new": "mid_block.temp_attentions.0"}
     assign_to_checkpoint(
         _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

     for i in range(_A):
          a											=   i // (config["layers_per_block"] + 1)
          a											=   i % (config["layers_per_block"] + 1)
          a											=   [shave_segments(_A	,      2) for name in output_blocks[i]]
          a											=   {}

          for layer in output_block_layers:
               a						,   a											=   layer.split(".")[0], shave_segments(_A	,      1)
               if layer_id in output_block_list:
                    output_block_list[layer_id].append(_A)
               else:
                    a											=   [layer_name]

          if len(_A) > 1:
               a											=   [key for key in output_blocks[i] if f'''output_blocks.{i}.0''' in key]
               a											=   [key for key in output_blocks[i] if f'''output_blocks.{i}.1''' in key]
               a											=   [key for key in output_blocks[i] if f'''output_blocks.{i}.2''' in key]

               a											=   renew_resnet_paths(_A)
               a											=   renew_resnet_paths(_A)

               a											=   {"old": f'''output_blocks.{i}.0''', "new": f'''up_blocks.{block_id}.resnets.{layer_in_block_id}'''}
               assign_to_checkpoint(
                   _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

               a											=   [key for key in resnets if "temopral_conv" in key]
               a											=   renew_temp_conv_paths(_A)
               a											=   {
                   "old": f'''output_blocks.{i}.0.temopral_conv''',
                   "new": f'''up_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
               }
               assign_to_checkpoint(
                   _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

               a											=   {k: sorted(_A) for k, v in output_block_list.items()}
               if ["conv.bias", "conv.weight"] in output_block_list.values():
                    a											=   list(output_block_list.values()).index(["conv.bias", "conv.weight"])
                    a											=   unet_state_dict[
                        f'''output_blocks.{i}.{index}.conv.weight'''
                    ]
                    a											=   unet_state_dict[
                        f'''output_blocks.{i}.{index}.conv.bias'''
                    ]

                    # Clear attentions as they have been attributed above.
                    if len(_A) == 2:
                         a											=   []

               if len(_A):
                    a											=   renew_attention_paths(_A)
                    a											=   {
                        "old": f'''output_blocks.{i}.1''',
                        "new": f'''up_blocks.{block_id}.attentions.{layer_in_block_id}''',
                    }
                    assign_to_checkpoint(
                        _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)

               if len(_A):
                    a											=   renew_attention_paths(_A)
                    a											=   {
                        "old": f'''output_blocks.{i}.2''',
                        "new": f'''up_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
                    }
                    assign_to_checkpoint(
                        _A	,      _A	,      _A	,      additional_replacements=[meta_path]	,      config=_A)
          else:
               a											=   renew_resnet_paths(_A	,      n_shave_prefix_segments=1)
               for path in resnet_0_paths:
                    a											=   ".".join(["output_blocks", str(_A), path["old"]])
                    a											=   ".".join(["up_blocks", str(_A), "resnets", str(_A), path["new"]])
                    a											=   unet_state_dict[old_path]

               a											=   [l for l in output_block_layers if "temopral_conv" in l]
               for path in temopral_conv_paths:
                    a											=   path.split("temopral_conv.")[-1]
                    a											=   ".".join(["output_blocks", str(_A), str(_A), "temopral_conv", pruned_path])
                    a											=   ".".join(["up_blocks", str(_A), "temp_convs", str(_A), pruned_path])
                    a											=   unet_state_dict[old_path]

     return new_checkpoint


if __name__ == "__main__":
   lowercase__   :    str 					=      argparse.ArgumentParser()

   parser.add_argument(
       "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert."
   )
   parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
   lowercase__   :    Union[str, Any] 					=      parser.parse_args()

   lowercase__   :    str 					=      torch.load(args.checkpoint_path, map_location="cpu")
   lowercase__   :    int 					=      UNetaDConditionModel()

   lowercase__   :    Dict 					=      convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)

   lowercase__   :    Dict 					=      set(unet.state_dict().keys()) - set(converted_ckpt.keys())
   lowercase__   :    List[Any] 					=      set(converted_ckpt.keys()) - set(unet.state_dict().keys())

   assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"

   # load state_dict
   unet.load_state_dict(converted_ckpt)

   unet.save_pretrained(args.dump_path)

   # -- finish converting the unet --







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				)				)				)

