


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 			snake_case__(							nn.Module							):







						'''simple docstring'''




						def __init__(			self				,  __lowercase				,  __lowercase=1.0							)   ->      List[Any]:
										super().__init__()

										lowerCAmelCase_	: Optional[int]      			=			cut_size
										lowerCAmelCase_	: Any      			=			cut_power




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






def      lowerCAmelCase	(						lowerCAmelCase_	,				lowerCAmelCase_	)->	List[Any]:
				lowerCAmelCase_	: Dict      			=			F.normalize(lowerCAmelCase_	,				dim=-1	)
				lowerCAmelCase_	: int      			=			F.normalize(lowerCAmelCase_	,				dim=-1	)
				return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)



def      lowerCAmelCase	(						lowerCAmelCase_	,				lowerCAmelCase_	)->	Optional[int]:
				for param in model.parameters():
								lowerCAmelCase_	: Optional[int]      			=			value







class 			snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						def __init__(			self				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  )   ->      Optional[int]:
										super().__init__()
										self.register_modules(
										    vae=__lowercase				,  text_encoder=__lowercase				,  clip_model=__lowercase				,  tokenizer=__lowercase				,  unet=__lowercase				,  scheduler=__lowercase				,  feature_extractor=__lowercase				,  )

										lowerCAmelCase_	: Optional[int]      			=			transforms.Normalize(mean=feature_extractor.image_mean				,  std=feature_extractor.image_std							)
										lowerCAmelCase_	: List[str]      			=			(
										    feature_extractor.size
										    if isinstance(feature_extractor.size				,  __lowercase							)
										    else feature_extractor.size['''shortest_edge''']
										)
										lowerCAmelCase_	: Optional[Any]      			=			MakeCutouts(self.cut_out_size							)

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




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




						def   lowercase_	(			self							)   ->      Optional[Any]:
										self.enable_attention_slicing(__lowercase							)




						def   lowercase_	(			self							)   ->      Optional[int]:
										set_requires_grad(self.vae				,  __lowercase							)




						def   lowercase_	(			self							)   ->      str:
										set_requires_grad(self.vae				,  __lowercase							)




						def   lowercase_	(			self							)   ->      Tuple:
										set_requires_grad(self.unet				,  __lowercase							)




						def   lowercase_	(			self							)   ->      Dict:
										set_requires_grad(self.unet				,  __lowercase							)




						@torch.enable_grad()
						def   lowercase_	(			self				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase=True				,  )   ->      List[Any]:
										lowerCAmelCase_	: Union[str, Any]      			=			latents.detach().requires_grad_()

										lowerCAmelCase_	: Tuple      			=			self.scheduler.scale_model_input(__lowercase				,  __lowercase							)

										# predict the noise residual
										lowerCAmelCase_	: List[Any]      			=			self.unet(__lowercase				,  __lowercase				,  encoder_hidden_states=__lowercase							).sample

										if isinstance(self.scheduler				,  (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)							):
														lowerCAmelCase_	: List[Any]      			=			self.scheduler.alphas_cumprod[timestep]
														lowerCAmelCase_	: str      			=			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_	: str      			=			(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

														lowerCAmelCase_	: Tuple      			=			torch.sqrt(__lowercase							)
														lowerCAmelCase_	: List[Any]      			=			pred_original_sample * (fac) + latents * (1 - fac)
										elif isinstance(self.scheduler				,  __lowercase							):
														lowerCAmelCase_	: Optional[int]      			=			self.scheduler.sigmas[index]
														lowerCAmelCase_	: List[str]      			=			latents - sigma * noise_pred
										else:
														raise ValueError(f"""scheduler type {type(self.scheduler							)} not supported"""							)

										lowerCAmelCase_	: List[Any]      			=			1 / self.vae.config.scaling_factor * sample
										lowerCAmelCase_	: Optional[int]      			=			self.vae.decode(__lowercase							).sample
										lowerCAmelCase_	: Union[str, Any]      			=			(image / 2 + 0.5).clamp(0				,  1							)

										if use_cutouts:
														lowerCAmelCase_	: str      			=			self.make_cutouts(__lowercase				,  __lowercase							)
										else:
														lowerCAmelCase_	: List[Any]      			=			transforms.Resize(self.cut_out_size							)(__lowercase							)
										lowerCAmelCase_	: Any      			=			self.normalize(__lowercase							).to(latents.dtype							)

										lowerCAmelCase_	: Union[str, Any]      			=			self.clip_model.get_image_features(__lowercase							)
										lowerCAmelCase_	: Any      			=			image_embeddings_clip / image_embeddings_clip.norm(p=2				,  dim=-1				,  keepdim=__lowercase							)

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

										lowerCAmelCase_	: str      			=			-torch.autograd.grad(__lowercase				,  __lowercase							)[0]

										if isinstance(self.scheduler				,  __lowercase							):
														lowerCAmelCase_	: List[str]      			=			latents.detach() + grads * (sigma**2)
														lowerCAmelCase_	: Union[str, Any]      			=			noise_pred_original
										else:
														lowerCAmelCase_	: List[str]      			=			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				,  )   ->      Union[str, Any]:
										if isinstance(__lowercase				,  __lowercase							):
														lowerCAmelCase_	: str      			=			1
										elif isinstance(__lowercase				,  __lowercase							):
														lowerCAmelCase_	: List[str]      			=			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_	: List[str]      			=			self.tokenizer(
										    __lowercase				,  padding='''max_length'''				,  max_length=self.tokenizer.model_max_length				,  truncation=__lowercase				,  return_tensors='''pt'''				,  )
										lowerCAmelCase_	: Tuple      			=			self.text_encoder(text_input.input_ids.to(self.device							)							)[0]
										# duplicate text embeddings for each generation per prompt
										lowerCAmelCase_	: Optional[Any]      			=			text_embeddings.repeat_interleave(__lowercase				,  dim=0							)

										if clip_guidance_scale > 0:
														if clip_prompt is not None:
																		lowerCAmelCase_	: Union[str, Any]      			=			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_	: Tuple      			=			text_input.input_ids.to(self.device							)
														lowerCAmelCase_	: Tuple      			=			self.clip_model.get_text_features(__lowercase							)
														lowerCAmelCase_	: Tuple      			=			text_embeddings_clip / text_embeddings_clip.norm(p=2				,  dim=-1				,  keepdim=__lowercase							)
														# duplicate text embeddings clip for each generation per prompt
														lowerCAmelCase_	: Optional[Any]      			=			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_	: Optional[int]      			=			guidance_scale > 1.0
										# get unconditional embeddings for classifier free guidance
										if do_classifier_free_guidance:
														lowerCAmelCase_	: Any      			=			text_input.input_ids.shape[-1]
														lowerCAmelCase_	: Optional[int]      			=			self.tokenizer(['''''']				,  padding='''max_length'''				,  max_length=__lowercase				,  return_tensors='''pt'''							)
														lowerCAmelCase_	: Tuple      			=			self.text_encoder(uncond_input.input_ids.to(self.device							)							)[0]
														# duplicate unconditional embeddings for each generation per prompt
														lowerCAmelCase_	: Optional[Any]      			=			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_	: Optional[Any]      			=			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_	: Dict      			=			(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
										lowerCAmelCase_	: List[str]      			=			text_embeddings.dtype
										if latents is None:
														if self.device.type == "mps":
																		# randn does not work reproducibly on mps
																		lowerCAmelCase_	: Optional[int]      			=			torch.randn(__lowercase				,  generator=__lowercase				,  device='''cpu'''				,  dtype=__lowercase							).to(
																		    self.device							)
														else:
																		lowerCAmelCase_	: int      			=			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_	: Union[str, Any]      			=			latents.to(self.device							)

										# set timesteps
										lowerCAmelCase_	: Union[str, Any]      			=			'''offset''' in set(inspect.signature(self.scheduler.set_timesteps							).parameters.keys()							)
										lowerCAmelCase_	: str      			=			{}
										if accepts_offset:
														lowerCAmelCase_	: Tuple      			=			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_	: int      			=			self.scheduler.timesteps.to(self.device							)

										# scale the initial noise by the standard deviation required by the scheduler
										lowerCAmelCase_	: Tuple      			=			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_	: Dict      			=			'''eta''' in set(inspect.signature(self.scheduler.step							).parameters.keys()							)
										lowerCAmelCase_	: Optional[int]      			=			{}
										if accepts_eta:
														lowerCAmelCase_	: List[str]      			=			eta

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

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

														# predict the noise residual
														lowerCAmelCase_	: Union[str, Any]      			=			self.unet(__lowercase				,  __lowercase				,  encoder_hidden_states=__lowercase							).sample

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

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

														# compute the previous noisy sample x_t -> x_t-1
														lowerCAmelCase_	: List[Any]      			=			self.scheduler.step(__lowercase				,  __lowercase				,  __lowercase				,  **__lowercase							).prev_sample

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

										lowerCAmelCase_	: Dict      			=			(image / 2 + 0.5).clamp(0				,  1							)
										lowerCAmelCase_	: List[Any]      			=			image.cpu().permute(0				,  2				,  3				,  1							).numpy()

										if output_type == "pil":
														lowerCAmelCase_	: List[str]      			=			self.numpy_to_pil(__lowercase							)

										if not return_dict:
														return (image, None)

										return StableDiffusionPipelineOutput(images=__lowercase				,  nsfw_content_detected=__lowercase							)



from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
						import torch

_UpperCAmelCase     :					List[Any]       =logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
						from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):



						def      lowerCAmelCase	(						cls	)->	int:
										return cls



