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




def 							__SCREAMING_SNAKE_CASE       (		SCREAMING_SNAKE_CASE	:				nn.Module      )					->						int:
				__lowercase								= {}
				__lowercase								= (
				    LoRAAttnProcessor2_0 if hasattr(SCREAMING_SNAKE_CASE   ,							'scaled_dot_product_attention'      ) else LoRAAttnProcessor
				)
				for name, module in text_encoder.named_modules():
								if name.endswith(SCREAMING_SNAKE_CASE      ):
												__lowercase								= lora_attn_processor_class(
												    hidden_size=module.out_proj.out_features   ,							cross_attention_dim=SCREAMING_SNAKE_CASE      )
				return text_lora_attn_procs




def 							__SCREAMING_SNAKE_CASE       (		SCREAMING_SNAKE_CASE	:				nn.Module      )					->						Tuple:
				__lowercase								= create_text_encoder_lora_attn_procs(SCREAMING_SNAKE_CASE      )
				__lowercase								= AttnProcsLayers(SCREAMING_SNAKE_CASE      )
				return text_encoder_lora_layers




def 							__SCREAMING_SNAKE_CASE       (		SCREAMING_SNAKE_CASE	:				List[Any]   ,							SCREAMING_SNAKE_CASE	:				int=False      )					->						Dict:
				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'      ):
																__lowercase								= (
																    torch.randn_like(layer_module.up.weight      )
																    if randn_weight
																    else torch.zeros_like(layer_module.up.weight      )
																)
																__lowercase								= torch.nn.Parameter(SCREAMING_SNAKE_CASE      )
class      A__       (						unittest.TestCase    ):





							def       a__			(			self  :       Dict    )						->		Optional[Any]:



											"""simple docstring"""


											torch.manual_seed(0    )
											__lowercase								= 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					,			)
											__lowercase								= 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    )
											__lowercase								= 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					,			)
											__lowercase								= CLIPTextConfig(
											    bos_token_id=0					,			eos_token_id=2					,			hidden_size=32					,			intermediate_size=37					,			layer_norm_eps=1e-0_5					,			num_attention_heads=4					,			num_hidden_layers=5					,			pad_token_id=1					,			vocab_size=10_00					,			)
											__lowercase								= CLIPTextModel(_UpperCAmelCase    )
											__lowercase								= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

											__lowercase		, __lowercase								= create_unet_lora_layers(_UpperCAmelCase    )
											__lowercase								= create_text_encoder_lora_layers(_UpperCAmelCase    )

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



											"""simple docstring"""


											__lowercase								= 1
											__lowercase								= 10
											__lowercase								= 4
											__lowercase								= (32, 32)

											__lowercase								= torch.manual_seed(0    )
											__lowercase								= floats_tensor((batch_size, num_channels) + sizes    )
											__lowercase								= torch.randint(1					,			_UpperCAmelCase					,			size=(batch_size, sequence_length)					,			generator=_UpperCAmelCase    )

											__lowercase								= {
											    '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       a__			(			self  :       int    )						->		Any:



											"""simple docstring"""


											__lowercase								= 77

											__lowercase								= torch.randint(2					,			56					,			size=(1, max_seq_length)					,			generator=torch.manual_seed(0    )    )

											__lowercase								= {}
											__lowercase								= inputs
											return prepared_inputs





							def       a__			(			self  :       Union[str, Any]					,			_UpperCAmelCase  :       str    )						->		str:



											"""simple docstring"""


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



											"""simple docstring"""


											__lowercase		, __lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionPipeline(**_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase		, __lowercase		, __lowercase								= self.get_dummy_inputs()

											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images
											__lowercase								= 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    )

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

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





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



											"""simple docstring"""


											__lowercase		, __lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionPipeline(**_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase		, __lowercase		, __lowercase								= self.get_dummy_inputs()

											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images
											__lowercase								= 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    )

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

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





							def       a__			(			self  :       Any    )						->		str:



											"""simple docstring"""


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

											__lowercase		, __lowercase		, __lowercase								= self.get_dummy_inputs()

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

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

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

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





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



											"""simple docstring"""


											__lowercase		, __lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionPipeline(**_UpperCAmelCase    )

											__lowercase								= self.get_dummy_tokens()

											# inference without lora
											__lowercase								= pipe.text_encoder(**_UpperCAmelCase    )[0]
											assert outputs_without_lora.shape == (1, 77, 32)

											# create lora_attn_procs with zeroed out up.weights
											__lowercase								= 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
											__lowercase								= 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
											__lowercase								= 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
											__lowercase								= 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       a__			(			self  :       List[Any]    )						->		str:



											"""simple docstring"""


											__lowercase		, __lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionPipeline(**_UpperCAmelCase    )

											__lowercase								= self.get_dummy_tokens()

											# inference without lora
											__lowercase								= pipe.text_encoder(**_UpperCAmelCase    )[0]
											assert outputs_without_lora.shape == (1, 77, 32)

											# create lora_attn_procs with randn up.weights
											__lowercase								= 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
											__lowercase								= 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
											__lowercase								= 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       a__			(			self  :       List[str]    )						->		Optional[Any]:



											"""simple docstring"""


											__lowercase		, __lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionPipeline(**_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase		, __lowercase		, __lowercase								= 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    )

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

											__lowercase								= sd_pipe(**_UpperCAmelCase					,			cross_attention_kwargs={'scale': 0.5}    ).images
											__lowercase								= 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       a__			(			self  :       List[Any]    )						->		str:



											"""simple docstring"""


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

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



											"""simple docstring"""


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

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



											"""simple docstring"""


											__lowercase		, __lowercase								= self.get_dummy_components()
											__lowercase								= StableDiffusionPipeline(**_UpperCAmelCase    )
											__lowercase								= sd_pipe.to(_UpperCAmelCase    )
											sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

											__lowercase		, __lowercase		, __lowercase								= self.get_dummy_inputs()

											# enable XFormers
											sd_pipe.enable_xformers_memory_efficient_attention()

											__lowercase								= sd_pipe(**_UpperCAmelCase    ).images
											__lowercase								= 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    )

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

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







from typing import TYPE_CHECKING

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


SCREAMING_SNAKE_CASE__							=    {"""configuration_hubert""": ["""HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """HubertConfig"""]}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__							=    [
							    """HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """HubertForCTC""",
							    """HubertForSequenceClassification""",
							    """HubertModel""",
							    """HubertPreTrainedModel""",
							]


try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__							=    [
							    """TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """TFHubertForCTC""",
							    """TFHubertModel""",
							    """TFHubertPreTrainedModel""",
							]

if TYPE_CHECKING:
							from .configuration_hubert import HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, HubertConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_hubert import (
														    HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    HubertForCTC,
														    HubertForSequenceClassification,
														    HubertModel,
														    HubertPreTrainedModel,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_hubert import (
														    TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFHubertForCTC,
														    TFHubertModel,
														    TFHubertPreTrainedModel,
														)


else:
							import sys

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





