
"""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 	lowerCamelCase__			(							UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_				)->							Any:







       """simple docstring"""
       if isinstance(UpperCAmelCase_		,      torch.Tensor				):
              return image
       elif isinstance(UpperCAmelCase_		,      PIL.Image.Image				):
              UpperCamelCase					 =		[image]

       if isinstance(image[0]		,      PIL.Image.Image				):
              UpperCamelCase					 =		[np.array(i.resize((w, h)		,      resample=PIL_INTERPOLATION["lanczos"]				)				)[None, :] for i in image]
              UpperCamelCase					 =		np.concatenate(UpperCAmelCase_		,      axis=0				)
              UpperCamelCase					 =		np.array(UpperCAmelCase_				).astype(np.floataa				) / 255.0
              UpperCamelCase					 =		image.transpose(0		,      3		,      1		,      2				)
              UpperCamelCase					 =		2.0 * image - 1.0
              UpperCamelCase					 =		torch.from_numpy(UpperCAmelCase_				)
       elif isinstance(image[0]		,      torch.Tensor				):
              UpperCamelCase					 =		torch.cat(UpperCAmelCase_		,      dim=0				)
       return image



def 	lowerCamelCase__			(							UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_=0.9995				)->							Tuple:







       """simple docstring"""
       if not isinstance(UpperCAmelCase_		,      np.ndarray				):
              UpperCamelCase					 =		True
              UpperCamelCase					 =		va.device
              UpperCamelCase					 =		va.cpu().numpy()
              UpperCamelCase					 =		va.cpu().numpy()

       UpperCamelCase					 =		np.sum(va * va / (np.linalg.norm(UpperCAmelCase_				) * np.linalg.norm(UpperCAmelCase_				))				)
       if np.abs(UpperCAmelCase_				) > DOT_THRESHOLD:
              UpperCamelCase					 =		(1 - t) * va + t * va
       else:
              UpperCamelCase					 =		np.arccos(UpperCAmelCase_				)
              UpperCamelCase					 =		np.sin(UpperCAmelCase_				)
              UpperCamelCase					 =		theta_a * t
              UpperCamelCase					 =		np.sin(UpperCAmelCase_				)
              UpperCamelCase					 =		np.sin(theta_a - theta_t				) / sin_theta_a
              UpperCamelCase					 =		sin_theta_t / sin_theta_a
              UpperCamelCase					 =		sa * va + sa * va

       if inputs_are_torch:
              UpperCamelCase					 =		torch.from_numpy(UpperCAmelCase_				).to(UpperCAmelCase_				)

       return va



def 	lowerCamelCase__			(							UpperCAmelCase_		,      UpperCAmelCase_				)->							Any:







       """simple docstring"""
       UpperCamelCase					 =		F.normalize(UpperCAmelCase_		,      dim=-1				)
       UpperCamelCase					 =		F.normalize(UpperCAmelCase_		,      dim=-1				)
       return (x - y).norm(dim=-1				).div(2				).arcsin().pow(2				).mul(2				)



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







       """simple docstring"""
       for param in model.parameters():
              UpperCamelCase					 =		value



class   __a  (      _lowerCAmelCase	):






       def __init__(  self :							str       ,      UpperCAmelCase_ :							AutoencoderKL       ,      UpperCAmelCase_ :							CLIPTextModel       ,      UpperCAmelCase_ :							CLIPModel       ,      UpperCAmelCase_ :							CLIPTokenizer       ,      UpperCAmelCase_ :							UNetaDConditionModel       ,      UpperCAmelCase_ :							Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]       ,      UpperCAmelCase_ :							CLIPFeatureExtractor       ,      UpperCAmelCase_ :							Tuple=None       ,      UpperCAmelCase_ :							List[Any]=None       ,      UpperCAmelCase_ :							Dict=None       ,      )->  Union[str, Any]:







              """simple docstring"""
              super().__init__()
              self.register_modules(
                  vae=UpperCAmelCase_       ,      text_encoder=UpperCAmelCase_       ,      clip_model=UpperCAmelCase_       ,      tokenizer=UpperCAmelCase_       ,      unet=UpperCAmelCase_       ,      scheduler=UpperCAmelCase_       ,      feature_extractor=UpperCAmelCase_       ,      coca_model=UpperCAmelCase_       ,      coca_tokenizer=UpperCAmelCase_       ,      coca_transform=UpperCAmelCase_       ,      )
              UpperCamelCase					 =		(
                  feature_extractor.size
                  if isinstance(feature_extractor.size       ,      UpperCAmelCase_			)
                  else feature_extractor.size["shortest_edge"]
              )
              UpperCamelCase					 =		transforms.Normalize(mean=feature_extractor.image_mean       ,      std=feature_extractor.image_std			)
              set_requires_grad(self.text_encoder       ,      UpperCAmelCase_			)
              set_requires_grad(self.clip_model       ,      UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[str]       ,      UpperCAmelCase_ :							Optional[Union[str, int]] = "auto"			)->  Tuple:







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






       def 	_SCREAMING_SNAKE_CASE       (  self :							Dict			)->  str:







              """simple docstring"""
              self.enable_attention_slicing(UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							Any			)->  Optional[int]:







              """simple docstring"""
              set_requires_grad(self.vae       ,      UpperCAmelCase_			)






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







              """simple docstring"""
              set_requires_grad(self.vae       ,      UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[Any]			)->  Optional[Any]:







              """simple docstring"""
              set_requires_grad(self.unet       ,      UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							int			)->  List[Any]:







              """simple docstring"""
              set_requires_grad(self.unet       ,      UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[Any]       ,      UpperCAmelCase_ :							Any       ,      UpperCAmelCase_ :							Optional[int]       ,      UpperCAmelCase_ :							Tuple			)->  int:







              """simple docstring"""
              # get the original timestep using init_timestep
              UpperCamelCase					 =		min(int(num_inference_steps * strength			)       ,      UpperCAmelCase_			)

              UpperCamelCase					 =		max(num_inference_steps - init_timestep       ,      0			)
              UpperCamelCase					 =		self.scheduler.timesteps[t_start:]

              return timesteps, num_inference_steps - t_start






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[int]       ,      UpperCAmelCase_ :							str       ,      UpperCAmelCase_ :							Any       ,      UpperCAmelCase_ :							Optional[int]       ,      UpperCAmelCase_ :							List[Any]       ,      UpperCAmelCase_ :							Union[str, Any]       ,      UpperCAmelCase_ :							Dict=None			)->  List[str]:







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

              UpperCamelCase					 =		image.to(device=UpperCAmelCase_       ,      dtype=UpperCAmelCase_			)

              if isinstance(UpperCAmelCase_       ,      UpperCAmelCase_			):
                     UpperCamelCase					 =		[
                         self.vae.encode(image[i : i + 1]			).latent_dist.sample(generator[i]			) for i in range(UpperCAmelCase_			)
                     ]
                     UpperCamelCase					 =		torch.cat(UpperCAmelCase_       ,      dim=0			)
              else:
                     UpperCamelCase					 =		self.vae.encode(UpperCAmelCase_			).latent_dist.sample(UpperCAmelCase_			)

              # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
              UpperCamelCase					 =		0.18215 * init_latents
              UpperCamelCase					 =		init_latents.repeat_interleave(UpperCAmelCase_       ,      dim=0			)

              UpperCamelCase					 =		randn_tensor(init_latents.shape       ,      generator=UpperCAmelCase_       ,      device=UpperCAmelCase_       ,      dtype=UpperCAmelCase_			)

              # get latents
              UpperCamelCase					 =		self.scheduler.add_noise(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)
              UpperCamelCase					 =		init_latents

              return latents






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[int]       ,      UpperCAmelCase_ :							Union[str, Any]			)->  Optional[int]:







              """simple docstring"""
              UpperCamelCase					 =		self.coca_transform(UpperCAmelCase_			).unsqueeze(0			)
              with torch.no_grad(), torch.cuda.amp.autocast():
                     UpperCamelCase					 =		self.coca_model.generate(transformed_image.to(device=self.device       ,      dtype=self.coca_model.dtype			)			)
              UpperCamelCase					 =		self.coca_tokenizer.decode(generated[0].cpu().numpy()			)
              return generated.split("<end_of_text>"			)[0].replace("<start_of_text>"       ,      ""			).rstrip(" .,"			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							Union[str, Any]       ,      UpperCAmelCase_ :							Tuple       ,      UpperCAmelCase_ :							Any			)->  Union[str, Any]:







              """simple docstring"""
              UpperCamelCase					 =		self.feature_extractor.preprocess(UpperCAmelCase_			)
              UpperCamelCase					 =		torch.from_numpy(clip_image_input["pixel_values"][0]			).unsqueeze(0			).to(self.device			).half()
              UpperCamelCase					 =		self.clip_model.get_image_features(UpperCAmelCase_			)
              UpperCamelCase					 =		image_embeddings_clip / image_embeddings_clip.norm(p=2       ,      dim=-1       ,      keepdim=UpperCAmelCase_			)
              UpperCamelCase					 =		image_embeddings_clip.repeat_interleave(UpperCAmelCase_       ,      dim=0			)
              return image_embeddings_clip






       @torch.enable_grad()
       def 	_SCREAMING_SNAKE_CASE       (  self :							Union[str, Any]       ,      UpperCAmelCase_ :							Tuple       ,      UpperCAmelCase_ :							Dict       ,      UpperCAmelCase_ :							int       ,      UpperCAmelCase_ :							str       ,      UpperCAmelCase_ :							Optional[Any]       ,      UpperCAmelCase_ :							Tuple       ,      UpperCAmelCase_ :							List[str]       ,      )->  Any:







              """simple docstring"""
              UpperCamelCase					 =		latents.detach().requires_grad_()

              UpperCamelCase					 =		self.scheduler.scale_model_input(UpperCAmelCase_       ,      UpperCAmelCase_			)

              # predict the noise residual
              UpperCamelCase					 =		self.unet(UpperCAmelCase_       ,      UpperCAmelCase_       ,      encoder_hidden_states=UpperCAmelCase_			).sample

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

                     UpperCamelCase					 =		torch.sqrt(UpperCAmelCase_			)
                     UpperCamelCase					 =		pred_original_sample * (fac) + latents * (1 - fac)
              elif isinstance(self.scheduler       ,      UpperCAmelCase_			):
                     UpperCamelCase					 =		self.scheduler.sigmas[index]
                     UpperCamelCase					 =		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					 =		1 / 0.18215 * sample
              UpperCamelCase					 =		self.vae.decode(UpperCAmelCase_			).sample
              UpperCamelCase					 =		(image / 2 + 0.5).clamp(0       ,      1			)

              UpperCamelCase					 =		transforms.Resize(self.feature_extractor_size			)(UpperCAmelCase_			)
              UpperCamelCase					 =		self.normalize(UpperCAmelCase_			).to(latents.dtype			)

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

              UpperCamelCase					 =		spherical_dist_loss(UpperCAmelCase_       ,      UpperCAmelCase_			).mean() * clip_guidance_scale

              UpperCamelCase					 =		-torch.autograd.grad(UpperCAmelCase_       ,      UpperCAmelCase_			)[0]

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






       @torch.no_grad()
       def __call__(  self :							int       ,      UpperCAmelCase_ :							Union[torch.FloatTensor, PIL.Image.Image]       ,      UpperCAmelCase_ :							Union[torch.FloatTensor, PIL.Image.Image]       ,      UpperCAmelCase_ :							Optional[str] = None       ,      UpperCAmelCase_ :							Optional[str] = None       ,      UpperCAmelCase_ :							Optional[int] = 512       ,      UpperCAmelCase_ :							Optional[int] = 512       ,      UpperCAmelCase_ :							float = 0.6       ,      UpperCAmelCase_ :							Optional[int] = 50       ,      UpperCAmelCase_ :							Optional[float] = 7.5       ,      UpperCAmelCase_ :							Optional[int] = 1       ,      UpperCAmelCase_ :							float = 0.0       ,      UpperCAmelCase_ :							Optional[float] = 100       ,      UpperCAmelCase_ :							Optional[torch.Generator] = None       ,      UpperCAmelCase_ :							Optional[str] = "pil"       ,      UpperCAmelCase_ :							bool = True       ,      UpperCAmelCase_ :							float = 0.8       ,      UpperCAmelCase_ :							float = 0.1       ,      UpperCAmelCase_ :							float = 0.1       ,      )->  Union[str, Any]:







              """simple docstring"""
              if isinstance(UpperCAmelCase_       ,      UpperCAmelCase_			) and len(UpperCAmelCase_			) != batch_size:
                     raise ValueError(f"You have passed {batch_size} batch_size, but only {len(UpperCAmelCase_			)} 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(UpperCAmelCase_       ,      torch.Generator			) and batch_size > 1:
                     UpperCamelCase					 =		[generator] + [None] * (batch_size - 1)

              UpperCamelCase					 =		[
                  ("model", self.coca_model is None),
                  ("tokenizer", self.coca_tokenizer is None),
                  ("transform", self.coca_transform is None),
              ]
              UpperCamelCase					 =		[x[0] for x in coca_is_none if x[1]]
              UpperCamelCase					 =		", ".join(UpperCAmelCase_			)
              # generate prompts with coca model if prompt is None
              if content_prompt is None:
                     if len(UpperCAmelCase_			):
                            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					 =		self.get_image_description(UpperCAmelCase_			)
              if style_prompt is None:
                     if len(UpperCAmelCase_			):
                            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					 =		self.get_image_description(UpperCAmelCase_			)

              # get prompt text embeddings for content and style
              UpperCamelCase					 =		self.tokenizer(
                  UpperCAmelCase_       ,      padding="max_length"       ,      max_length=self.tokenizer.model_max_length       ,      truncation=UpperCAmelCase_       ,      return_tensors="pt"       ,      )
              UpperCamelCase					 =		self.text_encoder(content_text_input.input_ids.to(self.device			)			)[0]

              UpperCamelCase					 =		self.tokenizer(
                  UpperCAmelCase_       ,      padding="max_length"       ,      max_length=self.tokenizer.model_max_length       ,      truncation=UpperCAmelCase_       ,      return_tensors="pt"       ,      )
              UpperCamelCase					 =		self.text_encoder(style_text_input.input_ids.to(self.device			)			)[0]

              UpperCamelCase					 =		slerp(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)

              # duplicate text embeddings for each generation per prompt
              UpperCamelCase					 =		text_embeddings.repeat_interleave(UpperCAmelCase_       ,      dim=0			)

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

              self.scheduler.set_timesteps(UpperCAmelCase_       ,      **UpperCAmelCase_			)
              # 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					 =		self.get_timesteps(UpperCAmelCase_       ,      UpperCAmelCase_       ,      self.device			)
              UpperCamelCase					 =		timesteps[:1].repeat(UpperCAmelCase_			)

              # Preprocess image
              UpperCamelCase					 =		preprocess(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)
              UpperCamelCase					 =		self.prepare_latents(
                  UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      text_embeddings.dtype       ,      self.device       ,      UpperCAmelCase_			)

              UpperCamelCase					 =		preprocess(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)
              UpperCamelCase					 =		self.prepare_latents(
                  UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      text_embeddings.dtype       ,      self.device       ,      UpperCAmelCase_			)

              UpperCamelCase					 =		slerp(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)

              if clip_guidance_scale > 0:
                     UpperCamelCase					 =		self.get_clip_image_embeddings(UpperCAmelCase_       ,      UpperCAmelCase_			)
                     UpperCamelCase					 =		self.get_clip_image_embeddings(UpperCAmelCase_       ,      UpperCAmelCase_			)
                     UpperCamelCase					 =		slerp(
                         UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_			)

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

              # scale the initial noise by the standard deviation required by the scheduler
              UpperCamelCase					 =		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					 =		"eta" in set(inspect.signature(self.scheduler.step			).parameters.keys()			)
              UpperCamelCase					 =		{}
              if accepts_eta:
                     UpperCamelCase					 =		eta

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

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

                            # predict the noise residual
                            UpperCamelCase					 =		self.unet(UpperCAmelCase_       ,      UpperCAmelCase_       ,      encoder_hidden_states=UpperCAmelCase_			).sample

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

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

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

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
              UpperCamelCase					 =		1 / 0.18215 * latents
              UpperCamelCase					 =		self.vae.decode(UpperCAmelCase_			).sample

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

              if output_type == "pil":
                     UpperCamelCase					 =		self.numpy_to_pil(UpperCAmelCase_			)

              if not return_dict:
                     return (image, None)

              return StableDiffusionPipelineOutput(images=UpperCAmelCase_       ,      nsfw_content_detected=UpperCAmelCase_			)







"""simple docstring"""


def 	lowerCamelCase__			(							UpperCAmelCase_		,      UpperCAmelCase_ = False				)->							bool:







       """simple docstring"""
       if n == 2:
              return True
       if not n % 2 or n < 2:
              return False
       if n > 5 and n % 10 not in (1, 3, 7, 9):  # can quickly check last digit
              return False
       if n > 3_31_70_44_06_46_79_88_73_85_96_19_81 and not allow_probable:
              raise ValueError(
                  "Warning: upper bound of deterministic test is exceeded. "
                  "Pass allow_probable=True to allow probabilistic test. "
                  "A return value of True indicates a probable prime."				)
       # array bounds provided by analysis
       UpperCamelCase					 =		[
           20_47,
           1_37_36_53,
           25_32_60_01,
           32_15_03_17_51,
           2_15_23_02_89_87_47,
           3_47_47_49_66_03_83,
           3_41_55_00_71_72_83_21,
           1,
           3_82_51_23_05_65_46_41_30_51,
           1,
           1,
           31_86_65_85_78_34_03_11_51_16_74_61,
           3_31_70_44_06_46_79_88_73_85_96_19_81,
       ]

       UpperCamelCase					 =		[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
       for idx, _p in enumerate(UpperCAmelCase_		,      1				):
              if n < _p:
                     # then we have our last prime to check
                     UpperCamelCase					 =		primes[:idx]
                     break
       UpperCamelCase ,       UpperCamelCase					 =		n - 1, 0
       # break up n -1 into a power of 2 (s) and
       # remaining odd component
       # essentially, solve for d * 2 ** s == n - 1
       while d % 2 == 0:
              d //= 2
              s += 1
       for prime in plist:
              UpperCamelCase					 =		False
              for r in range(UpperCAmelCase_				):
                     UpperCamelCase					 =		pow(UpperCAmelCase_		,      d * 2**r		,      UpperCAmelCase_				)
                     # see article for analysis explanation for m
                     if (r == 0 and m == 1) or ((m + 1) % n == 0):
                            UpperCamelCase					 =		True
                            # this loop will not determine compositeness
                            break
              if pr:
                     continue
              # if pr is False, then the above loop never evaluated to true,
              # and the n MUST be composite
              return False
       return True



def 	lowerCamelCase__			(							)->							None:







       """simple docstring"""
       assert not miller_rabin(5_61				)
       assert miller_rabin(5_63				)
       # 2047

       assert not miller_rabin(83_82_01				)
       assert miller_rabin(83_82_07				)
       # 1_373_653

       assert not miller_rabin(17_31_60_01				)
       assert miller_rabin(17_31_60_17				)
       # 25_326_001

       assert not miller_rabin(30_78_38_66_41				)
       assert miller_rabin(30_78_38_66_53				)
       # 3_215_031_751

       assert not miller_rabin(1_71_30_45_57_48_01				)
       assert miller_rabin(1_71_30_45_57_48_19				)
       # 2_152_302_898_747

       assert not miller_rabin(2_77_97_99_72_83_07				)
       assert miller_rabin(2_77_97_99_72_83_27				)
       # 3_474_749_660_383

       assert not miller_rabin(1_13_85_00_23_90_94_41				)
       assert miller_rabin(1_13_85_00_23_90_95_27				)
       # 341_550_071_728_321

       assert not miller_rabin(1_27_50_41_01_88_48_80_43_51				)
       assert miller_rabin(1_27_50_41_01_88_48_80_43_91				)
       # 3_825_123_056_546_413_051

       assert not miller_rabin(7_96_66_46_44_58_50_77_87_79_18_67				)
       assert miller_rabin(7_96_66_46_44_58_50_77_87_79_19_51				)
       # 318_665_857_834_031_151_167_461

       assert not miller_rabin(55_28_40_67_74_46_64_78_97_66_03_33				)
       assert miller_rabin(55_28_40_67_74_46_64_78_97_66_03_59				)
       # 3_317_044_064_679_887_385_961_981
       # upper limit for probabilistic test


if __name__ == "__main__":
      test_miller_rabin()





