'''simple docstring'''

import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

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







def __UpperCamelCase( _A       :					Optional[int]	,      _A       :					str	,      _A       :					Optional[int]     ):
   '''simple docstring'''
   if isinstance(_A	,      torch.Tensor     ):
      return image
   elif isinstance(_A	,      PIL.Image.Image     ):
      UpperCAmelCase__				:					str 					=	[image]

   if isinstance(image[0]	,      PIL.Image.Image     ):
      UpperCAmelCase__				:					str 					=	[np.array(i.resize((w, h)	,      resample=PIL_INTERPOLATION['''lanczos''']     )     )[None, :] for i in image]
      UpperCAmelCase__				:					Dict 					=	np.concatenate(_A	,      axis=0     )
      UpperCAmelCase__				:					Dict 					=	np.array(_A     ).astype(np.floataa     ) / 2_5_5.0
      UpperCAmelCase__				:					List[Any] 					=	image.transpose(0	,      3	,      1	,      2     )
      UpperCAmelCase__				:					Tuple 					=	2.0 * image - 1.0
      UpperCAmelCase__				:					Any 					=	torch.from_numpy(_A     )
   elif isinstance(image[0]	,      torch.Tensor     ):
      UpperCAmelCase__				:					Tuple 					=	torch.cat(_A	,      dim=0     )
   return image







def __UpperCamelCase( _A       :					Any	,      _A       :					Dict	,      _A       :					Optional[Any]	,      _A       :					Union[str, Any]=0.9_9_9_5     ):
   '''simple docstring'''
   if not isinstance(_A	,      np.ndarray     ):
      UpperCAmelCase__				:					List[Any] 					=	True
      UpperCAmelCase__				:					List[str] 					=	va.device
      UpperCAmelCase__				:					Union[str, Any] 					=	va.cpu().numpy()
      UpperCAmelCase__				:					List[Any] 					=	va.cpu().numpy()

   UpperCAmelCase__				:					Optional[int] 					=	np.sum(va * va / (np.linalg.norm(_A     ) * np.linalg.norm(_A     ))     )
   if np.abs(_A     ) > DOT_THRESHOLD:
      UpperCAmelCase__				:					Any 					=	(1 - t) * va + t * va
   else:
      UpperCAmelCase__				:					Union[str, Any] 					=	np.arccos(_A     )
      UpperCAmelCase__				:					str 					=	np.sin(_A     )
      UpperCAmelCase__				:					Tuple 					=	theta_a * t
      UpperCAmelCase__				:					str 					=	np.sin(_A     )
      UpperCAmelCase__				:					str 					=	np.sin(theta_a - theta_t     ) / sin_theta_a
      UpperCAmelCase__				:					Dict 					=	sin_theta_t / sin_theta_a
      UpperCAmelCase__				:					Optional[int] 					=	sa * va + sa * va

   if inputs_are_torch:
      UpperCAmelCase__				:					Dict 					=	torch.from_numpy(_A     ).to(_A     )

   return va







