




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



def 					UpperCamelCase__       (						lowercase__					:       nn.Module				):
		snake_case						:  Any        =			{}
		snake_case						:  List[str]        =			(
		    LoRAAttnProcessor2_0 if hasattr(_lowerCAmelCase    ,       "scaled_dot_product_attention"				) else LoRAAttnProcessor
		)
		for name, module in text_encoder.named_modules():
				if name.endswith(_lowerCAmelCase				):
						snake_case						:  List[Any]        =			lora_attn_processor_class(
						    hidden_size=module.out_proj.out_features    ,       cross_attention_dim=_lowerCAmelCase				)
		return text_lora_attn_procs



def 					UpperCamelCase__       (						lowercase__					:       nn.Module				):
		snake_case						:  Dict        =			create_text_encoder_lora_attn_procs(_lowerCAmelCase				)
		snake_case						:  Union[str, Any]        =			AttnProcsLayers(_lowerCAmelCase				)
		return text_encoder_lora_layers



def 					UpperCamelCase__       (						lowercase__					:       List[str]    ,       lowercase__					:       Optional[int]=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"				):
								snake_case						:  List[Any]        =			(
								    torch.randn_like(layer_module.up.weight				)
								    if randn_weight
								    else torch.zeros_like(layer_module.up.weight				)
								)
								snake_case						:  Optional[Any]        =			torch.nn.Parameter(_lowerCAmelCase				)
class      lowerCamelCase__      (					unittest.TestCase		):


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					torch.manual_seed(0							)
					snake_case						:  Optional[int]        =			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			,	)
					snake_case						:  Dict        =			DDIMScheduler(
					    beta_start=0.0_00_85			,	beta_end=0.0_12			,	beta_schedule="scaled_linear"			,	clip_sample=__snake_case			,	set_alpha_to_one=__snake_case			,	steps_offset=1			,	)
					torch.manual_seed(0							)
					snake_case						:  str        =			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			,	)
					snake_case						:  Union[str, Any]        =			CLIPTextConfig(
					    bos_token_id=0			,	eos_token_id=2			,	hidden_size=32			,	intermediate_size=37			,	layer_norm_eps=1E-05			,	num_attention_heads=4			,	num_hidden_layers=5			,	pad_token_id=1			,	vocab_size=1_000			,	)
					snake_case						:  Union[str, Any]        =			CLIPTextModel(__snake_case							)
					snake_case						:  List[str]        =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)

					snake_case						:  Tuple        =			create_unet_lora_layers(__snake_case							)
					snake_case						:  Union[str, Any]        =			create_text_encoder_lora_layers(__snake_case							)

					snake_case						:  Union[str, Any]        =			{
					    '''unet''': unet,
					    '''scheduler''': scheduler,
					    '''vae''': vae,
					    '''text_encoder''': text_encoder,
					    '''tokenizer''': tokenizer,
					    '''safety_checker''': None,
					    '''feature_extractor''': None,
					}
					snake_case						:  Dict        =			{
					    '''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							):



					"""simple docstring"""







					snake_case						:  Optional[int]        =			1
					snake_case						:  str        =			10
					snake_case						:  int        =			4
					snake_case						:  Tuple        =			(32, 32)

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

					snake_case						:  int        =			{
					    '''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							):



					"""simple docstring"""







					snake_case						:  Tuple        =			77

					snake_case						:  Any        =			torch.randint(2			,	56			,	size=(1, max_seq_length)			,	generator=torch.manual_seed(0							)							)

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


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Union[str, Any]        =			self.get_dummy_components()
					LoraLoaderMixin.save_lora_weights(
					    save_directory=__snake_case			,	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(__snake_case			,	"pytorch_lora_weights.bin"							)							)							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  List[str]        =			self.get_dummy_components()
					snake_case						:  Any        =			StableDiffusionPipeline(**__snake_case							)
					snake_case						:  Dict        =			sd_pipe.to(__snake_case							)
					sd_pipe.set_progress_bar_config(disable=__snake_case							)

					snake_case						:  int        =			self.get_dummy_inputs()

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

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

					snake_case						:  Optional[int]        =			sd_pipe(**__snake_case							).images
					snake_case						:  str        =			lora_images[0, -3:, -3:, -1]

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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  int        =			self.get_dummy_components()
					snake_case						:  Optional[Any]        =			StableDiffusionPipeline(**__snake_case							)
					snake_case						:  Any        =			sd_pipe.to(__snake_case							)
					sd_pipe.set_progress_bar_config(disable=__snake_case							)

					snake_case						:  Optional[int]        =			self.get_dummy_inputs()

					snake_case						:  str        =			sd_pipe(**__snake_case							).images
					snake_case						:  Any        =			original_images[0, -3:, -3:, -1]

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

					snake_case						:  Any        =			sd_pipe(**__snake_case							).images
					snake_case						:  Optional[int]        =			lora_images[0, -3:, -3:, -1]

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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Tuple        =			self.get_dummy_components()
					snake_case						:  Optional[int]        =			lora_components['''unet_lora_attn_procs''']
					snake_case						:  Any        =			StableDiffusionPipeline(**__snake_case							)
					snake_case						:  List[str]        =			sd_pipe.to(__snake_case							)
					sd_pipe.set_progress_bar_config(disable=__snake_case							)

					snake_case						:  int        =			self.get_dummy_inputs()

					snake_case						:  Tuple        =			sd_pipe(**__snake_case							).images
					snake_case						:  Optional[int]        =			original_images[0, -3:, -3:, -1]

					with tempfile.TemporaryDirectory() as tmpdirname:
							snake_case						:  List[Any]        =			sd_pipe.unet
							unet.set_attn_processor(__snake_case							)
							unet.save_attn_procs(__snake_case							)
							self.assertTrue(os.path.isfile(os.path.join(__snake_case			,	"pytorch_lora_weights.bin"							)							)							)
							sd_pipe.load_lora_weights(__snake_case							)

					snake_case						:  Dict        =			sd_pipe(**__snake_case							).images
					snake_case						:  Tuple        =			lora_images[0, -3:, -3:, -1]

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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Tuple        =			self.get_dummy_components()
					snake_case						:  Any        =			StableDiffusionPipeline(**__snake_case							)

					snake_case						:  Optional[Any]        =			self.get_dummy_tokens()

					# inference without lora
					snake_case						:  Optional[Any]        =			pipe.text_encoder(**__snake_case							)[0]
					assert outputs_without_lora.shape == (1, 77, 32)

					# create lora_attn_procs with zeroed out up.weights
					snake_case						:  str        =			create_text_encoder_lora_attn_procs(pipe.text_encoder							)
					set_lora_up_weights(__snake_case			,	randn_weight=__snake_case							)

					# monkey patch
					pipe._modify_text_encoder(__snake_case							)

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

					# inference with lora
					snake_case						:  Any        =			pipe.text_encoder(**__snake_case							)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

					assert torch.allclose(
					    __snake_case			,	__snake_case							), "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
					snake_case						:  List[Any]        =			create_text_encoder_lora_attn_procs(pipe.text_encoder							)
					set_lora_up_weights(__snake_case			,	randn_weight=__snake_case							)

					# monkey patch
					pipe._modify_text_encoder(__snake_case							)

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

					# inference with lora
					snake_case						:  Optional[int]        =			pipe.text_encoder(**__snake_case							)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  Tuple        =			self.get_dummy_components()
					snake_case						:  Optional[int]        =			StableDiffusionPipeline(**__snake_case							)

					snake_case						:  Tuple        =			self.get_dummy_tokens()

					# inference without lora
					snake_case						:  str        =			pipe.text_encoder(**__snake_case							)[0]
					assert outputs_without_lora.shape == (1, 77, 32)

					# create lora_attn_procs with randn up.weights
					snake_case						:  str        =			create_text_encoder_lora_attn_procs(pipe.text_encoder							)
					set_lora_up_weights(__snake_case			,	randn_weight=__snake_case							)

					# monkey patch
					pipe._modify_text_encoder(__snake_case							)

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

					# inference with lora
					snake_case						:  int        =			pipe.text_encoder(**__snake_case							)[0]
					assert outputs_with_lora.shape == (1, 77, 32)

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

					# remove monkey patch
					pipe._remove_text_encoder_monkey_patch()

					# inference with removed lora
					snake_case						:  Any        =			pipe.text_encoder(**__snake_case							)[0]
					assert outputs_without_lora_removed.shape == (1, 77, 32)

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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  int        =			self.get_dummy_components()
					snake_case						:  Tuple        =			StableDiffusionPipeline(**__snake_case							)
					snake_case						:  Union[str, Any]        =			sd_pipe.to(__snake_case							)
					sd_pipe.set_progress_bar_config(disable=__snake_case							)

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

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

					snake_case						:  str        =			sd_pipe(**__snake_case							).images
					snake_case						:  Union[str, Any]        =			lora_images[0, -3:, -3:, -1]

					snake_case						:  Dict        =			sd_pipe(**__snake_case			,	cross_attention_kwargs={"scale": 0.5}							).images
					snake_case						:  Dict        =			lora_images_with_scale[0, -3:, -3:, -1]

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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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

							snake_case						:  List[str]        =			self.get_dummy_components()
							snake_case						:  Tuple        =			StableDiffusionPipeline(**__snake_case							)
							snake_case						:  List[str]        =			sd_pipe.to(__snake_case							)
							sd_pipe.set_progress_bar_config(disable=__snake_case							)

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

            # load LoRA weight file
							sd_pipe.load_lora_weights(__snake_case							)

							# check if lora attention processors are used
							for _, module in sd_pipe.unet.named_modules():
									if isinstance(__snake_case			,	__snake_case							):
											snake_case						:  List[Any]        =			(
											    LoRAAttnProcessor2_0 if hasattr(__snake_case			,	"scaled_dot_product_attention"							) else LoRAAttnProcessor
											)
											self.assertIsInstance(module.processor			,	__snake_case							)


			@unittest.skipIf(torch_device != "cuda"			,	"This test is supposed to run on GPU"							)
			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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

							snake_case						:  Tuple        =			self.get_dummy_components()
							snake_case						:  List[str]        =			StableDiffusionPipeline(**__snake_case							)
							snake_case						:  Any        =			sd_pipe.to(__snake_case							)
							sd_pipe.set_progress_bar_config(disable=__snake_case							)

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

            # load LoRA weight file
							sd_pipe.load_lora_weights(__snake_case							)

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



			@unittest.skipIf(torch_device != "cuda"			,	"This test is supposed to run on GPU"							)
			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					snake_case						:  List[str]        =			self.get_dummy_components()
					snake_case						:  str        =			StableDiffusionPipeline(**__snake_case							)
					snake_case						:  Any        =			sd_pipe.to(__snake_case							)
					sd_pipe.set_progress_bar_config(disable=__snake_case							)

					snake_case						:  int        =			self.get_dummy_inputs()

					# enable XFormers
					sd_pipe.enable_xformers_memory_efficient_attention()

					snake_case						:  Tuple        =			sd_pipe(**__snake_case							).images
					snake_case						:  Dict        =			original_images[0, -3:, -3:, -1]

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

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

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







