



'''simple docstring'''






import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin

@dataclass
class     SCREAMING_SNAKE_CASE(							lowercase_      ):







	"""simple docstring"""






	lowerCamelCase__      			=	42
	lowerCamelCase__      			=	42


def 			snake_case_					( _lowerCAmelCase  :						List[str]       ,				_lowerCAmelCase  :						List[Any]=0.9_9_9       ,				_lowerCAmelCase  :						Dict="cosine"       ,				)				->       Dict:






							"""simple docstring"""




							if alpha_transform_type == "cosine":

														def alpha_bar_fn(_lowerCAmelCase  :						int  ):
																					return math.cos((t + 0.0_0_8) / 1.0_0_8 * math.pi / 2  ) ** 2

							elif alpha_transform_type == "exp":

														def alpha_bar_fn(_lowerCAmelCase  :						Union[str, Any]  ):
																					return math.exp(t * -1_2.0  )

							else:
														raise ValueError(f"""Unsupported alpha_tranform_type: {alpha_transform_type}"""  )

							UpperCAmelCase			:	Dict										=		[]
							for i in range(_lowerCAmelCase  ):
														UpperCAmelCase			:	Union[str, Any]										=		i / num_diffusion_timesteps
														UpperCAmelCase			:	Any										=		(i + 1) / num_diffusion_timesteps
														betas.append(min(1 - alpha_bar_fn(_lowerCAmelCase  ) / alpha_bar_fn(_lowerCAmelCase  )       ,				_lowerCAmelCase  )  )
							return torch.tensor(_lowerCAmelCase       ,				dtype=torch.floataa  )