def      lowerCAmelCase	(						lowerCAmelCase_	,				lowerCAmelCase_ = None	,				lowerCAmelCase_ = None	,				lowerCAmelCase_ = None	,				lowerCAmelCase_ = None	,				)->	int:
				lowerCAmelCase_	: Tuple      			=			device
				lowerCAmelCase_	: List[str]      			=			shape[0]

				lowerCAmelCase_	: Union[str, Any]      			=			layout or torch.strided
				lowerCAmelCase_	: Union[str, Any]      			=			device or torch.device('''cpu'''	)

				if generator is not None:
								lowerCAmelCase_	: List[str]      			=			generator.device.type if not isinstance(lowerCAmelCase_	,				lowerCAmelCase_	) else generator[0].device.type
								if gen_device_type != device.type and gen_device_type == "cpu":
												lowerCAmelCase_	: Optional[Any]      			=			'''cpu'''
												if device != "mps":
																logger.info(
																    f"""The passed generator was created on 'cpu' even though a tensor on {device} was expected."""
																    f""" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"""
																    f""" slighly speed up this function by passing a generator that was created on the {device} device."""	)
								elif gen_device_type != device.type and gen_device_type == "cuda":
												raise ValueError(f"""Cannot generate a {device} tensor from a generator of type {gen_device_type}."""	)

				if isinstance(lowerCAmelCase_	,				lowerCAmelCase_	):
								lowerCAmelCase_	: Dict      			=			(1,) + shape[1:]
								lowerCAmelCase_	: Union[str, Any]      			=			[
								    torch.randn(lowerCAmelCase_	,				generator=generator[i]	,				device=lowerCAmelCase_	,				dtype=lowerCAmelCase_	,				layout=lowerCAmelCase_	)
								    for i in range(lowerCAmelCase_	)
								]
								lowerCAmelCase_	: Tuple      			=			torch.cat(lowerCAmelCase_	,				dim=0	).to(lowerCAmelCase_	)
				else:
								lowerCAmelCase_	: int      			=			torch.randn(lowerCAmelCase_	,				generator=lowerCAmelCase_	,				device=lowerCAmelCase_	,				dtype=lowerCAmelCase_	,				layout=lowerCAmelCase_	).to(lowerCAmelCase_	)

				return latents




def      lowerCAmelCase	(						lowerCAmelCase_	)->	Optional[Any]:
				if is_torch_version('''<'''	,				'''2.0.0'''	) or not hasattr(lowerCAmelCase_	,				'''_dynamo'''	):
								return False
				return isinstance(lowerCAmelCase_	,				torch._dynamo.eval_frame.OptimizedModule	)