


"""simple docstring"""





import argparse
import time

import torch
import transformers
from measures_util import end_measure, log_measures, start_measure
from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeqaSeqLM, AutoTokenizer

from accelerate.utils import compute_module_sizes


snake_case__		:  Optional[Any]     					=					{
    '''gpt-j-6b''': {'''is_causal''': True, '''model''': '''sgugger/sharded-gpt-j-6B''', '''tokenizer''': '''EleutherAI/gpt-j-6B'''},
    '''gpt-neox''': {'''is_causal''': True, '''model''': '''EleutherAI/gpt-neox-20b'''},
    '''opt''': {'''is_causal''': True, '''model''': '''facebook/opt-30b'''},
    '''T0pp''': {'''is_causal''': False, '''model''': '''bigscience/T0pp''', '''model_revision''': '''sharded'''},
}

snake_case__		:  Any     					=					[
    '''Hello, my name is''',
    '''Are unicorns real? Unicorns are''',
    '''For the first time in several years,''',
    '''My name is Julien and I am''',
    '''The goal of life is''',
    '''Whenever I\'m sad, I like to''',
]


def   _snake_case    (    ):
      lowerCAmelCase       :						int							= argparse.ArgumentParser(description='''Run and time generations on a big model using Accelerate.'''      )
      parser.add_argument('''model_name'''			, type=_snake_case			, default=_snake_case			, help='''The name of the model to try.'''      )
      parser.add_argument(
          '''--tokenizer_name'''			, type=_snake_case			, default=_snake_case			, help='''The name of the tokenizer (if different from the model.'''      )
      parser.add_argument('''--is_causal'''			, type=_snake_case			, default=_snake_case			, help='''Whether or not the model is causal.'''      )
      parser.add_argument(
          '''--model_revision'''			, type=_snake_case			, default=_snake_case			, help='''The revision to use for the model checkpoint.'''      )
      parser.add_argument('''--torch_dtype'''			, type=_snake_case			, default=_snake_case			, help='''The dtype for the model.'''      )
      parser.add_argument('''--disk_offload'''			, action='''store_true'''      )

      lowerCAmelCase       :						Dict							= parser.parse_args()

      # Sanitize args
      if args.model_name in DEFAULT_MODELS:
            lowerCAmelCase       :						Optional[int]							= DEFAULT_MODELS[args.model_name]
            lowerCAmelCase       :						Union[str, Any]							= defaults['''model''']
            if args.tokenizer_name is None:
                  lowerCAmelCase       :						Union[str, Any]							= defaults.get('''tokenizer'''			, args.model_name      )
            if args.is_causal is None:
                  lowerCAmelCase       :						List[Any]							= defaults['''is_causal''']
            if args.model_revision is None:
                  lowerCAmelCase       :						Tuple							= defaults.get('''model_revision'''			, '''main'''      )

      if args.is_causal is None:
            raise ValueError('''Could not infer the default for `--is_causal`, pass either True or False for it.'''      )
      if args.tokenizer_name is None:
            lowerCAmelCase       :						Union[str, Any]							= args.model_name
      if args.model_revision is None:
            lowerCAmelCase       :						List[Any]							= '''main'''

      return args


