

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__							(     _A :					nn.Module				):





  '''simple docstring'''







  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(_A      , '''scaled_dot_product_attention'''				) else LoRAAttnProcessor
    )
    a__    =lora_attn_processor_class(
        hidden_size=_A      , cross_attention_dim=_A				)
  a__    =AttnProcsLayers(_A				)
  return lora_attn_procs, unet_lora_layers




def      UpperCAmelCase__							(     _A :					nn.Module				):





  '''simple docstring'''







  a__    ={}
  a__    =(
      LoRAAttnProcessor2_0 if hasattr(_A      , '''scaled_dot_product_attention'''				) else LoRAAttnProcessor
  )
  for name, module in text_encoder.named_modules():
    if name.endswith(_A				):
      a__    =lora_attn_processor_class(
          hidden_size=module.out_proj.out_features      , cross_attention_dim=_A				)
  return text_lora_attn_procs




def      UpperCAmelCase__							(     _A :					nn.Module				):





  '''simple docstring'''







  a__    =create_text_encoder_lora_attn_procs(_A				)
  a__    =AttnProcsLayers(_A				)
  return text_encoder_lora_layers




def      UpperCAmelCase__							(     _A :					Union[str, Any]      , _A :					Optional[Any]=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'''				):
        a__    =(
            torch.randn_like(layer_module.up.weight				)
            if randn_weight
            else torch.zeros_like(layer_module.up.weight				)
        )
        a__    =torch.nn.Parameter(_A				)







class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """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.00085,       beta_end=0.012,       beta_schedule='''scaled_linear''',       clip_sample=lowercase_,       set_alpha_to_one=lowercase_,       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(lowercase_			)
     a__    =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

     a__, a__    =create_unet_lora_layers(lowercase_			)
     a__    =create_text_encoder_lora_layers(lowercase_			)

     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 					_UpperCAmelCase					(							self			)				->       int:






     """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,       lowercase_,       size=(batch_size, sequence_length),       generator=lowercase_			)

     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 					_UpperCAmelCase					(							self			)				->       str:






     """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 					_UpperCAmelCase					(							self,       lowercase_			)				->       str:






     """simple docstring"""






     a__, a__    =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 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__, a__    =self.get_dummy_components()
     a__    =StableDiffusionPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__, a__, a__    =self.get_dummy_inputs()

     a__    =sd_pipe(**lowercase_			).images
     a__    =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_			)

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

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



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__, a__    =self.get_dummy_components()
     a__    =StableDiffusionPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__, a__, a__    =self.get_dummy_inputs()

     a__    =sd_pipe(**lowercase_			).images
     a__    =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_			)

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

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



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






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

     a__, a__, a__    =self.get_dummy_inputs()

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

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

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

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



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__, a__    =self.get_dummy_components()
     a__    =StableDiffusionPipeline(**lowercase_			)

     a__    =self.get_dummy_tokens()

     # inference without lora
     a__    =pipe.text_encoder(**lowercase_			)[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(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
     a__    =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
     a__    =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
     a__    =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 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__, a__    =self.get_dummy_components()
     a__    =StableDiffusionPipeline(**lowercase_			)

     a__    =self.get_dummy_tokens()

     # inference without lora
     a__    =pipe.text_encoder(**lowercase_			)[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(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
     a__    =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
     a__    =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 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__, a__    =self.get_dummy_components()
     a__    =StableDiffusionPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__, a__, a__    =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_			)

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

     a__    =sd_pipe(**lowercase_,       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(lowercase_			),       torch.from_numpy(lowercase_			)			)			)



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






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

       a__, a__    =self.get_dummy_components()
       a__    =StableDiffusionPipeline(**lowercase_			)
       a__    =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_			):
           a__    =(
               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 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






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

       a__, a__    =self.get_dummy_components()
       a__    =StableDiffusionPipeline(**lowercase_			)
       a__    =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 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__, a__    =self.get_dummy_components()
     a__    =StableDiffusionPipeline(**lowercase_			)
     a__    =sd_pipe.to(lowercase_			)
     sd_pipe.set_progress_bar_config(disable=lowercase_			)

     a__, a__, a__    =self.get_dummy_inputs()

     # enable XFormers
     sd_pipe.enable_xformers_memory_efficient_attention()

     a__    =sd_pipe(**lowercase_			).images
     a__    =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_			)

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

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



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase 	=   logging.get_logger(__name__)

lowerCamelCase 	=   {
    '''facebook/sam-vit-huge''': '''https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json''',
    '''facebook/sam-vit-large''': '''https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json''',
    '''facebook/sam-vit-base''': '''https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json''',
}







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=256,       lowercase_=1024,       lowercase_=16,       lowercase_=16,       lowercase_=4,       lowercase_="gelu",       lowercase_=1E-6,       **lowercase_,       )				->       int:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =hidden_size
     a__    =image_size
     a__    =patch_size
     a__    =image_size // patch_size
     a__    =mask_input_channels
     a__    =num_point_embeddings
     a__    =hidden_act
     a__    =layer_norm_eps







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=256,       lowercase_="relu",       lowercase_=2048,       lowercase_=2,       lowercase_=8,       lowercase_=2,       lowercase_=3,       lowercase_=3,       lowercase_=256,       lowercase_=1E-6,       **lowercase_,       )				->       int:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =hidden_size
     a__    =hidden_act
     a__    =mlp_dim
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =attention_downsample_rate
     a__    =num_multimask_outputs
     a__    =iou_head_depth
     a__    =iou_head_hidden_dim
     a__    =layer_norm_eps







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=768,       lowercase_=256,       lowercase_=12,       lowercase_=12,       lowercase_=3,       lowercase_=1024,       lowercase_=16,       lowercase_="gelu",       lowercase_=1E-06,       lowercase_=0.0,       lowercase_=1E-10,       lowercase_=True,       lowercase_=4.0,       lowercase_=True,       lowercase_=True,       lowercase_=14,       lowercase_=[2, 5, 8, 11],       lowercase_=128,       lowercase_=None,       **lowercase_,       )				->       Union[str, Any]:






     """simple docstring"""






     super().__init__(**lowercase_			)

     a__    =hidden_size
     a__    =output_channels
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =num_channels
     a__    =image_size
     a__    =patch_size
     a__    =hidden_act
     a__    =layer_norm_eps
     a__    =attention_dropout
     a__    =initializer_range
     a__    =qkv_bias
     a__    =mlp_ratio
     a__    =use_abs_pos
     a__    =use_rel_pos
     a__    =window_size
     a__    =global_attn_indexes
     a__    =num_pos_feats
     a__    =int(hidden_size * mlp_ratio			) if mlp_dim is None else mlp_dim







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   lowerCamelCase__					:      Optional[int]        =  'sam'
   lowerCamelCase__					:      str        =  True



   def __init__(							self,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_=0.02,       **lowercase_,       )				->       Tuple:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =vision_config if vision_config is not None else {}
     a__    =prompt_encoder_config if prompt_encoder_config is not None else {}
     a__    =mask_decoder_config if mask_decoder_config is not None else {}

     if isinstance(lowercase_,       lowercase_			):
       a__    =vision_config.to_dict()
     if isinstance(lowercase_,       lowercase_			):
       a__    =prompt_encoder_config.to_dict()
     if isinstance(lowercase_,       lowercase_			):
       a__    =mask_decoder_config.to_dict()

     a__    =SamVisionConfig(**lowercase_			)
     a__    =SamPromptEncoderConfig(**lowercase_			)
     a__    =SamMaskDecoderConfig(**lowercase_			)
     a__    =initializer_range



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =copy.deepcopy(self.__dict__			)
     a__    =self.vision_config.to_dict()
     a__    =self.prompt_encoder_config.to_dict()
     a__    =self.mask_decoder_config.to_dict()
     a__    =self.__class__.model_type
     return output
