






import inspect
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_=1.0   )					->					Any:
														'''simple docstring'''

														super().__init__()

														lowerCAmelCase_										=     cut_size
														lowerCAmelCase_										=     cut_power






							def      _lowercase       (		self			,			lowercase_			,			lowercase_   )					->					int:
														'''simple docstring'''

														lowerCAmelCase_						,   lowerCAmelCase_										=     pixel_values.shape[2:4]
														lowerCAmelCase_										=     min(lowercase_			,			lowercase_   )
														lowerCAmelCase_										=     min(lowercase_			,			lowercase_			,			self.cut_size   )
														lowerCAmelCase_										=     []
														for _ in range(lowercase_   ):
																					lowerCAmelCase_										=     int(torch.rand([]   ) ** self.cut_power * (max_size - min_size) + min_size   )
																					lowerCAmelCase_										=     torch.randint(0			,			sideX - size + 1			,			()   )
																					lowerCAmelCase_										=     torch.randint(0			,			sideY - size + 1			,			()   )
																					lowerCAmelCase_										=     pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
																					cutouts.append(F.adaptive_avg_poolad(lowercase_			,			self.cut_size   )   )
														return torch.cat(lowercase_   )


def      lowerCamelCase					(						a_ ,						a_	)	->			Optional[int]:
							lowerCAmelCase_										=     F.normalize(a_ ,						dim=-1	)
							lowerCAmelCase_										=     F.normalize(a_ ,						dim=-1	)
							return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)
def      lowerCamelCase					(						a_ ,						a_	)	->			int:
							for param in model.parameters():
														lowerCAmelCase_										=     value