def __UpperCamelCase( _A       :					int	,      _A       :					Optional[int]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					Optional[int] 					=	F.normalize(_A	,      dim=-1     )
   UpperCAmelCase__				:					Union[str, Any] 					=	F.normalize(_A	,      dim=-1     )
   return (x - y).norm(dim=-1     ).div(2     ).arcsin().pow(2     ).mul(2     )







def __UpperCamelCase( _A       :					Any	,      _A       :					List[str]     ):
   '''simple docstring'''
   for param in model.parameters():
      UpperCAmelCase__				:					Tuple 					=	value


class 					_lowercase						(   lowerCAmelCase ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,)      ->      Tuple:



         '''simple docstring'''



         super().__init__()
         self.register_modules(
             vae=lowerCamelCase_       ,text_encoder=lowerCamelCase_       ,clip_model=lowerCamelCase_       ,tokenizer=lowerCamelCase_       ,unet=lowerCamelCase_       ,scheduler=lowerCamelCase_       ,feature_extractor=lowerCamelCase_       ,coca_model=lowerCamelCase_       ,coca_tokenizer=lowerCamelCase_       ,coca_transform=lowerCamelCase_       ,)
         UpperCAmelCase__				:					Tuple 					=	(
             feature_extractor.size
             if isinstance(feature_extractor.size       ,lowerCamelCase_				)
             else feature_extractor.size['''shortest_edge''']
         )
         UpperCAmelCase__				:					List[Any] 					=	transforms.Normalize(mean=feature_extractor.image_mean       ,std=feature_extractor.image_std				)
         set_requires_grad(self.text_encoder       ,lowerCamelCase_				)
         set_requires_grad(self.clip_model       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_ = "auto"				)      ->      Optional[Any]:



         '''simple docstring'''



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




      def 		lowerCAmelCase__       (					self				)      ->      List[Any]:



         '''simple docstring'''



         self.enable_attention_slicing(lowerCamelCase_				)




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



         '''simple docstring'''



         set_requires_grad(self.vae       ,lowerCamelCase_				)




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



         '''simple docstring'''



         set_requires_grad(self.vae       ,lowerCamelCase_				)




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



         '''simple docstring'''



         set_requires_grad(self.unet       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self				)      ->      Any:



         '''simple docstring'''



         set_requires_grad(self.unet       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	min(int(num_inference_steps * strength				)       ,lowerCamelCase_				)

         UpperCAmelCase__				:					Any 					=	max(num_inference_steps - init_timestep       ,0				)
         UpperCAmelCase__				:					List[str] 					=	self.scheduler.timesteps[t_start:]

         return timesteps, num_inference_steps - t_start




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_=None				)      ->      Optional[int]:



         '''simple docstring'''



         if not isinstance(lowerCamelCase_       ,torch.Tensor				):
            raise ValueError(f'''`image` has to be of type `torch.Tensor` but is {type(lowerCamelCase_				)}'''				)

         UpperCAmelCase__				:					List[Any] 					=	image.to(device=lowerCamelCase_       ,dtype=lowerCamelCase_				)

         if isinstance(lowerCamelCase_       ,lowerCamelCase_				):
            UpperCAmelCase__				:					Optional[int] 					=	[
                self.vae.encode(image[i : i + 1]				).latent_dist.sample(generator[i]				) for i in range(lowerCamelCase_				)
            ]
            UpperCAmelCase__				:					List[Any] 					=	torch.cat(lowerCamelCase_       ,dim=0				)
         else:
            UpperCAmelCase__				:					Any 					=	self.vae.encode(lowerCamelCase_				).latent_dist.sample(lowerCamelCase_				)

         # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
         UpperCAmelCase__				:					Union[str, Any] 					=	0.1_8215 * init_latents
         UpperCAmelCase__				:					Optional[Any] 					=	init_latents.repeat_interleave(lowerCamelCase_       ,dim=0				)

         UpperCAmelCase__				:					str 					=	randn_tensor(init_latents.shape       ,generator=lowerCamelCase_       ,device=lowerCamelCase_       ,dtype=lowerCamelCase_				)

         # get latents
         UpperCAmelCase__				:					List[str] 					=	self.scheduler.add_noise(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	init_latents

         return latents




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	self.coca_transform(lowerCamelCase_				).unsqueeze(0				)
         with torch.no_grad(), torch.cuda.amp.autocast():
            UpperCAmelCase__				:					Optional[int] 					=	self.coca_model.generate(transformed_image.to(device=self.device       ,dtype=self.coca_model.dtype				)				)
         UpperCAmelCase__				:					List[str] 					=	self.coca_tokenizer.decode(generated[0].cpu().numpy()				)
         return generated.split('''<end_of_text>'''				)[0].replace('''<start_of_text>'''       ,''''''				).rstrip(''' .,'''				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					List[str] 					=	self.feature_extractor.preprocess(lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	torch.from_numpy(clip_image_input['''pixel_values'''][0]				).unsqueeze(0				).to(self.device				).half()
         UpperCAmelCase__				:					Any 					=	self.clip_model.get_image_features(lowerCamelCase_				)
         UpperCAmelCase__				:					Any 					=	image_embeddings_clip / image_embeddings_clip.norm(p=2       ,dim=-1       ,keepdim=lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	image_embeddings_clip.repeat_interleave(lowerCamelCase_       ,dim=0				)
         return image_embeddings_clip




      @torch.enable_grad()
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,)      ->      Any:



         '''simple docstring'''



         UpperCAmelCase__				:					int 					=	latents.detach().requires_grad_()

         UpperCAmelCase__				:					Union[str, Any] 					=	self.scheduler.scale_model_input(lowerCamelCase_       ,lowerCamelCase_				)

         # predict the noise residual
         UpperCAmelCase__				:					Optional[int] 					=	self.unet(lowerCamelCase_       ,lowerCamelCase_       ,encoder_hidden_states=lowerCamelCase_				).sample

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

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

         # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
         UpperCAmelCase__				:					int 					=	1 / 0.1_8215 * sample
         UpperCAmelCase__				:					Any 					=	self.vae.decode(lowerCamelCase_				).sample
         UpperCAmelCase__				:					Tuple 					=	(image / 2 + 0.5).clamp(0       ,1				)

         UpperCAmelCase__				:					List[str] 					=	transforms.Resize(self.feature_extractor_size				)(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	self.normalize(lowerCamelCase_				).to(latents.dtype				)

         UpperCAmelCase__				:					int 					=	self.clip_model.get_image_features(lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	image_embeddings_clip / image_embeddings_clip.norm(p=2       ,dim=-1       ,keepdim=lowerCamelCase_				)

         UpperCAmelCase__				:					List[Any] 					=	spherical_dist_loss(lowerCamelCase_       ,lowerCamelCase_				).mean() * clip_guidance_scale

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

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




      @torch.no_grad()
      def __call__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_ = None       ,lowerCamelCase_ = None       ,lowerCamelCase_ = 512       ,lowerCamelCase_ = 512       ,lowerCamelCase_ = 0.6       ,lowerCamelCase_ = 50       ,lowerCamelCase_ = 7.5       ,lowerCamelCase_ = 1       ,lowerCamelCase_ = 0.0       ,lowerCamelCase_ = 100       ,lowerCamelCase_ = None       ,lowerCamelCase_ = "pil"       ,lowerCamelCase_ = True       ,lowerCamelCase_ = 0.8       ,lowerCamelCase_ = 0.1       ,lowerCamelCase_ = 0.1       ,)      ->      Dict:



         '''simple docstring'''



         if isinstance(lowerCamelCase_       ,lowerCamelCase_				) and len(lowerCamelCase_				) != batch_size:
            raise ValueError(f'''You have passed {batch_size} batch_size, but only {len(lowerCamelCase_				)} generators.'''				)

         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}.'''				)

         if isinstance(lowerCamelCase_       ,torch.Generator				) and batch_size > 1:
            UpperCAmelCase__				:					Optional[Any] 					=	[generator] + [None] * (batch_size - 1)

         UpperCAmelCase__				:					int 					=	[
             ('''model''', self.coca_model is None),
             ('''tokenizer''', self.coca_tokenizer is None),
             ('''transform''', self.coca_transform is None),
         ]
         UpperCAmelCase__				:					List[str] 					=	[x[0] for x in coca_is_none if x[1]]
         UpperCAmelCase__				:					List[str] 					=	''', '''.join(lowerCamelCase_				)
         # generate prompts with coca model if prompt is None
         if content_prompt is None:
            if len(lowerCamelCase_				):
               raise ValueError(
                   f'''Content prompt is None and CoCa [{coca_is_none_str}] is None.'''
                   f'''Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''				)
            UpperCAmelCase__				:					Dict 					=	self.get_image_description(lowerCamelCase_				)
         if style_prompt is None:
            if len(lowerCamelCase_				):
               raise ValueError(
                   f'''Style prompt is None and CoCa [{coca_is_none_str}] is None.'''
                   f''' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''				)
            UpperCAmelCase__				:					int 					=	self.get_image_description(lowerCamelCase_				)

         # get prompt text embeddings for content and style
         UpperCAmelCase__				:					Union[str, Any] 					=	self.tokenizer(
             lowerCamelCase_       ,padding='''max_length'''       ,max_length=self.tokenizer.model_max_length       ,truncation=lowerCamelCase_       ,return_tensors='''pt'''       ,)
         UpperCAmelCase__				:					Union[str, Any] 					=	self.text_encoder(content_text_input.input_ids.to(self.device				)				)[0]

         UpperCAmelCase__				:					Optional[int] 					=	self.tokenizer(
             lowerCamelCase_       ,padding='''max_length'''       ,max_length=self.tokenizer.model_max_length       ,truncation=lowerCamelCase_       ,return_tensors='''pt'''       ,)
         UpperCAmelCase__				:					Optional[Any] 					=	self.text_encoder(style_text_input.input_ids.to(self.device				)				)[0]

         UpperCAmelCase__				:					str 					=	slerp(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)

         # duplicate text embeddings for each generation per prompt
         UpperCAmelCase__				:					Union[str, Any] 					=	text_embeddings.repeat_interleave(lowerCamelCase_       ,dim=0				)

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

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

         UpperCAmelCase__					,   UpperCAmelCase__				:					List[str] 					=	self.get_timesteps(lowerCamelCase_       ,lowerCamelCase_       ,self.device				)
         UpperCAmelCase__				:					Optional[int] 					=	timesteps[:1].repeat(lowerCamelCase_				)

         # Preprocess image
         UpperCAmelCase__				:					Any 					=	preprocess(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	self.prepare_latents(
             lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,text_embeddings.dtype       ,self.device       ,lowerCamelCase_				)

         UpperCAmelCase__				:					List[Any] 					=	preprocess(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)
         UpperCAmelCase__				:					Union[str, Any] 					=	self.prepare_latents(
             lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,text_embeddings.dtype       ,self.device       ,lowerCamelCase_				)

         UpperCAmelCase__				:					List[str] 					=	slerp(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)

         if clip_guidance_scale > 0:
            UpperCAmelCase__				:					List[str] 					=	self.get_clip_image_embeddings(lowerCamelCase_       ,lowerCamelCase_				)
            UpperCAmelCase__				:					Any 					=	self.get_clip_image_embeddings(lowerCamelCase_       ,lowerCamelCase_				)
            UpperCAmelCase__				:					Any 					=	slerp(
                lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)

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

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

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

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

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

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

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

               # compute the previous noisy sample x_t -> x_t-1
               UpperCAmelCase__				:					Dict 					=	self.scheduler.step(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,**lowerCamelCase_				).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
         UpperCAmelCase__				:					Dict 					=	1 / 0.1_8215 * latents
         UpperCAmelCase__				:					Tuple 					=	self.vae.decode(lowerCamelCase_				).sample

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

         if output_type == "pil":
            UpperCAmelCase__				:					Dict 					=	self.numpy_to_pil(lowerCamelCase_				)

         if not return_dict:
            return (image, None)

         return StableDiffusionPipelineOutput(images=lowerCamelCase_       ,nsfw_content_detected=lowerCamelCase_				)





'''simple docstring'''

import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin


class 					_lowercase						(   lowerCAmelCase		,lowerCAmelCase ):
      '''simple docstring'''


      @register_to_config
      def __init__(					self       ,lowerCamelCase_ = 128       ,lowerCamelCase_ = 256       ,lowerCamelCase_ = 2000.0       ,lowerCamelCase_ = 768       ,lowerCamelCase_ = 12       ,lowerCamelCase_ = 12       ,lowerCamelCase_ = 64       ,lowerCamelCase_ = 2048       ,lowerCamelCase_ = 0.1       ,)      ->      List[Any]:



         '''simple docstring'''



         super().__init__()

         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Sequential(
             nn.Linear(lowerCamelCase_       ,d_model * 4       ,bias=lowerCamelCase_				)       ,nn.SiLU()       ,nn.Linear(d_model * 4       ,d_model * 4       ,bias=lowerCamelCase_				)       ,nn.SiLU()       ,)

         UpperCAmelCase__				:					List[str] 					=	nn.Embedding(lowerCamelCase_       ,lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	False

         UpperCAmelCase__				:					List[str] 					=	nn.Linear(lowerCamelCase_       ,lowerCamelCase_       ,bias=lowerCamelCase_				)

         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Dropout(p=lowerCamelCase_				)

         UpperCAmelCase__				:					int 					=	nn.ModuleList()
         for lyr_num in range(lowerCamelCase_				):
            # FiLM conditional T5 decoder
            UpperCAmelCase__				:					str 					=	DecoderLayer(d_model=lowerCamelCase_       ,d_kv=lowerCamelCase_       ,num_heads=lowerCamelCase_       ,d_ff=lowerCamelCase_       ,dropout_rate=lowerCamelCase_				)
            self.decoders.append(lowerCamelCase_				)

         UpperCAmelCase__				:					Tuple 					=	TaLayerNorm(lowerCamelCase_				)

         UpperCAmelCase__				:					Any 					=	nn.Dropout(p=lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	nn.Linear(lowerCamelCase_       ,lowerCamelCase_       ,bias=lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Dict:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	torch.mul(query_input.unsqueeze(-1				)       ,key_input.unsqueeze(-2				)				)
         return mask.unsqueeze(-3				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__					,   UpperCAmelCase__					,   UpperCAmelCase__				:					Union[str, Any] 					=	decoder_input_tokens.shape
         assert decoder_noise_time.shape == (batch,)

         # decoder_noise_time is in [0, 1), so rescale to expected timing range.
         UpperCAmelCase__				:					str 					=	get_timestep_embedding(
             decoder_noise_time * self.config.max_decoder_noise_time       ,embedding_dim=self.config.d_model       ,max_period=self.config.max_decoder_noise_time       ,).to(dtype=self.dtype				)

         UpperCAmelCase__				:					List[Any] 					=	self.conditioning_emb(lowerCamelCase_				).unsqueeze(1				)

         assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

         UpperCAmelCase__				:					List[Any] 					=	decoder_input_tokens.shape[1]

         # If we want to use relative positions for audio context, we can just offset
         # this sequence by the length of encodings_and_masks.
         UpperCAmelCase__				:					int 					=	torch.broadcast_to(
             torch.arange(lowerCamelCase_       ,device=decoder_input_tokens.device				)       ,(batch, seq_length)       ,)

         UpperCAmelCase__				:					Optional[Any] 					=	self.position_encoding(lowerCamelCase_				)

         UpperCAmelCase__				:					Union[str, Any] 					=	self.continuous_inputs_projection(lowerCamelCase_				)
         inputs += position_encodings
         UpperCAmelCase__				:					str 					=	self.dropout(lowerCamelCase_				)

         # decoder: No padding present.
         UpperCAmelCase__				:					Optional[Any] 					=	torch.ones(
             decoder_input_tokens.shape[:2]       ,device=decoder_input_tokens.device       ,dtype=inputs.dtype				)

         # Translate encoding masks to encoder-decoder masks.
         UpperCAmelCase__				:					Dict 					=	[(x, self.encoder_decoder_mask(lowerCamelCase_       ,lowerCamelCase_				)) for x, y in encodings_and_masks]

         # cross attend style: concat encodings
         UpperCAmelCase__				:					str 					=	torch.cat([x[0] for x in encodings_and_encdec_masks]       ,dim=1				)
         UpperCAmelCase__				:					Dict 					=	torch.cat([x[1] for x in encodings_and_encdec_masks]       ,dim=-1				)

         for lyr in self.decoders:
            UpperCAmelCase__				:					Optional[Any] 					=	lyr(
                lowerCamelCase_       ,conditioning_emb=lowerCamelCase_       ,encoder_hidden_states=lowerCamelCase_       ,encoder_attention_mask=lowerCamelCase_       ,)[0]

         UpperCAmelCase__				:					List[Any] 					=	self.decoder_norm(lowerCamelCase_				)
         UpperCAmelCase__				:					int 					=	self.post_dropout(lowerCamelCase_				)

         UpperCAmelCase__				:					Union[str, Any] 					=	self.spec_out(lowerCamelCase_				)
         return spec_out




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_=1e-6				)      ->      str:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					int 					=	nn.ModuleList()

         # cond self attention: layer 0
         self.layer.append(
             TaLayerSelfAttentionCond(d_model=lowerCamelCase_       ,d_kv=lowerCamelCase_       ,num_heads=lowerCamelCase_       ,dropout_rate=lowerCamelCase_				)				)

         # cross attention: layer 1
         self.layer.append(
             TaLayerCrossAttention(
                 d_model=lowerCamelCase_       ,d_kv=lowerCamelCase_       ,num_heads=lowerCamelCase_       ,dropout_rate=lowerCamelCase_       ,layer_norm_epsilon=lowerCamelCase_       ,)				)

         # Film Cond MLP + dropout: last layer
         self.layer.append(
             TaLayerFFCond(d_model=lowerCamelCase_       ,d_ff=lowerCamelCase_       ,dropout_rate=lowerCamelCase_       ,layer_norm_epsilon=lowerCamelCase_				)				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	self.layer[0](
             lowerCamelCase_       ,conditioning_emb=lowerCamelCase_       ,attention_mask=lowerCamelCase_       ,)

         if encoder_hidden_states is not None:
            UpperCAmelCase__				:					Any 					=	torch.where(encoder_attention_mask > 0       ,0       ,-1e10				).to(
                encoder_hidden_states.dtype				)

            UpperCAmelCase__				:					Union[str, Any] 					=	self.layer[1](
                lowerCamelCase_       ,key_value_states=lowerCamelCase_       ,attention_mask=lowerCamelCase_       ,)

         # Apply Film Conditional Feed Forward layer
         UpperCAmelCase__				:					Optional[Any] 					=	self.layer[-1](lowerCamelCase_       ,lowerCamelCase_				)

         return (hidden_states,)




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Tuple:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					str 					=	TaLayerNorm(lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	TaFiLMLayer(in_features=d_model * 4       ,out_features=lowerCamelCase_				)
         UpperCAmelCase__				:					Any 					=	Attention(query_dim=lowerCamelCase_       ,heads=lowerCamelCase_       ,dim_head=lowerCamelCase_       ,out_bias=lowerCamelCase_       ,scale_qk=lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	nn.Dropout(lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	self.layer_norm(lowerCamelCase_				)

         if conditioning_emb is not None:
            UpperCAmelCase__				:					Any 					=	self.FiLMLayer(lowerCamelCase_       ,lowerCamelCase_				)

         # Self-attention block
         UpperCAmelCase__				:					List[Any] 					=	self.attention(lowerCamelCase_				)

         UpperCAmelCase__				:					str 					=	hidden_states + self.dropout(lowerCamelCase_				)

         return hidden_states




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Tuple:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					Optional[int] 					=	Attention(query_dim=lowerCamelCase_       ,heads=lowerCamelCase_       ,dim_head=lowerCamelCase_       ,out_bias=lowerCamelCase_       ,scale_qk=lowerCamelCase_				)
         UpperCAmelCase__				:					int 					=	TaLayerNorm(lowerCamelCase_       ,eps=lowerCamelCase_				)
         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Dropout(lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	self.layer_norm(lowerCamelCase_				)
         UpperCAmelCase__				:					List[Any] 					=	self.attention(
             lowerCamelCase_       ,encoder_hidden_states=lowerCamelCase_       ,attention_mask=attention_mask.squeeze(1				)       ,)
         UpperCAmelCase__				:					int 					=	hidden_states + self.dropout(lowerCamelCase_				)
         return layer_output




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Optional[int]:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					Union[str, Any] 					=	TaDenseGatedActDense(d_model=lowerCamelCase_       ,d_ff=lowerCamelCase_       ,dropout_rate=lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	TaFiLMLayer(in_features=d_model * 4       ,out_features=lowerCamelCase_				)
         UpperCAmelCase__				:					int 					=	TaLayerNorm(lowerCamelCase_       ,eps=lowerCamelCase_				)
         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Dropout(lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_=None				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	self.layer_norm(lowerCamelCase_				)
         if conditioning_emb is not None:
            UpperCAmelCase__				:					Optional[Any] 					=	self.film(lowerCamelCase_       ,lowerCamelCase_				)

         UpperCAmelCase__				:					List[str] 					=	self.DenseReluDense(lowerCamelCase_				)
         UpperCAmelCase__				:					Any 					=	hidden_states + self.dropout(lowerCamelCase_				)
         return hidden_states




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      List[str]:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Linear(lowerCamelCase_       ,lowerCamelCase_       ,bias=lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	nn.Linear(lowerCamelCase_       ,lowerCamelCase_       ,bias=lowerCamelCase_				)
         UpperCAmelCase__				:					Tuple 					=	nn.Linear(lowerCamelCase_       ,lowerCamelCase_       ,bias=lowerCamelCase_				)
         UpperCAmelCase__				:					Tuple 					=	nn.Dropout(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	NewGELUActivation()




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	self.act(self.wi_a(lowerCamelCase_				)				)
         UpperCAmelCase__				:					List[str] 					=	self.wi_a(lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	hidden_gelu * hidden_linear
         UpperCAmelCase__				:					int 					=	self.dropout(lowerCamelCase_				)

         UpperCAmelCase__				:					Optional[Any] 					=	self.wo(lowerCamelCase_				)
         return hidden_states




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_=1e-6				)      ->      int:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					Tuple 					=	nn.Parameter(torch.ones(lowerCamelCase_				)				)
         UpperCAmelCase__				:					Dict 					=	eps




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	hidden_states.to(torch.floataa				).pow(2				).mean(-1       ,keepdim=lowerCamelCase_				)
         UpperCAmelCase__				:					List[Any] 					=	hidden_states * torch.rsqrt(variance + self.variance_epsilon				)

         # convert into half-precision if necessary
         if self.weight.dtype in [torch.floataa, torch.bfloataa]:
            UpperCAmelCase__				:					Union[str, Any] 					=	hidden_states.to(self.weight.dtype				)

         return self.weight * hidden_states




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      torch.Tensor:



         '''simple docstring'''



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




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Dict:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					str 					=	nn.Linear(lowerCamelCase_       ,out_features * 2       ,bias=lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Any 					=	self.scale_bias(lowerCamelCase_				)
         UpperCAmelCase__					,   UpperCAmelCase__				:					List[Any] 					=	torch.chunk(lowerCamelCase_       ,2       ,-1				)
         UpperCAmelCase__				:					List[Any] 					=	x * (1 + scale) + shift
         return x