def   _snake_case    (    ):
      transformers.utils.logging.set_verbosity_error()
      lowerCAmelCase       :						Union[str, Any]							= parse_args()

      if args.torch_dtype is None:
            lowerCAmelCase       :						List[str]							= AutoConfig.from_pretrained(args.model_name      )
            lowerCAmelCase       :						Optional[Any]							= getattr(_snake_case			, '''torch_dtype'''			, torch.floataa      )
      else:
            lowerCAmelCase       :						Optional[int]							= getattr(_snake_case			, args.torch_dtype      )
      lowerCAmelCase       :						Tuple							= AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
      lowerCAmelCase       :						str							= {
          '''torch_dtype''': torch_dtype,
          '''revision''': args.model_revision,
      }
      if args.disk_offload:
            lowerCAmelCase       :						Optional[int]							= '''tmp_offload'''
            lowerCAmelCase       :						Optional[int]							= True

      lowerCAmelCase       :						str							= start_measure()
      lowerCAmelCase       :						List[str]							= model_cls.from_pretrained(args.model_name			, device_map='''auto'''			, **_snake_case      )
      lowerCAmelCase       :						Tuple							= end_measure(_snake_case      )
      log_measures(_snake_case			, '''Model loading'''      )

      lowerCAmelCase       :						str							= compute_module_sizes(_snake_case      )
      lowerCAmelCase       :						Tuple							= {v: 0 for v in model.hf_device_map.values()}
      for module, device in model.hf_device_map.items():
            device_size[device] += module_sizes[module]
      lowerCAmelCase       :						int							= '''\n'''.join([f'''- {device}: {size // 2**20}MiB''' for device, size in device_size.items()]      )
      print(f'''\nTheoretical use:\n{message}'''      )

      lowerCAmelCase       :						int							= AutoTokenizer.from_pretrained(args.tokenizer_name      )

      lowerCAmelCase       :						List[str]							= start_measure()
      lowerCAmelCase       :						Union[str, Any]							= []
      lowerCAmelCase       :						int							= []
      lowerCAmelCase       :						Dict							= []
      for prompt in PROMPTS:
            lowerCAmelCase       :						str							= tokenizer(_snake_case			, return_tensors='''pt'''      ).to(0      )
            lowerCAmelCase       :						int							= inputs['''input_ids'''][0].tolist()
            lowerCAmelCase       :						Optional[Any]							= time.time()
            lowerCAmelCase       :						Optional[int]							= model.generate(inputs['''input_ids''']      )
            lowerCAmelCase       :						Any							= time.time()
            lowerCAmelCase       :						str							= outputs[0].tolist()
            lowerCAmelCase       :						int							= len(_snake_case      ) if outputs[: len(_snake_case      )] != tokens else len(_snake_case      ) - len(_snake_case      )
            lowerCAmelCase       :						Tuple							= after_generate - before_generate

            lowerCAmelCase       :						Tuple							= tokenizer.decode(_snake_case			, skip_special_tokens=_snake_case      )
            texts_outs.append(_snake_case      )
            generation_times.append(_snake_case      )
            gen_tokens.append(_snake_case      )
            print(f'''Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n'''      )

      lowerCAmelCase       :						List[Any]							= end_measure(_snake_case      )
      log_measures(_snake_case			, '''Model generation'''      )

      lowerCAmelCase       :						str							= [gen / tok for gen, tok in zip(_snake_case			, _snake_case      )]
      lowerCAmelCase       :						Any							= sum(_snake_case      ) / len(_snake_case      )
      print(f'''Average time of generation per token: {avg_gen:.2f}s'''      )
      print(f'''First generation (avg time per token): {generation_times_per_token[0]:.2f}s'''      )
      lowerCAmelCase       :						List[Any]							= sum(generation_times_per_token[1:]      ) / (len(_snake_case      ) - 1)
      print(f'''Average time of generation per token (excluding the first): {avg_gen:.2f}s'''      )


if __name__ == "__main__":
 main()






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


def   _snake_case    (    _snake_case					:      nn.Module      ):
      lowerCAmelCase       :						List[str]							= {}
      lowerCAmelCase       :						Optional[Any]							= (
          LoRAAttnProcessor2_0 if hasattr(_snake_case			, '''scaled_dot_product_attention'''      ) else LoRAAttnProcessor
      )
      for name, module in text_encoder.named_modules():
            if name.endswith(_snake_case      ):
                  lowerCAmelCase       :						Dict							= lora_attn_processor_class(
                      hidden_size=module.out_proj.out_features			, cross_attention_dim=_snake_case      )
      return text_lora_attn_procs


def   _snake_case    (    _snake_case					:      nn.Module      ):
      lowerCAmelCase       :						Tuple							= create_text_encoder_lora_attn_procs(_snake_case      )
      lowerCAmelCase       :						Optional[int]							= AttnProcsLayers(_snake_case      )
      return text_encoder_lora_layers