class     SCREAMING_SNAKE_CASE(							lowercase_							,		lowercase_      ):







	"""simple docstring"""






	lowerCamelCase__      			=	1







	@register_to_config
	def __init__(						self      :					str		,      __snake_case      :					int = 1000		,      __snake_case      :					float = 0.00_01		,      __snake_case      :					float = 0.02		,      __snake_case      :					str = "linear"		,      __snake_case      :					float = 0.0		,      __snake_case      :					Optional[np.ndarray] = None		,      __snake_case      :					bool = True		,      )      ->		Tuple:
								if trained_betas is not None:
															UpperCAmelCase			:	List[Any]										=		torch.from_numpy(__UpperCamelCase		)
								elif beta_schedule == "linear":
															UpperCAmelCase			:	List[str]										=		torch.linspace(__UpperCamelCase		,      __UpperCamelCase		,      __UpperCamelCase		,      dtype=torch.floataa		)
								elif beta_schedule == "scaled_linear":
															# this schedule is very specific to the latent diffusion model.
															UpperCAmelCase			:	Dict										=		(
															    torch.linspace(beta_start**0.5		,      beta_end**0.5		,      __UpperCamelCase		,      dtype=torch.floataa		) ** 2
															)
								elif beta_schedule == "squaredcos_cap_v2":
															# Glide cosine schedule
															UpperCAmelCase			:	Tuple										=		betas_for_alpha_bar(__UpperCamelCase		)
								elif beta_schedule == "sigmoid":
															# GeoDiff sigmoid schedule
															UpperCAmelCase			:	str										=		torch.linspace(-6		,      6		,      __UpperCamelCase		)
															UpperCAmelCase			:	Optional[int]										=		torch.sigmoid(__UpperCamelCase		) * (beta_end - beta_start) + beta_start
								else:
															raise NotImplementedError(F"""{beta_schedule} does is not implemented for {self.__class__}"""		)

								UpperCAmelCase			:	Optional[Any]										=		1.0 - self.betas
								UpperCAmelCase			:	Tuple										=		torch.cumprod(self.alphas		,      dim=0		)
								UpperCAmelCase			:	Union[str, Any]										=		torch.tensor(1.0		)

								UpperCAmelCase			:	List[str]										=		torch.tensor(1.0		)

								# standard deviation of the initial noise distribution
								UpperCAmelCase			:	Optional[Any]										=		1.0

								# setable values
								UpperCAmelCase			:	Any										=		None
								UpperCAmelCase			:	Tuple										=		torch.from_numpy(np.arange(0		,      __UpperCamelCase		)[::-1].copy()		)

								UpperCAmelCase			:	Optional[int]										=		eta







	def 					A     (						self      :					Tuple		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					Optional[int] = None		)      ->		Optional[Any]:
								return sample







	def 					A     (						self      :					Any		,      __snake_case      :					int		,      __snake_case      :					int = 10		,      __snake_case      :					int = 10		,      __snake_case      :					Union[str, torch.device] = None		,      )      ->		Tuple:
								UpperCAmelCase			:	str										=		min(self.config.num_train_timesteps		,      __UpperCamelCase		)
								UpperCAmelCase			:	Union[str, Any]										=		num_inference_steps

								UpperCAmelCase			:	int										=		[]

								UpperCAmelCase			:	Dict										=		{}
								for j in range(0		,      num_inference_steps - jump_length		,      __UpperCamelCase		):
															UpperCAmelCase			:	List[str]										=		jump_n_sample - 1

								UpperCAmelCase			:	Dict										=		num_inference_steps
								while t >= 1:
															UpperCAmelCase			:	List[Any]										=		t - 1
															timesteps.append(__UpperCamelCase		)

															if jumps.get(__UpperCamelCase		,      0		) > 0:
																						UpperCAmelCase			:	Union[str, Any]										=		jumps[t] - 1
																						for _ in range(__UpperCamelCase		):
																													UpperCAmelCase			:	int										=		t + 1
																													timesteps.append(__UpperCamelCase		)

								UpperCAmelCase			:	Union[str, Any]										=		np.array(__UpperCamelCase		) * (self.config.num_train_timesteps // self.num_inference_steps)
								UpperCAmelCase			:	Optional[int]										=		torch.from_numpy(__UpperCamelCase		).to(__UpperCamelCase		)







	def 					A     (						self      :					List[Any]		,      __snake_case      :					Optional[Any]		)      ->		Optional[Any]:
								UpperCAmelCase			:	Optional[int]										=		t - self.config.num_train_timesteps // self.num_inference_steps

								UpperCAmelCase			:	List[str]										=		self.alphas_cumprod[t]
								UpperCAmelCase			:	Dict										=		self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
								UpperCAmelCase			:	List[Any]										=		1 - alpha_prod_t
								UpperCAmelCase			:	List[Any]										=		1 - alpha_prod_t_prev

								# For t > 0, compute predicted variance βt (see formula (6) and (7) from
								# https://arxiv.org/pdf/2006.11239.pdf) and sample from it to get
								# previous sample x_{t-1} ~ N(pred_prev_sample, variance) == add
								# variance to pred_sample
								# Is equivalent to formula (16) in https://arxiv.org/pdf/2010.02502.pdf
								# without eta.
								# variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * self.betas[t]
								UpperCAmelCase			:	Optional[Any]										=		(beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

								return variance







	def 					A     (						self      :					Union[str, Any]		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					int		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					Optional[torch.Generator] = None		,      __snake_case      :					bool = True		,      )      ->		List[Any]:
								UpperCAmelCase			:	Tuple										=		timestep
								UpperCAmelCase			:	Tuple										=		timestep - self.config.num_train_timesteps // self.num_inference_steps

								# 1. compute alphas, betas
								UpperCAmelCase			:	Tuple										=		self.alphas_cumprod[t]
								UpperCAmelCase			:	Any										=		self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
								UpperCAmelCase			:	str										=		1 - alpha_prod_t

								# 2. compute predicted original sample from predicted noise also called
								# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
								UpperCAmelCase			:	List[Any]										=		(sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5

								# 3. Clip "predicted x_0"
								if self.config.clip_sample:
															UpperCAmelCase			:	List[str]										=		torch.clamp(__UpperCamelCase		,      -1		,      1		)

								# We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we
								# substitute formula (7) in the algorithm coming from DDPM paper
								# (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper.
								# DDIM schedule gives the same results as DDPM with eta = 1.0
								# Noise is being reused in 7. and 8., but no impact on quality has
								# been observed.

								# 5. Add noise
								UpperCAmelCase			:	int										=		model_output.device
								UpperCAmelCase			:	Union[str, Any]										=		randn_tensor(model_output.shape		,      generator=__UpperCamelCase		,      device=__UpperCamelCase		,      dtype=model_output.dtype		)
								UpperCAmelCase			:	Union[str, Any]										=		self.eta * self._get_variance(__UpperCamelCase		) ** 0.5

								UpperCAmelCase			:	Any										=		0
								if t > 0 and self.eta > 0:
															UpperCAmelCase			:	Optional[int]										=		std_dev_t * noise

								# 6. compute "direction pointing to x_t" of formula (12)
								# from https://arxiv.org/pdf/2010.02502.pdf
								UpperCAmelCase			:	str										=		(1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * model_output

								# 7. compute x_{t-1} of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
								UpperCAmelCase			:	Optional[Any]										=		alpha_prod_t_prev**0.5 * pred_original_sample + pred_sample_direction + variance

								# 8. Algorithm 1 Line 5 https://arxiv.org/pdf/2201.09865.pdf
								UpperCAmelCase			:	List[Any]										=		(alpha_prod_t_prev**0.5) * original_image + ((1 - alpha_prod_t_prev) ** 0.5) * noise

								# 9. Algorithm 1 Line 8 https://arxiv.org/pdf/2201.09865.pdf
								UpperCAmelCase			:	Dict										=		mask * prev_known_part + (1.0 - mask) * prev_unknown_part

								if not return_dict:
															return (
															    pred_prev_sample,
															    pred_original_sample,
															)

								return RePaintSchedulerOutput(prev_sample=__UpperCamelCase		,      pred_original_sample=__UpperCamelCase		)







	def 					A     (						self      :					str		,      __snake_case      :					Optional[Any]		,      __snake_case      :					List[str]		,      __snake_case      :					List[Any]=None		)      ->		Optional[Any]:
								UpperCAmelCase			:	str										=		self.config.num_train_timesteps // self.num_inference_steps

								for i in range(__UpperCamelCase		):
															UpperCAmelCase			:	List[Any]										=		self.betas[timestep + i]
															if sample.device.type == "mps":
																						# randn does not work reproducibly on mps
																						UpperCAmelCase			:	int										=		randn_tensor(sample.shape		,      dtype=sample.dtype		,      generator=__UpperCamelCase		)
																						UpperCAmelCase			:	str										=		noise.to(sample.device		)
															else:
																						UpperCAmelCase			:	Optional[Any]										=		randn_tensor(sample.shape		,      generator=__UpperCamelCase		,      device=sample.device		,      dtype=sample.dtype		)

															# 10. Algorithm 1 Line 10 https://arxiv.org/pdf/2201.09865.pdf
															UpperCAmelCase			:	Any										=		(1 - beta) ** 0.5 * sample + beta**0.5 * noise

								return sample







	def 					A     (						self      :					List[Any]		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					torch.FloatTensor		,      __snake_case      :					torch.IntTensor		,      )      ->		Tuple:
								raise NotImplementedError('''Use `DDPMScheduler.add_noise()` to train for sampling with RePaint.'''		)







	def __len__(						self      :					List[str]		)      ->		str:
								return self.config.num_train_timesteps










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


def 			snake_case_					( _lowerCAmelCase  :						nn.Module  )				->       int:
							UpperCAmelCase			:	Any										=		{}
							UpperCAmelCase			:	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  ):
																					UpperCAmelCase			:	List[Any]										=		lora_attn_processor_class(
																					    hidden_size=module.out_proj.out_features       ,				cross_attention_dim=_lowerCAmelCase  )
							return text_lora_attn_procs


def 			snake_case_					( _lowerCAmelCase  :						nn.Module  )				->       Optional[Any]:
							UpperCAmelCase			:	Dict										=		create_text_encoder_lora_attn_procs(_lowerCAmelCase  )
							UpperCAmelCase			:	Union[str, Any]										=		AttnProcsLayers(_lowerCAmelCase  )
							return text_encoder_lora_layers


def 			snake_case_					( _lowerCAmelCase  :						List[str]       ,				_lowerCAmelCase  :						Optional[int]=False  )				->       int:
							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'''  ):
																												UpperCAmelCase			:	List[Any]										=		(
																												    torch.randn_like(layer_module.up.weight  )
																												    if randn_weight
																												    else torch.zeros_like(layer_module.up.weight  )
																												)
																												UpperCAmelCase			:	Optional[Any]										=		torch.nn.Parameter(_lowerCAmelCase  )




class     SCREAMING_SNAKE_CASE(							unittest.TestCase      ):







	"""simple docstring"""






	def 					A     (						self      :					str		)      ->		Dict:
								torch.manual_seed(0		)
								UpperCAmelCase			:	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		,      )
								UpperCAmelCase			:	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		)
								UpperCAmelCase			:	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		,      )
								UpperCAmelCase			:	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=1000		,      )
								UpperCAmelCase			:	Union[str, Any]										=		CLIPTextModel(__snake_case		)
								UpperCAmelCase			:	List[str]										=		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

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

								UpperCAmelCase			:	Union[str, Any]										=		{
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								UpperCAmelCase			:	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 					A     (						self      :					Union[str, Any]		)      ->		Tuple:
								UpperCAmelCase			:	Optional[int]										=		1
								UpperCAmelCase			:	str										=		10
								UpperCAmelCase			:	int										=		4
								UpperCAmelCase			:	Tuple										=		(32, 32)

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

								UpperCAmelCase			:	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 					A     (						self      :					str		)      ->		Tuple:
								UpperCAmelCase			:	Tuple										=		77

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

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







	def 					A     (						self      :					Any		,      __snake_case      :					Dict		)      ->		Union[str, Any]:
								UpperCAmelCase							,		UpperCAmelCase			:	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 					A     (						self      :					List[str]		)      ->		List[Any]:
								UpperCAmelCase							,		UpperCAmelCase			:	List[str]										=		self.get_dummy_components()
								UpperCAmelCase			:	Any										=		StableDiffusionPipeline(**__snake_case		)
								UpperCAmelCase			:	Dict										=		sd_pipe.to(__snake_case		)
								sd_pipe.set_progress_bar_config(disable=__snake_case		)

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

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

								UpperCAmelCase			:	Optional[int]										=		sd_pipe(**__snake_case		).images
								UpperCAmelCase			:	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 					A     (						self      :					Any		)      ->		List[Any]:
								UpperCAmelCase							,		UpperCAmelCase			:	int										=		self.get_dummy_components()
								UpperCAmelCase			:	Optional[Any]										=		StableDiffusionPipeline(**__snake_case		)
								UpperCAmelCase			:	Any										=		sd_pipe.to(__snake_case		)
								sd_pipe.set_progress_bar_config(disable=__snake_case		)

								UpperCAmelCase							,		UpperCAmelCase							,		UpperCAmelCase			:	Optional[int]										=		self.get_dummy_inputs()

								UpperCAmelCase			:	str										=		sd_pipe(**__snake_case		).images
								UpperCAmelCase			:	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		)

								UpperCAmelCase			:	Any										=		sd_pipe(**__snake_case		).images
								UpperCAmelCase			:	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 					A     (						self      :					List[str]		)      ->		int:
								UpperCAmelCase							,		UpperCAmelCase			:	Tuple										=		self.get_dummy_components()
								UpperCAmelCase			:	Optional[int]										=		lora_components['''unet_lora_attn_procs''']
								UpperCAmelCase			:	Any										=		StableDiffusionPipeline(**__snake_case		)
								UpperCAmelCase			:	List[str]										=		sd_pipe.to(__snake_case		)
								sd_pipe.set_progress_bar_config(disable=__snake_case		)

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

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

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

								UpperCAmelCase			:	Dict										=		sd_pipe(**__snake_case		).images
								UpperCAmelCase			:	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 					A     (						self      :					Optional[int]		)      ->		List[Any]:
								UpperCAmelCase							,		UpperCAmelCase			:	Tuple										=		self.get_dummy_components()
								UpperCAmelCase			:	Any										=		StableDiffusionPipeline(**__snake_case		)

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

								# inference without lora
								UpperCAmelCase			:	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
								UpperCAmelCase			:	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
								UpperCAmelCase			:	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
								UpperCAmelCase			:	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
								UpperCAmelCase			:	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 					A     (						self      :					Optional[int]		)      ->		str:
								UpperCAmelCase							,		UpperCAmelCase			:	Tuple										=		self.get_dummy_components()
								UpperCAmelCase			:	Optional[int]										=		StableDiffusionPipeline(**__snake_case		)

								UpperCAmelCase			:	Tuple										=		self.get_dummy_tokens()

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

								# create lora_attn_procs with randn up.weights
								UpperCAmelCase			:	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
								UpperCAmelCase			:	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
								UpperCAmelCase			:	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 					A     (						self      :					str		)      ->		List[Any]:
								UpperCAmelCase							,		UpperCAmelCase			:	int										=		self.get_dummy_components()
								UpperCAmelCase			:	Tuple										=		StableDiffusionPipeline(**__snake_case		)
								UpperCAmelCase			:	Union[str, Any]										=		sd_pipe.to(__snake_case		)
								sd_pipe.set_progress_bar_config(disable=__snake_case		)

								UpperCAmelCase							,		UpperCAmelCase							,		UpperCAmelCase			:	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		)

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

								UpperCAmelCase			:	Dict										=		sd_pipe(**__snake_case		,      cross_attention_kwargs={'''scale''': 0.5}		).images
								UpperCAmelCase			:	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 					A     (						self      :					Tuple		)      ->		Dict:
								with tempfile.TemporaryDirectory() as tmpdirname:
															self.create_lora_weight_file(__snake_case		)

															UpperCAmelCase							,		UpperCAmelCase			:	List[str]										=		self.get_dummy_components()
															UpperCAmelCase			:	Tuple										=		StableDiffusionPipeline(**__snake_case		)
															UpperCAmelCase			:	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		):
																													UpperCAmelCase			:	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 					A     (						self      :					List[str]		)      ->		Optional[Any]:
								with tempfile.TemporaryDirectory() as tmpdirname:
															self.create_lora_weight_file(__snake_case		)

															UpperCAmelCase							,		UpperCAmelCase			:	Tuple										=		self.get_dummy_components()
															UpperCAmelCase			:	List[str]										=		StableDiffusionPipeline(**__snake_case		)
															UpperCAmelCase			:	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 					A     (						self      :					List[Any]		)      ->		Any:
								UpperCAmelCase							,		UpperCAmelCase			:	List[str]										=		self.get_dummy_components()
								UpperCAmelCase			:	str										=		StableDiffusionPipeline(**__snake_case		)
								UpperCAmelCase			:	Any										=		sd_pipe.to(__snake_case		)
								sd_pipe.set_progress_bar_config(disable=__snake_case		)

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

								# enable XFormers
								sd_pipe.enable_xformers_memory_efficient_attention()

								UpperCAmelCase			:	Tuple										=		sd_pipe(**__snake_case		).images
								UpperCAmelCase			:	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		)

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





