






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				(lowercase__				)		->				Tuple:

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



def 			__UpperCAmelCase				(lowercase__				)		->				List[Any]:

						'''simple docstring'''


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



def 			__UpperCAmelCase				(lowercase__				)		->				Optional[int]:

						'''simple docstring'''


						a_				     =   create_text_encoder_lora_attn_procs(lowercase__				)
						a_				     =   AttnProcsLayers(lowercase__				)
						return text_encoder_lora_layers



def 			__UpperCAmelCase				(lowercase__	,lowercase__=False				)		->				Optional[Any]:

						'''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(lowercase__				)





class 		SCREAMING_SNAKE_CASE__   (     unittest.TestCase    ):



		def _lowerCAmelCase     (    self:    Optional[int])  ->List[str]:






								'''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.0_0085				,   beta_end=0.012				,   beta_schedule="scaled_linear"				,   clip_sample=UpperCamelCase__				,   set_alpha_to_one=UpperCamelCase__				,   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=10_00				,   )
								a_				     =   CLIPTextModel(UpperCamelCase__)
								a_				     =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

								a_     ,       a_				     =   create_unet_lora_layers(UpperCamelCase__)
								a_				     =   create_text_encoder_lora_layers(UpperCamelCase__)

								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 _lowerCAmelCase     (    self:    Dict)  ->List[Any]:






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

								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 _lowerCAmelCase     (    self:    List[str])  ->Tuple:






								'''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 _lowerCAmelCase     (    self:    Any				,   a:    List[Any])  ->Tuple:






								'''simple docstring'''


								a_     ,       a_				     =   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:    str)  ->Dict:






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)
								a_				     =   sd_pipe.to(UpperCamelCase__)
								sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

								a_     ,       a_     ,       a_				     =   self.get_dummy_inputs()

								a_				     =   sd_pipe(**UpperCamelCase__).images
								a_				     =   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__)

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






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)
								a_				     =   sd_pipe.to(UpperCamelCase__)
								sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

								a_     ,       a_     ,       a_				     =   self.get_dummy_inputs()

								a_				     =   sd_pipe(**UpperCamelCase__).images
								a_				     =   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__)

								a_				     =   sd_pipe(**UpperCamelCase__).images
								a_				     =   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:    int)  ->Any:






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   lora_components["unet_lora_attn_procs"]
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)
								a_				     =   sd_pipe.to(UpperCamelCase__)
								sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

								a_     ,       a_     ,       a_				     =   self.get_dummy_inputs()

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

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

								a_				     =   sd_pipe(**UpperCamelCase__).images
								a_				     =   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:    int)  ->str:






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)

								a_				     =   self.get_dummy_tokens()

								# inference without lora
								a_				     =   pipe.text_encoder(**UpperCamelCase__)[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(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
								a_				     =   pipe.text_encoder(**UpperCamelCase__)[0]
								assert outputs_with_lora.shape == (1, 77, 32)

								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
								a_				     =   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
								a_				     =   pipe.text_encoder(**UpperCamelCase__)[0]
								assert outputs_with_lora.shape == (1, 77, 32)

								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:    Any)  ->Tuple:






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)

								a_				     =   self.get_dummy_tokens()

								# inference without lora
								a_				     =   pipe.text_encoder(**UpperCamelCase__)[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(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
								a_				     =   pipe.text_encoder(**UpperCamelCase__)[0]
								assert outputs_with_lora.shape == (1, 77, 32)

								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
								a_				     =   pipe.text_encoder(**UpperCamelCase__)[0]
								assert outputs_without_lora_removed.shape == (1, 77, 32)

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



		def _lowerCAmelCase     (    self:    List[str])  ->Tuple:






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)
								a_				     =   sd_pipe.to(UpperCamelCase__)
								sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

								a_     ,       a_     ,       a_				     =   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__)

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

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



		def _lowerCAmelCase     (    self:    List[Any])  ->Optional[int]:






								'''simple docstring'''


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

														a_     ,       a_				     =   self.get_dummy_components()
														a_				     =   StableDiffusionPipeline(**UpperCamelCase__)
														a_				     =   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__):
																										a_				     =   (
																										    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:    int)  ->Any:






								'''simple docstring'''


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

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






								'''simple docstring'''


								a_     ,       a_				     =   self.get_dummy_components()
								a_				     =   StableDiffusionPipeline(**UpperCamelCase__)
								a_				     =   sd_pipe.to(UpperCamelCase__)
								sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

								a_     ,       a_     ,       a_				     =   self.get_dummy_inputs()

								# enable XFormers
								sd_pipe.enable_xformers_memory_efficient_attention()

								a_				     =   sd_pipe(**UpperCamelCase__).images
								a_				     =   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__)

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

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











'''simple docstring'''





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


a_      = logging.get_logger(__name__)

a_      = {
    'google/mobilebert-uncased': 'https://huggingface.co/google/mobilebert-uncased/resolve/main/config.json'
}





class 		SCREAMING_SNAKE_CASE__   (     lowercase_    ):
		_UpperCAmelCase         =MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
		_UpperCAmelCase         ='''mobilebert'''



		def __init__(    self:    Union[str, Any]				,   a:    str=3_05_22				,   a:    List[str]=5_12				,   a:    List[str]=24				,   a:    Dict=4				,   a:    List[str]=5_12				,   a:    Optional[int]="relu"				,   a:    Dict=0.0				,   a:    Union[str, Any]=0.1				,   a:    int=5_12				,   a:    Optional[Any]=2				,   a:    List[Any]=0.02				,   a:    Optional[Any]=1e-12				,   a:    List[Any]=0				,   a:    Optional[int]=1_28				,   a:    Union[str, Any]=True				,   a:    Dict=True				,   a:    List[Any]=1_28				,   a:    Any=False				,   a:    Any=True				,   a:    Dict=4				,   a:    List[Any]="no_norm"				,   a:    Any=True				,   a:    Tuple=None				,   **a:    Any				,   )  ->Any:






								'''simple docstring'''


								super().__init__(pad_token_id=a				,   **a)

								a_				     =   vocab_size
								a_				     =   hidden_size
								a_				     =   num_hidden_layers
								a_				     =   num_attention_heads
								a_				     =   hidden_act
								a_				     =   intermediate_size
								a_				     =   hidden_dropout_prob
								a_				     =   attention_probs_dropout_prob
								a_				     =   max_position_embeddings
								a_				     =   type_vocab_size
								a_				     =   initializer_range
								a_				     =   layer_norm_eps
								a_				     =   embedding_size
								a_				     =   trigram_input
								a_				     =   use_bottleneck
								a_				     =   intra_bottleneck_size
								a_				     =   use_bottleneck_attention
								a_				     =   key_query_shared_bottleneck
								a_				     =   num_feedforward_networks
								a_				     =   normalization_type
								a_				     =   classifier_activation

								if self.use_bottleneck:
														a_				     =   intra_bottleneck_size
								else:
														a_				     =   hidden_size

								a_				     =   classifier_dropout





class 		SCREAMING_SNAKE_CASE__   (     lowercase_    ):



		@property
		def _lowerCAmelCase     (    self:    Optional[int])  ->Mapping[str, Mapping[int, str]]:






								'''simple docstring'''


								if self.task == "multiple-choice":
														a_				     =   {0: "batch", 1: "choice", 2: "sequence"}
								else:
														a_				     =   {0: "batch", 1: "sequence"}
								return OrderedDict(
								    [
								        ("input_ids", dynamic_axis),
								        ("attention_mask", dynamic_axis),
								        ("token_type_ids", dynamic_axis),
								    ])