def   _snake_case    (    _snake_case					:      Optional[Any]			, _snake_case					:      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'''      ):
                        lowerCAmelCase       :						Dict							= (
                            torch.randn_like(layer_module.up.weight      )
                            if randn_weight
                            else torch.zeros_like(layer_module.up.weight      )
                        )
                        lowerCAmelCase       :						Optional[int]							= torch.nn.Parameter(_snake_case      )






class 		snake_case_(							unittest.TestCase      ):






     def   lowerCamelCase__  (						self	:       int					):
           torch.manual_seed(0					)
           lowerCAmelCase       :						str							= UNetaDConditionModel(
               block_out_channels=(3_2, 6_4)       ,					layers_per_block=2       ,					sample_size=3_2       ,					in_channels=4       ,					out_channels=4       ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,					up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,					cross_attention_dim=3_2       ,					)
           lowerCAmelCase       :						Dict							= DDIMScheduler(
               beta_start=0.00_085       ,					beta_end=0.012       ,					beta_schedule='''scaled_linear'''       ,					clip_sample=UpperCamelCase_       ,					set_alpha_to_one=UpperCamelCase_       ,					steps_offset=1       ,					)
           torch.manual_seed(0					)
           lowerCAmelCase       :						Tuple							= AutoencoderKL(
               block_out_channels=[3_2, 6_4]       ,					in_channels=3       ,					out_channels=3       ,					down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,					up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,					latent_channels=4       ,					)
           lowerCAmelCase       :						str							= CLIPTextConfig(
               bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=3_2       ,					intermediate_size=3_7       ,					layer_norm_eps=1E-05       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=1_0_0_0       ,					)
           lowerCAmelCase       :						List[Any]							= CLIPTextModel(UpperCamelCase_					)
           lowerCAmelCase       :						Dict							= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

           lowerCAmelCase,   lowerCAmelCase       :						List[str]							= create_unet_lora_layers(UpperCamelCase_					)
           lowerCAmelCase       :						int							= create_text_encoder_lora_layers(UpperCamelCase_					)

           lowerCAmelCase       :						Any							= {
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           lowerCAmelCase       :						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   lowerCamelCase__  (						self	:       List[str]					):
           lowerCAmelCase       :						str							= 1
           lowerCAmelCase       :						Union[str, Any]							= 1_0
           lowerCAmelCase       :						List[str]							= 4
           lowerCAmelCase       :						Dict							= (3_2, 3_2)

           lowerCAmelCase       :						Optional[Any]							= torch.manual_seed(0					)
           lowerCAmelCase       :						Any							= floats_tensor((batch_size, num_channels) + sizes					)
           lowerCAmelCase       :						Tuple							= torch.randint(1       ,					UpperCamelCase_       ,					size=(batch_size, sequence_length)       ,					generator=UpperCamelCase_					)

           lowerCAmelCase       :						Optional[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   lowerCamelCase__  (						self	:       Optional[Any]					):
           lowerCAmelCase       :						str							= 7_7

           lowerCAmelCase       :						str							= torch.randint(2       ,					5_6       ,					size=(1, max_seq_length)       ,					generator=torch.manual_seed(0					)					)

           lowerCAmelCase       :						Union[str, Any]							= {}
           lowerCAmelCase       :						List[str]							= inputs
           return prepared_inputs






     def   lowerCamelCase__  (						self	:       int       ,					UpperCamelCase_	:       Optional[Any]					):
           lowerCAmelCase,   lowerCAmelCase       :						Tuple							= self.get_dummy_components()
           LoraLoaderMixin.save_lora_weights(
               save_directory=UpperCamelCase_       ,					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(UpperCamelCase_       ,					'''pytorch_lora_weights.bin'''					)					)					)






     def   lowerCamelCase__  (						self	:       Optional[int]					):
           lowerCAmelCase,   lowerCAmelCase       :						Union[str, Any]							= self.get_dummy_components()
           lowerCAmelCase       :						str							= StableDiffusionPipeline(**UpperCamelCase_					)
           lowerCAmelCase       :						Union[str, Any]							= sd_pipe.to(UpperCamelCase_					)
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase_					)

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

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

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=UpperCamelCase_       ,					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(UpperCamelCase_       ,					'''pytorch_lora_weights.bin'''					)					)					)
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

           lowerCAmelCase       :						Optional[Any]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Tuple							= lora_images[0, -3:, -3:, -1]

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






     def   lowerCamelCase__  (						self	:       List[str]					):
           lowerCAmelCase,   lowerCAmelCase       :						int							= self.get_dummy_components()
           lowerCAmelCase       :						Dict							= StableDiffusionPipeline(**UpperCamelCase_					)
           lowerCAmelCase       :						List[str]							= sd_pipe.to(UpperCamelCase_					)
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase_					)

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

           lowerCAmelCase       :						Union[str, Any]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Union[str, Any]							= original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=UpperCamelCase_       ,					unet_lora_layers=lora_components['''unet_lora_layers''']       ,					text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']       ,					safe_serialization=UpperCamelCase_       ,					)
                 self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_       ,					'''pytorch_lora_weights.safetensors'''					)					)					)
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

           lowerCAmelCase       :						Union[str, Any]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						List[str]							= lora_images[0, -3:, -3:, -1]

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






     def   lowerCamelCase__  (						self	:       Union[str, Any]					):
           lowerCAmelCase,   lowerCAmelCase       :						Tuple							= self.get_dummy_components()
           lowerCAmelCase       :						Dict							= lora_components['''unet_lora_attn_procs''']
           lowerCAmelCase       :						str							= StableDiffusionPipeline(**UpperCamelCase_					)
           lowerCAmelCase       :						str							= sd_pipe.to(UpperCamelCase_					)
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase_					)

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

           lowerCAmelCase       :						List[str]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Optional[Any]							= original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 lowerCAmelCase       :						Tuple							= sd_pipe.unet
                 unet.set_attn_processor(UpperCamelCase_					)
                 unet.save_attn_procs(UpperCamelCase_					)
                 self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase_       ,					'''pytorch_lora_weights.bin'''					)					)					)
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

           lowerCAmelCase       :						List[Any]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Optional[int]							= lora_images[0, -3:, -3:, -1]

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






     def   lowerCamelCase__  (						self	:       List[Any]					):
           lowerCAmelCase,   lowerCAmelCase       :						str							= self.get_dummy_components()
           lowerCAmelCase       :						Optional[int]							= StableDiffusionPipeline(**UpperCamelCase_					)

           lowerCAmelCase       :						Union[str, Any]							= self.get_dummy_tokens()

           # inference without lora
           lowerCAmelCase       :						Optional[int]							= pipe.text_encoder(**UpperCamelCase_					)[0]
           assert outputs_without_lora.shape == (1, 7_7, 3_2)

           # create lora_attn_procs with zeroed out up.weights
           lowerCAmelCase       :						Tuple							= create_text_encoder_lora_attn_procs(pipe.text_encoder					)
           set_lora_up_weights(UpperCamelCase_       ,					randn_weight=UpperCamelCase_					)

           # monkey patch
           pipe._modify_text_encoder(UpperCamelCase_					)

           # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
           del text_attn_procs
           gc.collect()

           # inference with lora
           lowerCAmelCase       :						Any							= pipe.text_encoder(**UpperCamelCase_					)[0]
           assert outputs_with_lora.shape == (1, 7_7, 3_2)

           assert torch.allclose(
               UpperCamelCase_       ,					UpperCamelCase_					), "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
           lowerCAmelCase       :						Any							= create_text_encoder_lora_attn_procs(pipe.text_encoder					)
           set_lora_up_weights(UpperCamelCase_       ,					randn_weight=UpperCamelCase_					)

           # monkey patch
           pipe._modify_text_encoder(UpperCamelCase_					)

           # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
           del text_attn_procs
           gc.collect()

           # inference with lora
           lowerCAmelCase       :						Tuple							= pipe.text_encoder(**UpperCamelCase_					)[0]
           assert outputs_with_lora.shape == (1, 7_7, 3_2)

           assert not torch.allclose(
               UpperCamelCase_       ,					UpperCamelCase_					), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"






     def   lowerCamelCase__  (						self	:       Optional[Any]					):
           lowerCAmelCase,   lowerCAmelCase       :						Optional[Any]							= self.get_dummy_components()
           lowerCAmelCase       :						Optional[int]							= StableDiffusionPipeline(**UpperCamelCase_					)

           lowerCAmelCase       :						Dict							= self.get_dummy_tokens()

           # inference without lora
           lowerCAmelCase       :						List[str]							= pipe.text_encoder(**UpperCamelCase_					)[0]
           assert outputs_without_lora.shape == (1, 7_7, 3_2)

           # create lora_attn_procs with randn up.weights
           lowerCAmelCase       :						Any							= create_text_encoder_lora_attn_procs(pipe.text_encoder					)
           set_lora_up_weights(UpperCamelCase_       ,					randn_weight=UpperCamelCase_					)

           # monkey patch
           pipe._modify_text_encoder(UpperCamelCase_					)

           # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
           del text_attn_procs
           gc.collect()

           # inference with lora
           lowerCAmelCase       :						Union[str, Any]							= pipe.text_encoder(**UpperCamelCase_					)[0]
           assert outputs_with_lora.shape == (1, 7_7, 3_2)

           assert not torch.allclose(
               UpperCamelCase_       ,					UpperCamelCase_					), "lora outputs should be different to without lora outputs"

           # remove monkey patch
           pipe._remove_text_encoder_monkey_patch()

           # inference with removed lora
           lowerCAmelCase       :						List[str]							= pipe.text_encoder(**UpperCamelCase_					)[0]
           assert outputs_without_lora_removed.shape == (1, 7_7, 3_2)

           assert torch.allclose(
               UpperCamelCase_       ,					UpperCamelCase_					), "remove lora monkey patch should restore the original outputs"






     def   lowerCamelCase__  (						self	:       str					):
           lowerCAmelCase,   lowerCAmelCase       :						List[Any]							= self.get_dummy_components()
           lowerCAmelCase       :						Optional[Any]							= StableDiffusionPipeline(**UpperCamelCase_					)
           lowerCAmelCase       :						Optional[Any]							= sd_pipe.to(UpperCamelCase_					)
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase_					)

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

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=UpperCamelCase_       ,					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(UpperCamelCase_       ,					'''pytorch_lora_weights.bin'''					)					)					)
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

           lowerCAmelCase       :						List[str]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Union[str, Any]							= lora_images[0, -3:, -3:, -1]

           lowerCAmelCase       :						List[str]							= sd_pipe(**UpperCamelCase_       ,					cross_attention_kwargs={'''scale''': 0.5}					).images
           lowerCAmelCase       :						List[str]							= lora_images_with_scale[0, -3:, -3:, -1]

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






     def   lowerCamelCase__  (						self	:       List[str]					):
           with tempfile.TemporaryDirectory() as tmpdirname:
                 self.create_lora_weight_file(UpperCamelCase_					)

                 lowerCAmelCase,   lowerCAmelCase       :						List[str]							= self.get_dummy_components()
                 lowerCAmelCase       :						List[str]							= StableDiffusionPipeline(**UpperCamelCase_					)
                 lowerCAmelCase       :						Dict							= sd_pipe.to(UpperCamelCase_					)
                 sd_pipe.set_progress_bar_config(disable=UpperCamelCase_					)

                 # check if vanilla attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(UpperCamelCase_       ,					UpperCamelCase_					):
                             self.assertIsInstance(module.processor       ,					(AttnProcessor, AttnProcessor2_0)					)

            # load LoRA weight file
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

                 # check if lora attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(UpperCamelCase_       ,					UpperCamelCase_					):
                             lowerCAmelCase       :						Any							= (
                                 LoRAAttnProcessor2_0 if hasattr(UpperCamelCase_       ,					'''scaled_dot_product_attention'''					) else LoRAAttnProcessor
                             )
                             self.assertIsInstance(module.processor       ,					UpperCamelCase_					)






     @unittest.skipIf(torch_device != '''cuda'''       ,					'''This test is supposed to run on GPU'''					)
     def   lowerCamelCase__  (						self	:       str					):
           with tempfile.TemporaryDirectory() as tmpdirname:
                 self.create_lora_weight_file(UpperCamelCase_					)

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

                 # 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(UpperCamelCase_       ,					UpperCamelCase_					):
                             self.assertIsInstance(module.processor       ,					UpperCamelCase_					)

            # load LoRA weight file
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

                 # check if lora attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(UpperCamelCase_       ,					UpperCamelCase_					):
                             self.assertIsInstance(module.processor       ,					UpperCamelCase_					)






     @unittest.skipIf(torch_device != '''cuda'''       ,					'''This test is supposed to run on GPU'''					)
     def   lowerCamelCase__  (						self	:       int					):
           lowerCAmelCase,   lowerCAmelCase       :						Union[str, Any]							= self.get_dummy_components()
           lowerCAmelCase       :						Tuple							= StableDiffusionPipeline(**UpperCamelCase_					)
           lowerCAmelCase       :						Optional[Any]							= sd_pipe.to(UpperCamelCase_					)
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase_					)

           lowerCAmelCase,   lowerCAmelCase,   lowerCAmelCase       :						Union[str, Any]							= self.get_dummy_inputs()

           # enable XFormers
           sd_pipe.enable_xformers_memory_efficient_attention()

           lowerCAmelCase       :						List[Any]							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Optional[int]							= original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=UpperCamelCase_       ,					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(UpperCamelCase_       ,					'''pytorch_lora_weights.bin'''					)					)					)
                 sd_pipe.load_lora_weights(UpperCamelCase_					)

           lowerCAmelCase       :						Any							= sd_pipe(**UpperCamelCase_					).images
           lowerCAmelCase       :						Optional[Any]							= lora_images[0, -3:, -3:, -1]

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


