


'''simple docstring'''



from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import is_torch_available, is_transformers_available


if is_transformers_available() and is_torch_available():
      from .image_encoder import PaintByExampleImageEncoder
      from .pipeline_paint_by_example import PaintByExamplePipeline










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





    """simple docstring"""






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

         __lowerCamelCase       :     int        =		cut_size
         __lowerCamelCase       :     Optional[int]        =		cut_power







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







def       UpperCAmelCase__ (     UpperCAmelCase_   :    List[Any]				,   UpperCAmelCase_   :    Tuple							)					->		str:
     __lowerCamelCase       :     List[str]        =		F.normalize(UpperCAmelCase_				,   dim=-1							)
     __lowerCamelCase       :     Optional[Any]        =		F.normalize(UpperCAmelCase_				,   dim=-1							)
     return (x - y).norm(dim=-1							).div(2							).arcsin().pow(2							).mul(2							)



def       UpperCAmelCase__ (     UpperCAmelCase_   :    Optional[Any]				,   UpperCAmelCase_   :    Any							)					->		List[str]:
     for param in model.parameters():
          __lowerCamelCase       :     Dict        =		value





class        UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					)    ->					List[str]:
         super().__init__()
         self.register_modules(
             vae=SCREAMING_SNAKE_CASE_      ,					text_encoder=SCREAMING_SNAKE_CASE_      ,					clip_model=SCREAMING_SNAKE_CASE_      ,					tokenizer=SCREAMING_SNAKE_CASE_      ,					unet=SCREAMING_SNAKE_CASE_      ,					scheduler=SCREAMING_SNAKE_CASE_      ,					feature_extractor=SCREAMING_SNAKE_CASE_      ,					)

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

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







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







    def 				lowercase_    (							self     )    ->					Union[str, Any]:
         self.enable_attention_slicing(SCREAMING_SNAKE_CASE_     )







    def 				lowercase_    (							self     )    ->					Union[str, Any]:
         set_requires_grad(self.vae      ,					SCREAMING_SNAKE_CASE_     )







    def 				lowercase_    (							self     )    ->					List[str]:
         set_requires_grad(self.vae      ,					SCREAMING_SNAKE_CASE_     )







    def 				lowercase_    (							self     )    ->					Optional[int]:
         set_requires_grad(self.unet      ,					SCREAMING_SNAKE_CASE_     )







    def 				lowercase_    (							self     )    ->					Union[str, Any]:
         set_requires_grad(self.unet      ,					SCREAMING_SNAKE_CASE_     )







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

         __lowerCamelCase       :     List[str]        =		self.scheduler.scale_model_input(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     )

         # predict the noise residual
         __lowerCamelCase       :     Tuple        =		self.unet(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					encoder_hidden_states=SCREAMING_SNAKE_CASE_     ).sample

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

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

         __lowerCamelCase       :     Optional[Any]        =		1 / self.vae.config.scaling_factor * sample
         __lowerCamelCase       :     Optional[Any]        =		self.vae.decode(SCREAMING_SNAKE_CASE_     ).sample
         __lowerCamelCase       :     List[Any]        =		(image / 2 + 0.5).clamp(0      ,					1     )

         if use_cutouts:
              __lowerCamelCase       :     Union[str, Any]        =		self.make_cutouts(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     )
         else:
              __lowerCamelCase       :     Any        =		transforms.Resize(self.cut_out_size     )(SCREAMING_SNAKE_CASE_     )
         __lowerCamelCase       :     Any        =		self.normalize(SCREAMING_SNAKE_CASE_     ).to(latents.dtype     )

         __lowerCamelCase       :     Optional[Any]        =		self.clip_model.get_image_features(SCREAMING_SNAKE_CASE_     )
         __lowerCamelCase       :     Optional[Any]        =		image_embeddings_clip / image_embeddings_clip.norm(p=2      ,					dim=-1      ,					keepdim=SCREAMING_SNAKE_CASE_     )

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

         __lowerCamelCase       :     Union[str, Any]        =		-torch.autograd.grad(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     )[0]

         if isinstance(self.scheduler      ,					SCREAMING_SNAKE_CASE_     ):
              __lowerCamelCase       :     Optional[int]        =		latents.detach() + grads * (sigma**2)
              __lowerCamelCase       :     int        =		noise_pred_original
         else:
              __lowerCamelCase       :     Dict        =		noise_pred_original - torch.sqrt(SCREAMING_SNAKE_CASE_     ) * grads
         return noise_pred, latents







    @torch.no_grad()
    def __call__(							self      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_ = 5_12      ,					SCREAMING_SNAKE_CASE_ = 5_12      ,					SCREAMING_SNAKE_CASE_ = 50      ,					SCREAMING_SNAKE_CASE_ = 7.5      ,					SCREAMING_SNAKE_CASE_ = 1      ,					SCREAMING_SNAKE_CASE_ = 0.0      ,					SCREAMING_SNAKE_CASE_ = 1_00      ,					SCREAMING_SNAKE_CASE_ = None      ,					SCREAMING_SNAKE_CASE_ = 4      ,					SCREAMING_SNAKE_CASE_ = True      ,					SCREAMING_SNAKE_CASE_ = None      ,					SCREAMING_SNAKE_CASE_ = None      ,					SCREAMING_SNAKE_CASE_ = "pil"      ,					SCREAMING_SNAKE_CASE_ = True      ,					)    ->					Tuple:
         if isinstance(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     ):
              __lowerCamelCase       :     List[str]        =		1
         elif isinstance(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     ):
              __lowerCamelCase       :     Tuple        =		len(SCREAMING_SNAKE_CASE_     )
         else:
              raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(SCREAMING_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
         __lowerCamelCase       :     str        =		self.tokenizer(
             SCREAMING_SNAKE_CASE_      ,					padding='max_length'      ,					max_length=self.tokenizer.model_max_length      ,					truncation=SCREAMING_SNAKE_CASE_      ,					return_tensors='pt'      ,					)
         __lowerCamelCase       :     List[str]        =		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(SCREAMING_SNAKE_CASE_      ,					dim=0     )

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

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

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

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

         # check if the scheduler accepts generator
         __lowerCamelCase       :     Optional[int]        =		'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(SCREAMING_SNAKE_CASE_     )     ):
              # expand the latents if we are doing classifier free guidance
              __lowerCamelCase       :     Optional[int]        =		torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents
              __lowerCamelCase       :     List[str]        =		self.scheduler.scale_model_input(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_     )

              # predict the noise residual
              __lowerCamelCase       :     Tuple        =		self.unet(SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					encoder_hidden_states=SCREAMING_SNAKE_CASE_     ).sample

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

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

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

         # scale and decode the image latents with vae
         __lowerCamelCase       :     Any        =		1 / self.vae.config.scaling_factor * latents
         __lowerCamelCase       :     Tuple        =		self.vae.decode(SCREAMING_SNAKE_CASE_     ).sample

         __lowerCamelCase       :     str        =		(image / 2 + 0.5).clamp(0      ,					1     )
         __lowerCamelCase       :     Optional[Any]        =		image.cpu().permute(0      ,					2      ,					3      ,					1     ).numpy()

         if output_type == "pil":
              __lowerCamelCase       :     Optional[Any]        =		self.numpy_to_pil(SCREAMING_SNAKE_CASE_     )

         if not return_dict:
              return (image, None)

         return StableDiffusionPipelineOutput(images=SCREAMING_SNAKE_CASE_      ,					nsfw_content_detected=SCREAMING_SNAKE_CASE_     )






