
"""simple docstring"""
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       ):






							def __init__(					self						:			Optional[int]						, snake_case__						:			Union[str, Any]						, snake_case__						:			List[Any]=1.0    ):






												"""simple docstring"""



												super().__init__()

												__a:						Dict		    =					cut_size
												__a:						List[Any]		    =					cut_power






							def 							__lowercase  (					self						:			Optional[Any]						, snake_case__						:			str						, snake_case__						:			str    ):






												"""simple docstring"""



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






def     lowerCamelCase_    (	lowerCamelCase_	,				lowerCamelCase_ )     ->			Any:


					'''simple docstring'''
					__a:						Optional[int]		    =					F.normalize(lowerCamelCase_	,				dim=-1 )
					__a:						Dict		    =					F.normalize(lowerCamelCase_	,				dim=-1 )
					return (x - y).norm(dim=-1 ).div(2 ).arcsin().pow(2 ).mul(2 )


def     lowerCamelCase_    (	lowerCamelCase_	,				lowerCamelCase_ )     ->			Union[str, Any]:


					'''simple docstring'''
					for param in model.parameters():
										__a:						Tuple		    =					value




class    _snake_case  (						_lowercase       ):






							def __init__(					self						:			Any						, snake_case__						:			AutoencoderKL						, snake_case__						:			CLIPTextModel						, snake_case__						:			CLIPModel						, snake_case__						:			CLIPTokenizer						, snake_case__						:			UNetaDConditionModel						, snake_case__						:			Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]						, snake_case__						:			CLIPImageProcessor						, ):






												"""simple docstring"""



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

												__a:						Tuple		    =					transforms.Normalize(mean=feature_extractor.image_mean						, std=feature_extractor.image_std    )
												__a:						Tuple		    =					(
												    feature_extractor.size
												    if isinstance(feature_extractor.size						, snake_case__    )
												    else feature_extractor.size["""shortest_edge"""]
												)
												__a:						int		    =					MakeCutouts(self.cut_out_size    )

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






							def 							__lowercase  (					self						:			Any						, snake_case__						:			Optional[Union[str, int]] = "auto"    ):






												"""simple docstring"""



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






							def 							__lowercase  (					self						:			Optional[Any]    ):






												"""simple docstring"""



												self.enable_attention_slicing(snake_case__    )






							def 							__lowercase  (					self						:			int    ):






												"""simple docstring"""



												set_requires_grad(self.vae						, snake_case__    )






							def 							__lowercase  (					self						:			List[str]    ):






												"""simple docstring"""



												set_requires_grad(self.vae						, snake_case__    )






							def 							__lowercase  (					self						:			int    ):






												"""simple docstring"""



												set_requires_grad(self.unet						, snake_case__    )






							def 							__lowercase  (					self						:			Any    ):






												"""simple docstring"""



												set_requires_grad(self.unet						, snake_case__    )






							@torch.enable_grad()
							def 							__lowercase  (					self						:			Optional[Any]						, snake_case__						:			Union[str, Any]						, snake_case__						:			Optional[int]						, snake_case__						:			str						, snake_case__						:			List[Any]						, snake_case__						:			Dict						, snake_case__						:			Union[str, Any]						, snake_case__						:			str						, snake_case__						:			Any						, snake_case__						:			List[Any]=True						, ):






												"""simple docstring"""



												__a:						Dict		    =					latents.detach().requires_grad_()

												__a:						Optional[int]		    =					self.scheduler.scale_model_input(snake_case__						, snake_case__    )

												# predict the noise residual
												__a:						List[str]		    =					self.unet(snake_case__						, snake_case__						, encoder_hidden_states=snake_case__    ).sample

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

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

												__a:						Optional[Any]		    =					1 / self.vae.config.scaling_factor * sample
												__a:						Dict		    =					self.vae.decode(snake_case__    ).sample
												__a:						str		    =					(image / 2 + 0.5).clamp(0						, 1    )

												if use_cutouts:
																	__a:						int		    =					self.make_cutouts(snake_case__						, snake_case__    )
												else:
																	__a:						List[Any]		    =					transforms.Resize(self.cut_out_size    )(snake_case__    )
												__a:						Tuple		    =					self.normalize(snake_case__    ).to(latents.dtype    )

												__a:						Optional[int]		    =					self.clip_model.get_image_features(snake_case__    )
												__a:						List[str]		    =					image_embeddings_clip / image_embeddings_clip.norm(p=2						, dim=-1						, keepdim=snake_case__    )

												if use_cutouts:
																	__a:						List[Any]		    =					spherical_dist_loss(snake_case__						, snake_case__    )
																	__a:						Optional[Any]		    =					dists.view([num_cutouts, sample.shape[0], -1]    )
																	__a:						Optional[int]		    =					dists.sum(2    ).mean(0    ).sum() * clip_guidance_scale
												else:
																	__a:						Optional[Any]		    =					spherical_dist_loss(snake_case__						, snake_case__    ).mean() * clip_guidance_scale

												__a:						Any		    =					-torch.autograd.grad(snake_case__						, snake_case__    )[0]

												if isinstance(self.scheduler						, snake_case__    ):
																	__a:						Optional[Any]		    =					latents.detach() + grads * (sigma**2)
																	__a:						List[str]		    =					noise_pred_original
												else:
																	__a:						int		    =					noise_pred_original - torch.sqrt(snake_case__    ) * grads
												return noise_pred, latents






							@torch.no_grad()
							def __call__(					self						:			Dict						, snake_case__						:			Union[str, List[str]]						, snake_case__						:			Optional[int] = 5_1_2						, snake_case__						:			Optional[int] = 5_1_2						, snake_case__						:			Optional[int] = 5_0						, snake_case__						:			Optional[float] = 7.5						, snake_case__						:			Optional[int] = 1						, snake_case__						:			float = 0.0						, snake_case__						:			Optional[float] = 1_0_0						, snake_case__						:			Optional[Union[str, List[str]]] = None						, snake_case__						:			Optional[int] = 4						, snake_case__						:			Optional[bool] = True						, snake_case__						:			Optional[torch.Generator] = None						, snake_case__						:			Optional[torch.FloatTensor] = None						, snake_case__						:			Optional[str] = "pil"						, snake_case__						:			bool = True						, ):






												"""simple docstring"""



												if isinstance(snake_case__						, snake_case__    ):
																	__a:						List[Any]		    =					1
												elif isinstance(snake_case__						, snake_case__    ):
																	__a:						Dict		    =					len(snake_case__    )
												else:
																	raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(snake_case__    )}"""    )

												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
												__a:						Any		    =					self.tokenizer(
												    snake_case__						, padding="""max_length"""						, max_length=self.tokenizer.model_max_length						, truncation=snake_case__						, return_tensors="""pt"""						, )
												__a:						Optional[Any]		    =					self.text_encoder(text_input.input_ids.to(self.device    )    )[0]
												# duplicate text embeddings for each generation per prompt
												__a:						Tuple		    =					text_embeddings.repeat_interleave(snake_case__						, dim=0    )

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

												# set timesteps
												__a:						Optional[Any]		    =					"""offset""" in set(inspect.signature(self.scheduler.set_timesteps    ).parameters.keys()    )
												__a:						int		    =					{}
												if accepts_offset:
																	__a:						Dict		    =					1

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

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

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

												# check if the scheduler accepts generator
												__a:						Any		    =					"""generator""" in set(inspect.signature(self.scheduler.step    ).parameters.keys()    )
												if accepts_generator:
																	__a:						Union[str, Any]		    =					generator

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

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

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

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

																	# compute the previous noisy sample x_t -> x_t-1
																	__a:						Optional[int]		    =					self.scheduler.step(snake_case__						, snake_case__						, snake_case__						, **snake_case__    ).prev_sample

												# scale and decode the image latents with vae
												__a:						Union[str, Any]		    =					1 / self.vae.config.scaling_factor * latents
												__a:						Optional[int]		    =					self.vae.decode(snake_case__    ).sample

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

												if output_type == "pil":
																	__a:						Optional[Any]		    =					self.numpy_to_pil(snake_case__    )

												if not return_dict:
																	return (image, None)

												return StableDiffusionPipelineOutput(images=snake_case__						, nsfw_content_detected=snake_case__    )


"""simple docstring"""
from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


lowercase		:						Tuple =					logging.get_logger(__name__)




class    _snake_case  (						_lowercase       ):
							SCREAMING_SNAKE_CASE		     =						['''pixel_values''']






							def __init__(					self						:			Union[str, Any]						, snake_case__						:			bool = True						, snake_case__						:			Optional[Dict[str, int]] = None						, snake_case__						:			PILImageResampling = PILImageResampling.BILINEAR						, snake_case__						:			bool = True						, snake_case__						:			Union[int, float] = 1 / 2_5_5						, snake_case__						:			bool = True						, snake_case__						:			Optional[Union[float, List[float]]] = None						, snake_case__						:			Optional[Union[float, List[float]]] = None						, **snake_case__						:			Optional[int]						, ):






												"""simple docstring"""



												super().__init__(**snake_case__    )
												__a:						Dict		    =					size if size is not None else {"""height""": 2_2_4, """width""": 2_2_4}
												__a:						Optional[Any]		    =					get_size_dict(snake_case__    )
												__a:						Union[str, Any]		    =					do_resize
												__a:						str		    =					do_rescale
												__a:						Tuple		    =					do_normalize
												__a:						Dict		    =					size
												__a:						str		    =					resample
												__a:						Any		    =					rescale_factor
												__a:						Optional[int]		    =					image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
												__a:						int		    =					image_std if image_std is not None else IMAGENET_STANDARD_STD






							def 							__lowercase  (					self						:			Optional[Any]						, snake_case__						:			np.ndarray						, snake_case__						:			Dict[str, int]						, snake_case__						:			PILImageResampling = PILImageResampling.BILINEAR						, snake_case__						:			Optional[Union[str, ChannelDimension]] = None						, **snake_case__						:			List[str]						, ):






												"""simple docstring"""



												__a:						int		    =					get_size_dict(snake_case__    )
												if "height" not in size or "width" not in size:
																	raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""    )
												return resize(
												    snake_case__						, size=(size["""height"""], size["""width"""])						, resample=snake_case__						, data_format=snake_case__						, **snake_case__    )






							def 							__lowercase  (					self						:			List[str]						, snake_case__						:			np.ndarray						, snake_case__						:			float						, snake_case__						:			Optional[Union[str, ChannelDimension]] = None						, **snake_case__						:			Union[str, Any]    ):






												"""simple docstring"""



												return rescale(snake_case__						, scale=snake_case__						, data_format=snake_case__						, **snake_case__    )






							def 							__lowercase  (					self						:			Dict						, snake_case__						:			np.ndarray						, snake_case__						:			Union[float, List[float]]						, snake_case__						:			Union[float, List[float]]						, snake_case__						:			Optional[Union[str, ChannelDimension]] = None						, **snake_case__						:			Union[str, Any]						, ):






												"""simple docstring"""



												return normalize(snake_case__						, mean=snake_case__						, std=snake_case__						, data_format=snake_case__						, **snake_case__    )






							def 							__lowercase  (					self						:			Optional[int]						, snake_case__						:			ImageInput						, snake_case__						:			Optional[bool] = None						, snake_case__						:			Dict[str, int] = None						, snake_case__						:			PILImageResampling = None						, snake_case__						:			Optional[bool] = None						, snake_case__						:			Optional[float] = None						, snake_case__						:			Optional[bool] = None						, snake_case__						:			Optional[Union[float, List[float]]] = None						, snake_case__						:			Optional[Union[float, List[float]]] = None						, snake_case__						:			Optional[Union[str, TensorType]] = None						, snake_case__						:			Union[str, ChannelDimension] = ChannelDimension.FIRST						, **snake_case__						:			List[str]						, ):






												"""simple docstring"""



												__a:						Any		    =					do_resize if do_resize is not None else self.do_resize
												__a:						Any		    =					do_rescale if do_rescale is not None else self.do_rescale
												__a:						int		    =					do_normalize if do_normalize is not None else self.do_normalize
												__a:						Any		    =					resample if resample is not None else self.resample
												__a:						Any		    =					rescale_factor if rescale_factor is not None else self.rescale_factor
												__a:						Dict		    =					image_mean if image_mean is not None else self.image_mean
												__a:						List[str]		    =					image_std if image_std is not None else self.image_std

												__a:						int		    =					size if size is not None else self.size
												__a:						Any		    =					get_size_dict(snake_case__    )

												__a:						List[str]		    =					make_list_of_images(snake_case__    )

												if not valid_images(snake_case__    ):
																	raise ValueError(
																	    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
																	    """torch.Tensor, tf.Tensor or jax.ndarray."""    )

												if do_resize and size is None:
																	raise ValueError("""Size must be specified if do_resize is True."""    )

												if do_rescale and rescale_factor is None:
																	raise ValueError("""Rescale factor must be specified if do_rescale is True."""    )

												# All transformations expect numpy arrays.
												__a:						str		    =					[to_numpy_array(snake_case__    ) for image in images]

												if do_resize:
																	__a:						Dict		    =					[self.resize(image=snake_case__						, size=snake_case__						, resample=snake_case__    ) for image in images]

												if do_rescale:
																	__a:						str		    =					[self.rescale(image=snake_case__						, scale=snake_case__    ) for image in images]

												if do_normalize:
																	__a:						Dict		    =					[self.normalize(image=snake_case__						, mean=snake_case__						, std=snake_case__    ) for image in images]

												__a:						int		    =					[to_channel_dimension_format(snake_case__						, snake_case__    ) for image in images]

												__a:						str		    =					{"""pixel_values""": images}
												return BatchFeature(data=snake_case__						, tensor_type=snake_case__    )