class 							a_ (						a_		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			)					->					str:
														'''simple docstring'''

														super().__init__()
														self.register_modules(
														    vae=lowercase_			,			text_encoder=lowercase_			,			clip_model=lowercase_			,			tokenizer=lowercase_			,			unet=lowercase_			,			scheduler=lowercase_			,			feature_extractor=lowercase_			,			)

														lowerCAmelCase_										=     transforms.Normalize(mean=feature_extractor.image_mean			,			std=feature_extractor.image_std   )
														lowerCAmelCase_										=     (
														    feature_extractor.size
														    if isinstance(feature_extractor.size			,			lowercase_   )
														    else feature_extractor.size['shortest_edge']
														)
														lowerCAmelCase_										=     MakeCutouts(self.cut_out_size   )

														set_requires_grad(self.text_encoder			,			lowercase_   )
														set_requires_grad(self.clip_model			,			lowercase_   )

							def      _lowercase       (		self			,			lowercase_ = "auto"   )					->					Union[str, Any]:
														'''simple docstring'''

														if slice_size == "auto":
																					# half the attention head size is usually a good trade-off between
																					# speed and memory
																					lowerCAmelCase_										=     self.unet.config.attention_head_dim // 2
														self.unet.set_attention_slice(lowercase_   )

							def      _lowercase       (		self   )					->					List[Any]:
														'''simple docstring'''

														self.enable_attention_slicing(lowercase_   )

							def      _lowercase       (		self   )					->					str:
														'''simple docstring'''

														set_requires_grad(self.vae			,			lowercase_   )

							def      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

														set_requires_grad(self.vae			,			lowercase_   )

							def      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

														set_requires_grad(self.unet			,			lowercase_   )

							def      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

														set_requires_grad(self.unet			,			lowercase_   )

							@torch.enable_grad()
							def      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_=True			,			)					->					Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     latents.detach().requires_grad_()

														lowerCAmelCase_										=     self.scheduler.scale_model_input(lowercase_			,			lowercase_   )

														# predict the noise residual
														lowerCAmelCase_										=     self.unet(lowercase_			,			lowercase_			,			encoder_hidden_states=lowercase_   ).sample

														if isinstance(self.scheduler			,			(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)   ):
																					lowerCAmelCase_										=     self.scheduler.alphas_cumprod[timestep]
																					lowerCAmelCase_										=     1 - alpha_prod_t
																					# compute predicted original sample from predicted noise also called
																					# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
																					lowerCAmelCase_										=     (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

																					lowerCAmelCase_										=     torch.sqrt(lowercase_   )
																					lowerCAmelCase_										=     pred_original_sample * (fac) + latents * (1 - fac)
														elif isinstance(self.scheduler			,			lowercase_   ):
																					lowerCAmelCase_										=     self.scheduler.sigmas[index]
																					lowerCAmelCase_										=     latents - sigma * noise_pred
														else:
																					raise ValueError(f'''scheduler type {type(self.scheduler   )} not supported'''   )

														lowerCAmelCase_										=     1 / self.vae.config.scaling_factor * sample
														lowerCAmelCase_										=     self.vae.decode(lowercase_   ).sample
														lowerCAmelCase_										=     (image / 2 + 0.5).clamp(0			,			1   )

														if use_cutouts:
																					lowerCAmelCase_										=     self.make_cutouts(lowercase_			,			lowercase_   )
														else:
																					lowerCAmelCase_										=     transforms.Resize(self.cut_out_size   )(lowercase_   )
														lowerCAmelCase_										=     self.normalize(lowercase_   ).to(latents.dtype   )

														lowerCAmelCase_										=     self.clip_model.get_image_features(lowercase_   )
														lowerCAmelCase_										=     image_embeddings_clip / image_embeddings_clip.norm(p=2			,			dim=-1			,			keepdim=lowercase_   )

														if use_cutouts:
																					lowerCAmelCase_										=     spherical_dist_loss(lowercase_			,			lowercase_   )
																					lowerCAmelCase_										=     dists.view([num_cutouts, sample.shape[0], -1]   )
																					lowerCAmelCase_										=     dists.sum(2   ).mean(0   ).sum() * clip_guidance_scale
														else:
																					lowerCAmelCase_										=     spherical_dist_loss(lowercase_			,			lowercase_   ).mean() * clip_guidance_scale

														lowerCAmelCase_										=     -torch.autograd.grad(lowercase_			,			lowercase_   )[0]

														if isinstance(self.scheduler			,			lowercase_   ):
																					lowerCAmelCase_										=     latents.detach() + grads * (sigma**2)
																					lowerCAmelCase_										=     noise_pred_original
														else:
																					lowerCAmelCase_										=     noise_pred_original - torch.sqrt(lowercase_   ) * grads
														return noise_pred, latents






							@torch.no_grad()
							def __call__(		self			,			lowercase_			,			lowercase_ = 5_1_2			,			lowercase_ = 5_1_2			,			lowercase_ = 5_0			,			lowercase_ = 7.5			,			lowercase_ = 1			,			lowercase_ = 0.0			,			lowercase_ = 1_0_0			,			lowercase_ = None			,			lowercase_ = 4			,			lowercase_ = True			,			lowercase_ = None			,			lowercase_ = None			,			lowercase_ = "pil"			,			lowercase_ = True			,			)					->					str:
														'''simple docstring'''

														if isinstance(lowercase_			,			lowercase_   ):
																					lowerCAmelCase_										=     1
														elif isinstance(lowercase_			,			lowercase_   ):
																					lowerCAmelCase_										=     len(lowercase_   )
														else:
																					raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(lowercase_   )}'''   )

														if height % 8 != 0 or width % 8 != 0:
																					raise ValueError(f'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''   )

														# get prompt text embeddings
														lowerCAmelCase_										=     self.tokenizer(
														    lowercase_			,			padding='max_length'			,			max_length=self.tokenizer.model_max_length			,			truncation=lowercase_			,			return_tensors='pt'			,			)
														lowerCAmelCase_										=     self.text_encoder(text_input.input_ids.to(self.device   )   )[0]
														# duplicate text embeddings for each generation per prompt
														lowerCAmelCase_										=     text_embeddings.repeat_interleave(lowercase_			,			dim=0   )

														if clip_guidance_scale > 0:
																					if clip_prompt is not None:
																												lowerCAmelCase_										=     self.tokenizer(
																												    lowercase_			,			padding='max_length'			,			max_length=self.tokenizer.model_max_length			,			truncation=lowercase_			,			return_tensors='pt'			,			).input_ids.to(self.device   )
																					else:
																												lowerCAmelCase_										=     text_input.input_ids.to(self.device   )
																					lowerCAmelCase_										=     self.clip_model.get_text_features(lowercase_   )
																					lowerCAmelCase_										=     text_embeddings_clip / text_embeddings_clip.norm(p=2			,			dim=-1			,			keepdim=lowercase_   )
																					# duplicate text embeddings clip for each generation per prompt
																					lowerCAmelCase_										=     text_embeddings_clip.repeat_interleave(lowercase_			,			dim=0   )

														# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
														# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
														# corresponds to doing no classifier free guidance.
														lowerCAmelCase_										=     guidance_scale > 1.0
														# get unconditional embeddings for classifier free guidance
														if do_classifier_free_guidance:
																					lowerCAmelCase_										=     text_input.input_ids.shape[-1]
																					lowerCAmelCase_										=     self.tokenizer(['']			,			padding='max_length'			,			max_length=lowercase_			,			return_tensors='pt'   )
																					lowerCAmelCase_										=     self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]
																					# duplicate unconditional embeddings for each generation per prompt
																					lowerCAmelCase_										=     uncond_embeddings.repeat_interleave(lowercase_			,			dim=0   )

																					# For classifier free guidance, we need to do two forward passes.
																					# Here we concatenate the unconditional and text embeddings into a single batch
																					# to avoid doing two forward passes
																					lowerCAmelCase_										=     torch.cat([uncond_embeddings, text_embeddings]   )

														# get the initial random noise unless the user supplied it

														# Unlike in other pipelines, latents need to be generated in the target device
														# for 1-to-1 results reproducibility with the CompVis implementation.
														# However this currently doesn't work in `mps`.
														lowerCAmelCase_										=     (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
														lowerCAmelCase_										=     text_embeddings.dtype
														if latents is None:
																					if self.device.type == "mps":
																												# randn does not work reproducibly on mps
																												lowerCAmelCase_										=     torch.randn(lowercase_			,			generator=lowercase_			,			device='cpu'			,			dtype=lowercase_   ).to(
																												    self.device   )
																					else:
																												lowerCAmelCase_										=     torch.randn(lowercase_			,			generator=lowercase_			,			device=self.device			,			dtype=lowercase_   )
														else:
																					if latents.shape != latents_shape:
																												raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''   )
																					lowerCAmelCase_										=     latents.to(self.device   )

														# set timesteps
														lowerCAmelCase_										=     'offset' in set(inspect.signature(self.scheduler.set_timesteps   ).parameters.keys()   )
														lowerCAmelCase_										=     {}
														if accepts_offset:
																					lowerCAmelCase_										=     1

														self.scheduler.set_timesteps(lowercase_			,			**lowercase_   )

														# Some schedulers like PNDM have timesteps as arrays
														# It's more optimized to move all timesteps to correct device beforehand
														lowerCAmelCase_										=     self.scheduler.timesteps.to(self.device   )

														# scale the initial noise by the standard deviation required by the scheduler
														lowerCAmelCase_										=     latents * self.scheduler.init_noise_sigma

														# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
														# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
														# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
														# and should be between [0, 1]
														lowerCAmelCase_										=     'eta' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
														lowerCAmelCase_										=     {}
														if accepts_eta:
																					lowerCAmelCase_										=     eta

														# check if the scheduler accepts generator
														lowerCAmelCase_										=     'generator' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
														if accepts_generator:
																					lowerCAmelCase_										=     generator

														for i, t in enumerate(self.progress_bar(lowercase_   )   ):
																					# expand the latents if we are doing classifier free guidance
																					lowerCAmelCase_										=     torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents
																					lowerCAmelCase_										=     self.scheduler.scale_model_input(lowercase_			,			lowercase_   )

																					# predict the noise residual
																					lowerCAmelCase_										=     self.unet(lowercase_			,			lowercase_			,			encoder_hidden_states=lowercase_   ).sample

																					# perform classifier free guidance
																					if do_classifier_free_guidance:
																												lowerCAmelCase_						,   lowerCAmelCase_										=     noise_pred.chunk(2   )
																												lowerCAmelCase_										=     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																					# perform clip guidance
																					if clip_guidance_scale > 0:
																												lowerCAmelCase_										=     (
																												    text_embeddings.chunk(2   )[1] if do_classifier_free_guidance else text_embeddings
																												)
																												lowerCAmelCase_						,   lowerCAmelCase_										=     self.cond_fn(
																												    lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			)

																					# compute the previous noisy sample x_t -> x_t-1
																					lowerCAmelCase_										=     self.scheduler.step(lowercase_			,			lowercase_			,			lowercase_			,			**lowercase_   ).prev_sample

														# scale and decode the image latents with vae
														lowerCAmelCase_										=     1 / self.vae.config.scaling_factor * latents
														lowerCAmelCase_										=     self.vae.decode(lowercase_   ).sample

														lowerCAmelCase_										=     (image / 2 + 0.5).clamp(0			,			1   )
														lowerCAmelCase_										=     image.cpu().permute(0			,			2			,			3			,			1   ).numpy()

														if output_type == "pil":
																					lowerCAmelCase_										=     self.numpy_to_pil(lowercase_   )

														if not return_dict:
																					return (image, None)

														return StableDiffusionPipelineOutput(images=lowercase_			,			nsfw_content_detected=lowercase_   )











import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


lowerCamelCase_						     =  logging.get_logger(__name__)




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self   )					->					Any:
														'''simple docstring'''

														super().__init__()
														if version.parse(torch.__version__   ) < version.parse('1.12.0'   ):
																					raise ImportError(
																					    f'''You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use '''
																					    'PytorchGELUTanh. Please upgrade torch.'   )






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return nn.functional.gelu(lowercase_			,			approximate='tanh'   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi   ) * (input + 0.04_47_15 * torch.pow(lowercase_			,			3.0   ))   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_ = False   )					->					Optional[Any]:
														'''simple docstring'''

														super().__init__()
														if use_gelu_python:
																					lowerCAmelCase_										=     self._gelu_python
														else:
																					lowerCAmelCase_										=     nn.functional.gelu

							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0   )   ))






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return self.act(lowercase_   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return 0.5 * input * (1.0 + torch.tanh(input * 0.79_78_84_56_08 * (1.0 + 0.04_47_15 * input * input)   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input * torch.sigmoid(1.7_02 * input   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_   )					->					Optional[int]:
														'''simple docstring'''

														if min > max:
																					raise ValueError(f'''min should be < max (got min: {min}, max: {max})'''   )

														super().__init__()
														lowerCAmelCase_										=     min
														lowerCAmelCase_										=     max






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return torch.clip(gelu(lowercase_   )			,			self.min			,			self.max   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self   )					->					Union[str, Any]:
														'''simple docstring'''

														super().__init__()
														lowerCAmelCase_										=     math.sqrt(2 / math.pi   )






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.04_47_15 * torch.pow(lowercase_			,			3   ))   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return nn.functional.silu(lowercase_   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self   )					->					List[Any]:
														'''simple docstring'''

														super().__init__()
														if version.parse(torch.__version__   ) < version.parse('1.9.0'   ):
																					lowerCAmelCase_										=     self._mish_python
														else:
																					lowerCAmelCase_										=     nn.functional.mish

							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input * torch.tanh(nn.functional.softplus(lowercase_   )   )






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return self.act(lowercase_   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_			,			lowercase_=0.70_71_07			,			lowercase_=0.28_20_95   )					->					Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     (input - mu).div(sigma * math.sqrt(2.0   )   )
														return 0.5 * (1.0 + torch.erf(lowercase_   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Any:
														'''simple docstring'''

														lowerCAmelCase_										=     nn.functional.relu(lowercase_   )
														lowerCAmelCase_										=     torch.square(lowercase_   )
														return squared




class 							a_ (						a_		):


							'''simple docstring'''






							def __getitem__(		self			,			lowercase_   )					->					Any:
														'''simple docstring'''

														lowerCAmelCase_										=     super().__getitem__(lowercase_   )
														lowerCAmelCase_						,   lowerCAmelCase_										=     content if isinstance(lowercase_			,			lowercase_   ) else (content, {})
														return cls(**lowercase_   )


lowerCamelCase_						     =  {
    """gelu""": GELUActivation,
    """gelu_10""": (ClippedGELUActivation, {"""min""": -1_0, """max""": 1_0}),
    """gelu_fast""": FastGELUActivation,
    """gelu_new""": NewGELUActivation,
    """gelu_python""": (GELUActivation, {"""use_gelu_python""": True}),
    """gelu_pytorch_tanh""": PytorchGELUTanh,
    """gelu_accurate""": AccurateGELUActivation,
    """laplace""": LaplaceActivation,
    """linear""": LinearActivation,
    """mish""": MishActivation,
    """quick_gelu""": QuickGELUActivation,
    """relu""": nn.ReLU,
    """relu2""": ReLUSquaredActivation,
    """relu6""": nn.ReLUa,
    """sigmoid""": nn.Sigmoid,
    """silu""": SiLUActivation,
    """swish""": SiLUActivation,
    """tanh""": nn.Tanh,
}
lowerCamelCase_						     =  ClassInstantier(ACTaCLS)
def      lowerCamelCase					(						a_	)	->			str:
							if activation_string in ACTaFN:
														return ACTaFN[activation_string]
							else:
														raise KeyError(F'''function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()	)}'''	)


# For backwards compatibility with: from activations import gelu_python
lowerCamelCase_						     =  get_activation("""gelu_python""")
lowerCamelCase_						     =  get_activation("""gelu_new""")
lowerCamelCase_						     =  get_activation("""gelu""")
lowerCamelCase_						     =  get_activation("""gelu_fast""")
lowerCamelCase_						     =  get_activation("""quick_gelu""")
lowerCamelCase_						     =  get_activation("""silu""")
lowerCamelCase_						     =  get_activation("""mish""")



lowerCamelCase_						     =  get_activation("""linear""")