"""simple docstring"""

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


__A							=     logging.get_logger(__name__)

__A							=     {
    "dandelin/vilt-b32-mlm": "https://huggingface.co/dandelin/vilt-b32-mlm/blob/main/config.json"
}
class      lowerCamelCase__      (					lowerCamelCase_     ):
			a__	: Optional[int]        =       """vilt"""


			def __init__(       self			,	SCREAMING_SNAKE_CASE=30_522			,	SCREAMING_SNAKE_CASE=2			,	SCREAMING_SNAKE_CASE=2			,	SCREAMING_SNAKE_CASE=40			,	SCREAMING_SNAKE_CASE=768			,	SCREAMING_SNAKE_CASE=12			,	SCREAMING_SNAKE_CASE=12			,	SCREAMING_SNAKE_CASE=3_072			,	SCREAMING_SNAKE_CASE="gelu"			,	SCREAMING_SNAKE_CASE=0.0			,	SCREAMING_SNAKE_CASE=0.0			,	SCREAMING_SNAKE_CASE=0.02			,	SCREAMING_SNAKE_CASE=1E-12			,	SCREAMING_SNAKE_CASE=384			,	SCREAMING_SNAKE_CASE=32			,	SCREAMING_SNAKE_CASE=3			,	SCREAMING_SNAKE_CASE=True			,	SCREAMING_SNAKE_CASE=-1			,	SCREAMING_SNAKE_CASE=False			,	SCREAMING_SNAKE_CASE=-1			,	**SCREAMING_SNAKE_CASE			,	):



					"""simple docstring"""







					super().__init__(tie_word_embeddings=SCREAMING_SNAKE_CASE			,	**SCREAMING_SNAKE_CASE							)

					snake_case						:  str        =			vocab_size
					snake_case						:  Optional[Any]        =			type_vocab_size
					snake_case						:  List[str]        =			modality_type_vocab_size
					snake_case						:  Optional[Any]        =			max_position_embeddings

					snake_case						:  Tuple        =			hidden_size
					snake_case						:  List[str]        =			num_hidden_layers
					snake_case						:  int        =			num_attention_heads
					snake_case						:  Optional[int]        =			intermediate_size
					snake_case						:  List[Any]        =			hidden_act
					snake_case						:  List[str]        =			hidden_dropout_prob
					snake_case						:  Union[str, Any]        =			attention_probs_dropout_prob
					snake_case						:  Any        =			initializer_range
					snake_case						:  List[Any]        =			layer_norm_eps

					snake_case						:  Any        =			image_size
					snake_case						:  str        =			patch_size
					snake_case						:  Optional[int]        =			num_channels
					snake_case						:  List[str]        =			qkv_bias
					snake_case						:  Optional[int]        =			max_image_length
					snake_case						:  Any        =			num_images

