




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





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





def 						snake_case_							(_a     :							nn.Module					):
    UpperCAmelCase      	=		create_text_encoder_lora_attn_procs(_a					)
    UpperCAmelCase      	=		AttnProcsLayers(_a					)
    return text_encoder_lora_layers





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







class       _a  (     unittest.TestCase			):
       def    A				(     self						:     str    ):



           '''simple docstring'''

           torch.manual_seed(0    )
           UpperCAmelCase      	=		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       ,    )
           UpperCAmelCase      	=		DDIMScheduler(
               beta_start=0.0_0085       ,    beta_end=0.012       ,    beta_schedule='''scaled_linear'''       ,    clip_sample=lowercase       ,    set_alpha_to_one=lowercase       ,    steps_offset=1       ,    )
           torch.manual_seed(0    )
           UpperCAmelCase      	=		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       ,    )
           UpperCAmelCase      	=		CLIPTextConfig(
               bos_token_id=0       ,    eos_token_id=2       ,    hidden_size=32       ,    intermediate_size=37       ,    layer_norm_eps=1E-05       ,    num_attention_heads=4       ,    num_hidden_layers=5       ,    pad_token_id=1       ,    vocab_size=1_000       ,    )
           UpperCAmelCase      	=		CLIPTextModel(lowercase    )
           UpperCAmelCase      	=		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

           UpperCAmelCase	,					UpperCAmelCase      	=		create_unet_lora_layers(lowercase    )
           UpperCAmelCase      	=		create_text_encoder_lora_layers(lowercase    )

           UpperCAmelCase      	=		{
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           UpperCAmelCase      	=		{
               '''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    A				(     self						:     Optional[int]    ):



           '''simple docstring'''

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

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

           UpperCAmelCase      	=		{
               '''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    A				(     self						:     Dict    ):



           '''simple docstring'''

           UpperCAmelCase      	=		77

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

           UpperCAmelCase      	=		{}
           UpperCAmelCase      	=		inputs
           return prepared_inputs
       def    A				(     self						:     List[str]       ,    lowercase						:     str    ):



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		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    A				(     self						:     List[Any]    ):



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
           UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )
           UpperCAmelCase      	=		sd_pipe.to(lowercase    )
           sd_pipe.set_progress_bar_config(disable=lowercase    )

           UpperCAmelCase	,					UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_inputs()

           UpperCAmelCase      	=		sd_pipe(**lowercase    ).images
           UpperCAmelCase      	=		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    )

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

           # Outputs shouldn't match.
           self.assertFalse(torch.allclose(torch.from_numpy(lowercase    )       ,    torch.from_numpy(lowercase    )    )    )
       def    A				(     self						:     Optional[Any]    ):



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
           UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )
           UpperCAmelCase      	=		sd_pipe.to(lowercase    )
           sd_pipe.set_progress_bar_config(disable=lowercase    )

           UpperCAmelCase	,					UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_inputs()

           UpperCAmelCase      	=		sd_pipe(**lowercase    ).images
           UpperCAmelCase      	=		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    )

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

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



           '''simple docstring'''

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

           UpperCAmelCase	,					UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_inputs()

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

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

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

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



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
           UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )

           UpperCAmelCase      	=		self.get_dummy_tokens()

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

           # create lora_attn_procs with zeroed out up.weights
           UpperCAmelCase      	=		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
           UpperCAmelCase      	=		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
           UpperCAmelCase      	=		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
           UpperCAmelCase      	=		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    A				(     self						:     Tuple    ):



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
           UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )

           UpperCAmelCase      	=		self.get_dummy_tokens()

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

           # create lora_attn_procs with randn up.weights
           UpperCAmelCase      	=		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
           UpperCAmelCase      	=		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
           UpperCAmelCase      	=		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    A				(     self						:     List[Any]    ):



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
           UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )
           UpperCAmelCase      	=		sd_pipe.to(lowercase    )
           sd_pipe.set_progress_bar_config(disable=lowercase    )

           UpperCAmelCase	,					UpperCAmelCase	,					UpperCAmelCase      	=		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    )

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

           UpperCAmelCase      	=		sd_pipe(**lowercase       ,    cross_attention_kwargs={'''scale''': 0.5}    ).images
           UpperCAmelCase      	=		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    A				(     self						:     Union[str, Any]    ):



           '''simple docstring'''

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

               UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
               UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )
               UpperCAmelCase      	=		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    ):
                       UpperCAmelCase      	=		(
                           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    A				(     self						:     Tuple    ):



           '''simple docstring'''

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

               UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
               UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )
               UpperCAmelCase      	=		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    A				(     self						:     Tuple    ):



           '''simple docstring'''

           UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_components()
           UpperCAmelCase      	=		StableDiffusionPipeline(**lowercase    )
           UpperCAmelCase      	=		sd_pipe.to(lowercase    )
           sd_pipe.set_progress_bar_config(disable=lowercase    )

           UpperCAmelCase	,					UpperCAmelCase	,					UpperCAmelCase      	=		self.get_dummy_inputs()

           # enable XFormers
           sd_pipe.enable_xformers_memory_efficient_attention()

           UpperCAmelCase      	=		sd_pipe(**lowercase    ).images
           UpperCAmelCase      	=		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    )

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

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










'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A			   ={'configuration_dinat': ['DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DinatConfig']}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A			   =[
       'DINAT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'DinatForImageClassification',
       'DinatModel',
       'DinatPreTrainedModel',
       'DinatBackbone',
   ]

if TYPE_CHECKING:
   from .configuration_dinat import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP, DinatConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_dinat import (
          DINAT_PRETRAINED_MODEL_ARCHIVE_LIST,
          DinatBackbone,
          DinatForImageClassification,
          DinatModel,
          DinatPreTrainedModel,
      )

else:
   import sys

   A			   =_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




