



'''simple docstring'''







from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


SCREAMING_SNAKE_CASE__     : Any				=   {
    '''configuration_utils''': ['''GenerationConfig'''],
    '''streamers''': ['''TextIteratorStreamer''', '''TextStreamer'''],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__     : Dict				=   [
							    '''Constraint''',
							    '''ConstraintListState''',
							    '''DisjunctiveConstraint''',
							    '''PhrasalConstraint''',
							]
							SCREAMING_SNAKE_CASE__     : Optional[int]				=   [
							    '''BeamHypotheses''',
							    '''BeamScorer''',
							    '''BeamSearchScorer''',
							    '''ConstrainedBeamSearchScorer''',
							]
							SCREAMING_SNAKE_CASE__     : str				=   [
							    '''EpsilonLogitsWarper''',
							    '''EtaLogitsWarper''',
							    '''ForcedBOSTokenLogitsProcessor''',
							    '''ForcedEOSTokenLogitsProcessor''',
							    '''HammingDiversityLogitsProcessor''',
							    '''InfNanRemoveLogitsProcessor''',
							    '''LogitsProcessor''',
							    '''LogitsProcessorList''',
							    '''LogitsWarper''',
							    '''MinLengthLogitsProcessor''',
							    '''MinNewTokensLengthLogitsProcessor''',
							    '''NoBadWordsLogitsProcessor''',
							    '''NoRepeatNGramLogitsProcessor''',
							    '''PrefixConstrainedLogitsProcessor''',
							    '''RepetitionPenaltyLogitsProcessor''',
							    '''SequenceBiasLogitsProcessor''',
							    '''EncoderRepetitionPenaltyLogitsProcessor''',
							    '''TemperatureLogitsWarper''',
							    '''TopKLogitsWarper''',
							    '''TopPLogitsWarper''',
							    '''TypicalLogitsWarper''',
							    '''EncoderNoRepeatNGramLogitsProcessor''',
							    '''ExponentialDecayLengthPenalty''',
							    '''LogitNormalization''',
							]
							SCREAMING_SNAKE_CASE__     : Optional[int]				=   [
							    '''MaxNewTokensCriteria''',
							    '''MaxLengthCriteria''',
							    '''MaxTimeCriteria''',
							    '''StoppingCriteria''',
							    '''StoppingCriteriaList''',
							    '''validate_stopping_criteria''',
							]
							SCREAMING_SNAKE_CASE__     : Any				=   [
							    '''GenerationMixin''',
							    '''top_k_top_p_filtering''',
							    '''GreedySearchEncoderDecoderOutput''',
							    '''GreedySearchDecoderOnlyOutput''',
							    '''SampleEncoderDecoderOutput''',
							    '''SampleDecoderOnlyOutput''',
							    '''BeamSearchEncoderDecoderOutput''',
							    '''BeamSearchDecoderOnlyOutput''',
							    '''BeamSampleEncoderDecoderOutput''',
							    '''BeamSampleDecoderOnlyOutput''',
							    '''ContrastiveSearchEncoderDecoderOutput''',
							    '''ContrastiveSearchDecoderOnlyOutput''',
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__     : Tuple				=   [
							    '''TFForcedBOSTokenLogitsProcessor''',
							    '''TFForcedEOSTokenLogitsProcessor''',
							    '''TFLogitsProcessor''',
							    '''TFLogitsProcessorList''',
							    '''TFLogitsWarper''',
							    '''TFMinLengthLogitsProcessor''',
							    '''TFNoBadWordsLogitsProcessor''',
							    '''TFNoRepeatNGramLogitsProcessor''',
							    '''TFRepetitionPenaltyLogitsProcessor''',
							    '''TFTemperatureLogitsWarper''',
							    '''TFTopKLogitsWarper''',
							    '''TFTopPLogitsWarper''',
							    '''TFForceTokensLogitsProcessor''',
							    '''TFSuppressTokensAtBeginLogitsProcessor''',
							    '''TFSuppressTokensLogitsProcessor''',
							]
							SCREAMING_SNAKE_CASE__     : Tuple				=   [
							    '''TFGenerationMixin''',
							    '''tf_top_k_top_p_filtering''',
							    '''TFGreedySearchDecoderOnlyOutput''',
							    '''TFGreedySearchEncoderDecoderOutput''',
							    '''TFSampleEncoderDecoderOutput''',
							    '''TFSampleDecoderOnlyOutput''',
							    '''TFBeamSearchEncoderDecoderOutput''',
							    '''TFBeamSearchDecoderOnlyOutput''',
							    '''TFBeamSampleEncoderDecoderOutput''',
							    '''TFBeamSampleDecoderOnlyOutput''',
							    '''TFContrastiveSearchEncoderDecoderOutput''',
							    '''TFContrastiveSearchDecoderOnlyOutput''',
							]

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__     : List[str]				=   [
							    '''FlaxForcedBOSTokenLogitsProcessor''',
							    '''FlaxForcedEOSTokenLogitsProcessor''',
							    '''FlaxLogitsProcessor''',
							    '''FlaxLogitsProcessorList''',
							    '''FlaxLogitsWarper''',
							    '''FlaxMinLengthLogitsProcessor''',
							    '''FlaxTemperatureLogitsWarper''',
							    '''FlaxTopKLogitsWarper''',
							    '''FlaxTopPLogitsWarper''',
							]
							SCREAMING_SNAKE_CASE__     : int				=   [
							    '''FlaxGenerationMixin''',
							    '''FlaxGreedySearchOutput''',
							    '''FlaxSampleOutput''',
							    '''FlaxBeamSearchOutput''',
							]

if TYPE_CHECKING:
							from .configuration_utils import GenerationConfig
							from .streamers import TextIteratorStreamer, TextStreamer

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
														from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
														from .logits_process import (
														    EncoderNoRepeatNGramLogitsProcessor,
														    EncoderRepetitionPenaltyLogitsProcessor,
														    EpsilonLogitsWarper,
														    EtaLogitsWarper,
														    ExponentialDecayLengthPenalty,
														    ForcedBOSTokenLogitsProcessor,
														    ForcedEOSTokenLogitsProcessor,
														    HammingDiversityLogitsProcessor,
														    InfNanRemoveLogitsProcessor,
														    LogitNormalization,
														    LogitsProcessor,
														    LogitsProcessorList,
														    LogitsWarper,
														    MinLengthLogitsProcessor,
														    MinNewTokensLengthLogitsProcessor,
														    NoBadWordsLogitsProcessor,
														    NoRepeatNGramLogitsProcessor,
														    PrefixConstrainedLogitsProcessor,
														    RepetitionPenaltyLogitsProcessor,
														    SequenceBiasLogitsProcessor,
														    TemperatureLogitsWarper,
														    TopKLogitsWarper,
														    TopPLogitsWarper,
														    TypicalLogitsWarper,
														)
														from .stopping_criteria import (
														    MaxLengthCriteria,
														    MaxNewTokensCriteria,
														    MaxTimeCriteria,
														    StoppingCriteria,
														    StoppingCriteriaList,
														    validate_stopping_criteria,
														)
														from .utils import (
														    BeamSampleDecoderOnlyOutput,
														    BeamSampleEncoderDecoderOutput,
														    BeamSearchDecoderOnlyOutput,
														    BeamSearchEncoderDecoderOutput,
														    ContrastiveSearchDecoderOnlyOutput,
														    ContrastiveSearchEncoderDecoderOutput,
														    GenerationMixin,
														    GreedySearchDecoderOnlyOutput,
														    GreedySearchEncoderDecoderOutput,
														    SampleDecoderOnlyOutput,
														    SampleEncoderDecoderOutput,
														    top_k_top_p_filtering,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tf_logits_process import (
														    TFForcedBOSTokenLogitsProcessor,
														    TFForcedEOSTokenLogitsProcessor,
														    TFForceTokensLogitsProcessor,
														    TFLogitsProcessor,
														    TFLogitsProcessorList,
														    TFLogitsWarper,
														    TFMinLengthLogitsProcessor,
														    TFNoBadWordsLogitsProcessor,
														    TFNoRepeatNGramLogitsProcessor,
														    TFRepetitionPenaltyLogitsProcessor,
														    TFSuppressTokensAtBeginLogitsProcessor,
														    TFSuppressTokensLogitsProcessor,
														    TFTemperatureLogitsWarper,
														    TFTopKLogitsWarper,
														    TFTopPLogitsWarper,
														)
														from .tf_utils import (
														    TFBeamSampleDecoderOnlyOutput,
														    TFBeamSampleEncoderDecoderOutput,
														    TFBeamSearchDecoderOnlyOutput,
														    TFBeamSearchEncoderDecoderOutput,
														    TFContrastiveSearchDecoderOnlyOutput,
														    TFContrastiveSearchEncoderDecoderOutput,
														    TFGenerationMixin,
														    TFGreedySearchDecoderOnlyOutput,
														    TFGreedySearchEncoderDecoderOutput,
														    TFSampleDecoderOnlyOutput,
														    TFSampleEncoderDecoderOutput,
														    tf_top_k_top_p_filtering,
														)

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .flax_logits_process import (
														    FlaxForcedBOSTokenLogitsProcessor,
														    FlaxForcedEOSTokenLogitsProcessor,
														    FlaxLogitsProcessor,
														    FlaxLogitsProcessorList,
														    FlaxLogitsWarper,
														    FlaxMinLengthLogitsProcessor,
														    FlaxTemperatureLogitsWarper,
														    FlaxTopKLogitsWarper,
														    FlaxTopPLogitsWarper,
														)
														from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
							import sys

							SCREAMING_SNAKE_CASE__     : str				=   _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




"""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_					)					)					)


