





'''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 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[Any]				):
   lowercase_    :str					    =				{}
   for name in unet.attn_processors.keys():
      lowercase_    :List[str]					    =				None if name.endswith("attn1.processor"				) else unet.config.cross_attention_dim
      if name.startswith("mid_block"				):
         lowercase_    :Optional[int]					    =				unet.config.block_out_channels[-1]
      elif name.startswith("up_blocks"				):
         lowercase_    :List[str]					    =				int(name[len("up_blocks."				)]				)
         lowercase_    :Any					    =				list(reversed(unet.config.block_out_channels				)				)[block_id]
      elif name.startswith("down_blocks"				):
         lowercase_    :int					    =				int(name[len("down_blocks."				)]				)
         lowercase_    :List[Any]					    =				unet.config.block_out_channels[block_id]
      lowercase_    :List[Any]					    =				(
          LoRAAttnProcessor2_0 if hasattr(__lowerCamelCase			,"scaled_dot_product_attention"				) else LoRAAttnProcessor
      )
      lowercase_    :int					    =				lora_attn_processor_class(
          hidden_size=__lowerCamelCase			,cross_attention_dim=__lowerCamelCase				)
   lowercase_    :Optional[int]					    =				AttnProcsLayers(__lowerCamelCase				)
   return lora_attn_procs, unet_lora_layers
def 							UpperCAmelCase_						(		__lowerCamelCase      :      List[str]				):
   lowercase_    :Any					    =				{}
   lowercase_    :List[Any]					    =				(
       LoRAAttnProcessor2_0 if hasattr(__lowerCamelCase			,"scaled_dot_product_attention"				) else LoRAAttnProcessor
   )
   for name, module in text_encoder.named_modules():
      if name.endswith(__lowerCamelCase				):
         lowercase_    :Dict					    =				lora_attn_processor_class(
             hidden_size=module.out_proj.out_features			,cross_attention_dim=__lowerCamelCase				)
   return text_lora_attn_procs
def 							UpperCAmelCase_						(		__lowerCamelCase      :      List[str]				):
   lowercase_    :Optional[int]					    =				create_text_encoder_lora_attn_procs(__lowerCamelCase				)
   lowercase_    :List[Any]					    =				AttnProcsLayers(__lowerCamelCase				)
   return text_encoder_lora_layers
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Dict			,__lowerCamelCase      :      Dict=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"				):
            lowercase_    :Union[str, Any]					    =				(
                torch.randn_like(layer_module.up.weight				)
                if randn_weight
                else torch.zeros_like(layer_module.up.weight				)
            )
            lowercase_    :Dict					    =				torch.nn.Parameter(__lowerCamelCase				)


class    a_     (       unittest.TestCase							):





     def 							lowercase__   (				self	:      List[Any]		):





        """simple docstring"""




        torch.manual_seed(0		)
        lowercase_    :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			,   )
        lowercase_    :Union[str, Any]					    =				DDIMScheduler(
            beta_start=0.0_00_85			,   beta_end=0.0_12			,   beta_schedule="scaled_linear"			,   clip_sample=lowercase			,   set_alpha_to_one=lowercase			,   steps_offset=1			,   )
        torch.manual_seed(0		)
        lowercase_    :Dict					    =				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			,   )
        lowercase_    :Union[str, Any]					    =				CLIPTextConfig(
            bos_token_id=0			,   eos_token_id=2			,   hidden_size=32			,   intermediate_size=37			,   layer_norm_eps=1e-0_5			,   num_attention_heads=4			,   num_hidden_layers=5			,   pad_token_id=1			,   vocab_size=1_000			,   )
        lowercase_    :List[str]					    =				CLIPTextModel(lowercase		)
        lowercase_    :Optional[Any]					    =				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

        lowercase_    ,							lowercase_    :Optional[Any]					    =				create_unet_lora_layers(lowercase		)
        lowercase_    :Union[str, Any]					    =				create_text_encoder_lora_layers(lowercase		)

        lowercase_    :List[Any]					    =				{
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        lowercase_    :Optional[int]					    =				{
            "unet_lora_layers": unet_lora_layers,
            "text_encoder_lora_layers": text_encoder_lora_layers,
            "unet_lora_attn_procs": unet_lora_attn_procs,
        }
        return pipeline_components, lora_components





     def 							lowercase__   (				self	:      List[str]		):





        """simple docstring"""




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

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

        lowercase_    :List[Any]					    =				{
            "prompt": "A painting of a squirrel eating a burger",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 6.0,
            "output_type": "numpy",
        }

        return noise, input_ids, pipeline_inputs

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





     def 							lowercase__   (				self	:      Dict		):





        """simple docstring"""




        lowercase_    :Tuple					    =				77

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

        lowercase_    :str					    =				{}
        lowercase_    :Any					    =				inputs
        return prepared_inputs





     def 							lowercase__   (				self	:      int			,   lowercase	:      str		):





        """simple docstring"""




        lowercase_    ,							lowercase_    :int					    =				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 							lowercase__   (				self	:      str		):





        """simple docstring"""




        lowercase_    ,							lowercase_    :str					    =				self.get_dummy_components()
        lowercase_    :str					    =				StableDiffusionPipeline(**lowercase		)
        lowercase_    :Any					    =				sd_pipe.to(lowercase		)
        sd_pipe.set_progress_bar_config(disable=lowercase		)

        lowercase_    ,							lowercase_    ,							lowercase_    :List[str]					    =				self.get_dummy_inputs()

        lowercase_    :List[Any]					    =				sd_pipe(**lowercase		).images
        lowercase_    :Tuple					    =				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		)

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

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





     def 							lowercase__   (				self	:      Union[str, Any]		):





        """simple docstring"""




        lowercase_    ,							lowercase_    :int					    =				self.get_dummy_components()
        lowercase_    :Dict					    =				StableDiffusionPipeline(**lowercase		)
        lowercase_    :str					    =				sd_pipe.to(lowercase		)
        sd_pipe.set_progress_bar_config(disable=lowercase		)

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

        lowercase_    :Union[str, Any]					    =				sd_pipe(**lowercase		).images
        lowercase_    :Union[str, Any]					    =				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		)

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

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





     def 							lowercase__   (				self	:      int		):





        """simple docstring"""




        lowercase_    ,							lowercase_    :Dict					    =				self.get_dummy_components()
        lowercase_    :Dict					    =				lora_components["unet_lora_attn_procs"]
        lowercase_    :List[str]					    =				StableDiffusionPipeline(**lowercase		)
        lowercase_    :List[str]					    =				sd_pipe.to(lowercase		)
        sd_pipe.set_progress_bar_config(disable=lowercase		)

        lowercase_    ,							lowercase_    ,							lowercase_    :List[str]					    =				self.get_dummy_inputs()

        lowercase_    :Tuple					    =				sd_pipe(**lowercase		).images
        lowercase_    :int					    =				original_images[0, -3:, -3:, -1]

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

        lowercase_    :Union[str, Any]					    =				sd_pipe(**lowercase		).images
        lowercase_    :str					    =				lora_images[0, -3:, -3:, -1]

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





     def 							lowercase__   (				self	:      Optional[Any]		):





        """simple docstring"""




        lowercase_    ,							lowercase_    :int					    =				self.get_dummy_components()
        lowercase_    :Any					    =				StableDiffusionPipeline(**lowercase		)

        lowercase_    :int					    =				self.get_dummy_tokens()

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

        # create lora_attn_procs with zeroed out up.weights
        lowercase_    :List[Any]					    =				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
        lowercase_    :List[str]					    =				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
        lowercase_    :Tuple					    =				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
        lowercase_    :Optional[Any]					    =				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 							lowercase__   (				self	:      Optional[int]		):





        """simple docstring"""




        lowercase_    ,							lowercase_    :str					    =				self.get_dummy_components()
        lowercase_    :Any					    =				StableDiffusionPipeline(**lowercase		)

        lowercase_    :str					    =				self.get_dummy_tokens()

        # inference without lora
        lowercase_    :Union[str, Any]					    =				pipe.text_encoder(**lowercase		)[0]
        assert outputs_without_lora.shape == (1, 77, 32)

        # create lora_attn_procs with randn up.weights
        lowercase_    :Optional[int]					    =				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
        lowercase_    :List[Any]					    =				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
        lowercase_    :int					    =				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 							lowercase__   (				self	:      Union[str, Any]		):





        """simple docstring"""




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

        lowercase_    ,							lowercase_    ,							lowercase_    :Tuple					    =				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		)

        lowercase_    :Tuple					    =				sd_pipe(**lowercase		).images
        lowercase_    :List[str]					    =				lora_images[0, -3:, -3:, -1]

        lowercase_    :str					    =				sd_pipe(**lowercase			,   cross_attention_kwargs={"scale": 0.5}		).images
        lowercase_    :int					    =				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 							lowercase__   (				self	:      Any		):





        """simple docstring"""




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

           lowercase_    ,							lowercase_    :Union[str, Any]					    =				self.get_dummy_components()
           lowercase_    :Union[str, Any]					    =				StableDiffusionPipeline(**lowercase		)
           lowercase_    :Optional[Any]					    =				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		):
                 lowercase_    :List[Any]					    =				(
                     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 							lowercase__   (				self	:      str		):





        """simple docstring"""




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

           lowercase_    ,							lowercase_    :int					    =				self.get_dummy_components()
           lowercase_    :Union[str, Any]					    =				StableDiffusionPipeline(**lowercase		)
           lowercase_    :List[Any]					    =				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 							lowercase__   (				self	:      int		):





        """simple docstring"""




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

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

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

        lowercase_    :Any					    =				sd_pipe(**lowercase		).images
        lowercase_    :Union[str, Any]					    =				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		)

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

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









'''simple docstring'''


import platform
from argparse import ArgumentParser

import huggingface_hub
import pandas
import pyarrow

from datasets import __version__ as version
from datasets.commands import BaseDatasetsCLICommand
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]				):
   return EnvironmentCommand()


class    a_     (       _lowerCAmelCase							):





     @staticmethod
     def 							lowercase__   (				lowercase	:      ArgumentParser		):





        """simple docstring"""




        lowercase_    :Optional[int]					    =				parser.add_parser("env"			,   help="Print relevant system environment info."		)
        download_parser.set_defaults(func=lowercase		)





     def 							lowercase__   (				self	:      List[str]		):





        """simple docstring"""




        lowercase_    :Tuple					    =				{
            "`datasets` version": version,
            "Platform": platform.platform(),
            "Python version": platform.python_version(),
            "Huggingface_hub version": huggingface_hub.__version__,
            "PyArrow version": pyarrow.__version__,
            "Pandas version": pandas.__version__,
        }

        print("\nCopy-and-paste the text below in your GitHub issue.\n"		)
        print(self.format_dict(lowercase		)		)

        return info







     @staticmethod
     def 							lowercase__   (				lowercase	:      Dict		):





        """simple docstring"""




        return "\n".join([F'- {prop}: {val}' for prop, val in d.items()]		) + "\n"


