
"""simple docstring"""
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 :					Optional[Any]		):

  """simple docstring"""







  _snake_case  :    Dict            = {}
  for name in unet.attn_processors.keys():
    _snake_case  :    Any            = None if name.endswith("attn1.processor"		) else unet.config.cross_attention_dim
    if name.startswith("mid_block"		):
      _snake_case  :    Optional[Any]            = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"		):
      _snake_case  :    int            = int(name[len("up_blocks."		)]		)
      _snake_case  :    Any            = list(reversed(unet.config.block_out_channels		)		)[block_id]
    elif name.startswith("down_blocks"		):
      _snake_case  :    Optional[Any]            = int(name[len("down_blocks."		)]		)
      _snake_case  :    Dict            = unet.config.block_out_channels[block_id]
    _snake_case  :    int            = (
        LoRAAttnProcessor2_0 if hasattr(__lowercase							,  "scaled_dot_product_attention"		) else LoRAAttnProcessor
    )
    _snake_case  :    Optional[Any]            = lora_attn_processor_class(
        hidden_size=__lowercase							,  cross_attention_dim=__lowercase		)
  _snake_case  :    List[str]            = AttnProcsLayers(__lowercase		)
  return lora_attn_procs, unet_lora_layers
def 					a__ ( a :					Dict		):

  """simple docstring"""







  _snake_case  :    List[Any]            = {}
  _snake_case  :    Union[str, Any]            = (
      LoRAAttnProcessor2_0 if hasattr(__lowercase							,  "scaled_dot_product_attention"		) else LoRAAttnProcessor
  )
  for name, module in text_encoder.named_modules():
    if name.endswith(__lowercase		):
      _snake_case  :    Optional[Any]            = lora_attn_processor_class(
          hidden_size=module.out_proj.out_features							,  cross_attention_dim=__lowercase		)
  return text_lora_attn_procs
def 					a__ ( a :					List[Any]		):

  """simple docstring"""







  _snake_case  :    List[Any]            = create_text_encoder_lora_attn_procs(__lowercase		)
  _snake_case  :    Dict            = AttnProcsLayers(__lowercase		)
  return text_encoder_lora_layers
def 					a__ ( a :					List[Any]							,  a :					int=False		):

  """simple docstring"""







  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"		):
        _snake_case  :    int            = (
            torch.randn_like(layer_module.up.weight		)
            if randn_weight
            else torch.zeros_like(layer_module.up.weight		)
        )
        _snake_case  :    List[str]            = torch.nn.Parameter(__lowercase		)






class   _UpperCAmelCase					(       unittest.TestCase):
    def        lowerCamelCase__						(     self  ):
      torch.manual_seed(0  )
      _snake_case  :    List[str]            = 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    ,     )
      _snake_case  :    List[Any]            = DDIMScheduler(
          beta_start=0.00085    ,     beta_end=0.012    ,     beta_schedule="scaled_linear"    ,     clip_sample=a__    ,     set_alpha_to_one=a__    ,     steps_offset=1    ,     )
      torch.manual_seed(0  )
      _snake_case  :    Optional[int]            = 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    ,     )
      _snake_case  :    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=10_00    ,     )
      _snake_case  :    List[Any]            = CLIPTextModel(a__  )
      _snake_case  :    Any            = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

      _snake_case      ,       _snake_case  :    Optional[int]            = create_unet_lora_layers(a__  )
      _snake_case  :    Tuple            = create_text_encoder_lora_layers(a__  )

      _snake_case  :    Union[str, Any]            = {
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "safety_checker": None,
          "feature_extractor": None,
      }
      _snake_case  :    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        lowerCamelCase__						(     self  ):
      _snake_case  :    Optional[Any]            = 1
      _snake_case  :    int            = 10
      _snake_case  :    Union[str, Any]            = 4
      _snake_case  :    Tuple            = (32, 32)

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

      _snake_case  :    str            = {
          "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  ):
      _snake_case  :    str            = 77

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

      _snake_case  :    str            = {}
      _snake_case  :    Optional[Any]            = inputs
      return prepared_inputs
    def        lowerCamelCase__						(     self    ,     snake_case_  ):
      _snake_case      ,       _snake_case  :    Optional[Any]            = 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  ):
      _snake_case      ,       _snake_case  :    Optional[int]            = self.get_dummy_components()
      _snake_case  :    int            = StableDiffusionPipeline(**a__  )
      _snake_case  :    Tuple            = sd_pipe.to(a__  )
      sd_pipe.set_progress_bar_config(disable=a__  )

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

      _snake_case  :    List[Any]            = sd_pipe(**a__  ).images
      _snake_case  :    Optional[int]            = 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__  )

      _snake_case  :    Dict            = sd_pipe(**a__  ).images
      _snake_case  :    int            = 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  ):
      _snake_case      ,       _snake_case  :    List[str]            = self.get_dummy_components()
      _snake_case  :    str            = StableDiffusionPipeline(**a__  )
      _snake_case  :    str            = sd_pipe.to(a__  )
      sd_pipe.set_progress_bar_config(disable=a__  )

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

      _snake_case  :    List[str]            = sd_pipe(**a__  ).images
      _snake_case  :    str            = 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__  )

      _snake_case  :    Union[str, Any]            = sd_pipe(**a__  ).images
      _snake_case  :    Optional[Any]            = 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  ):
      _snake_case      ,       _snake_case  :    Union[str, Any]            = self.get_dummy_components()
      _snake_case  :    Tuple            = lora_components["unet_lora_attn_procs"]
      _snake_case  :    Optional[int]            = StableDiffusionPipeline(**a__  )
      _snake_case  :    str            = sd_pipe.to(a__  )
      sd_pipe.set_progress_bar_config(disable=a__  )

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

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

      with tempfile.TemporaryDirectory() as tmpdirname:
        _snake_case  :    Tuple            = 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__  )

      _snake_case  :    Tuple            = sd_pipe(**a__  ).images
      _snake_case  :    List[Any]            = 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  ):
      _snake_case      ,       _snake_case  :    Optional[Any]            = self.get_dummy_components()
      _snake_case  :    int            = StableDiffusionPipeline(**a__  )

      _snake_case  :    List[str]            = self.get_dummy_tokens()

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

      # create lora_attn_procs with zeroed out up.weights
      _snake_case  :    List[str]            = 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
      _snake_case  :    int            = 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
      _snake_case  :    Optional[Any]            = 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
      _snake_case  :    Union[str, Any]            = 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  ):
      _snake_case      ,       _snake_case  :    List[str]            = self.get_dummy_components()
      _snake_case  :    Tuple            = StableDiffusionPipeline(**a__  )

      _snake_case  :    str            = self.get_dummy_tokens()

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

      # create lora_attn_procs with randn up.weights
      _snake_case  :    str            = 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
      _snake_case  :    Optional[int]            = 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
      _snake_case  :    Tuple            = 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  ):
      _snake_case      ,       _snake_case  :    Dict            = self.get_dummy_components()
      _snake_case  :    Optional[int]            = StableDiffusionPipeline(**a__  )
      _snake_case  :    str            = sd_pipe.to(a__  )
      sd_pipe.set_progress_bar_config(disable=a__  )

      _snake_case      ,       _snake_case      ,       _snake_case  :    Optional[Any]            = 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__  )

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

      _snake_case  :    Dict            = sd_pipe(**a__    ,     cross_attention_kwargs={"scale": 0.5}  ).images
      _snake_case  :    Tuple            = 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  ):







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

        _snake_case      ,       _snake_case  :    Any            = self.get_dummy_components()
        _snake_case  :    Optional[int]            = StableDiffusionPipeline(**a__  )
        _snake_case  :    Any            = 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__  ):
            _snake_case  :    Optional[int]            = (
                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  ):







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

        _snake_case      ,       _snake_case  :    int            = self.get_dummy_components()
        _snake_case  :    Optional[int]            = StableDiffusionPipeline(**a__  )
        _snake_case  :    Union[str, Any]            = 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  ):
      _snake_case      ,       _snake_case  :    List[str]            = self.get_dummy_components()
      _snake_case  :    List[str]            = StableDiffusionPipeline(**a__  )
      _snake_case  :    Optional[int]            = sd_pipe.to(a__  )
      sd_pipe.set_progress_bar_config(disable=a__  )

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

      # enable XFormers
      sd_pipe.enable_xformers_memory_efficient_attention()

      _snake_case  :    Any            = sd_pipe(**a__  ).images
      _snake_case  :    Optional[Any]            = 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__  )

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

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





"""simple docstring"""


import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import Mock, patch

import diffusers.utils.hub_utils
class snake_case   (	unittest.TestCase):







		@patch("diffusers.utils.hub_utils.get_full_repo_name"		)
		def 			a_			(       self   :							Any     , a__   :							Mock		)							->    None:





						'''simple docstring'''







						_A       	=      "full_repo_name"
						with TemporaryDirectory() as tmpdir:
										# Dummy args values
										_A       	=      Mock()
										_A       	=      tmpdir
										_A       	=      0
										_A       	=      "hub_token"
										_A       	=      "dataset_name"
										_A       	=      0.0_1
										_A       	=      10_00_00
										_A       	=      1_00_00
										_A       	=      0.0_1
										_A       	=      0.0_2
										_A       	=      0.0_3
										_A       	=      0.0_0_0_5
										_A       	=      0.0_0_0_0_0_1
										_A       	=      1
										_A       	=      10
										_A       	=      0.0_0_1
										_A       	=      0.1
										_A       	=      0.2
										_A       	=      True

										# Model card mush be rendered and saved
										diffusers.utils.hub_utils.create_model_card(a__     , model_name="model_name"		)
										self.assertTrue((Path(a__		) / "README.md").is_file()		)