




"""simple docstring"""
import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class  _SCREAMING_SNAKE_CASE       (     SCREAMING_SNAKE_CASE__				):



 '''simple docstring'''





 __UpperCAmelCase		   =							(HeunDiscreteScheduler,)
 __UpperCAmelCase		   =							10






 def lowercase_      (self					,     **lowerCAmelCase__   ):

  '''simple docstring'''







  _UpperCamelCase      :	Any					      =  {
      "num_train_timesteps": 11_00,
      "beta_start": 0.0001,
      "beta_end": 0.02,
      "beta_schedule": "linear",
  }

  config.update(**lowerCAmelCase__   )
  return config






 def lowercase_      (self   ):

  '''simple docstring'''







  for timesteps in [10, 50, 1_00, 10_00]:
   self.check_over_configs(num_train_timesteps=lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]					,     [0.0002, 0.002, 0.02]   ):
   self.check_over_configs(beta_start=lowerCAmelCase__					,     beta_end=lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  for schedule in ["linear", "scaled_linear", "exp"]:
   self.check_over_configs(beta_schedule=lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  for clip_sample_range in [1.0, 2.0, 3.0]:
   self.check_over_configs(clip_sample_range=lowerCAmelCase__					,     clip_sample=lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  for prediction_type in ["epsilon", "v_prediction", "sample"]:
   self.check_over_configs(prediction_type=lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  self.scheduler_classes[0]
  _UpperCamelCase      :	List[Any]					      =  self.get_scheduler_config()
  _UpperCamelCase      :	List[str]					      =  scheduler_class(**lowerCAmelCase__   )

  scheduler.set_timesteps(self.num_inference_steps   )

  _UpperCamelCase      :	List[str]					      =  self.dummy_model()
  _UpperCamelCase      :	str					      =  self.dummy_sample_deter * scheduler.init_noise_sigma
  _UpperCamelCase      :	int					      =  sample.to(lowerCAmelCase__   )

  for i, t in enumerate(scheduler.timesteps   ):
   _UpperCamelCase      :	Tuple					      =  scheduler.scale_model_input(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	Union[str, Any]					      =  model(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	Union[str, Any]					      =  scheduler.step(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   )
   _UpperCamelCase      :	int					      =  output.prev_sample

  _UpperCamelCase      :	Optional[Any]					      =  torch.sum(torch.abs(lowerCAmelCase__   )   )
  _UpperCamelCase      :	Optional[Any]					      =  torch.mean(torch.abs(lowerCAmelCase__   )   )

  if torch_device in ["cpu", "mps"]:
   assert abs(result_sum.item() - 0.1233   ) < 1E-2
   assert abs(result_mean.item() - 0.0002   ) < 1E-3
  else:
   # CUDA
   assert abs(result_sum.item() - 0.1233   ) < 1E-2
   assert abs(result_mean.item() - 0.0002   ) < 1E-3






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  self.scheduler_classes[0]
  _UpperCamelCase      :	List[Any]					      =  self.get_scheduler_config(prediction_type="v_prediction"   )
  _UpperCamelCase      :	int					      =  scheduler_class(**lowerCAmelCase__   )

  scheduler.set_timesteps(self.num_inference_steps   )

  _UpperCamelCase      :	Any					      =  self.dummy_model()
  _UpperCamelCase      :	int					      =  self.dummy_sample_deter * scheduler.init_noise_sigma
  _UpperCamelCase      :	Tuple					      =  sample.to(lowerCAmelCase__   )

  for i, t in enumerate(scheduler.timesteps   ):
   _UpperCamelCase      :	Dict					      =  scheduler.scale_model_input(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	Tuple					      =  model(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	Optional[Any]					      =  scheduler.step(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   )
   _UpperCamelCase      :	Any					      =  output.prev_sample

  _UpperCamelCase      :	str					      =  torch.sum(torch.abs(lowerCAmelCase__   )   )
  _UpperCamelCase      :	str					      =  torch.mean(torch.abs(lowerCAmelCase__   )   )

  if torch_device in ["cpu", "mps"]:
   assert abs(result_sum.item() - 4.69_34E-07   ) < 1E-2
   assert abs(result_mean.item() - 6.11_12E-10   ) < 1E-3
  else:
   # CUDA
   assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72E-07   ) < 1E-2
   assert abs(result_mean.item() - 0.0002   ) < 1E-3






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Any					      =  self.scheduler_classes[0]
  _UpperCamelCase      :	str					      =  self.get_scheduler_config()
  _UpperCamelCase      :	Any					      =  scheduler_class(**lowerCAmelCase__   )

  scheduler.set_timesteps(self.num_inference_steps					,     device=lowerCAmelCase__   )

  _UpperCamelCase      :	Tuple					      =  self.dummy_model()
  _UpperCamelCase      :	List[str]					      =  self.dummy_sample_deter.to(lowerCAmelCase__   ) * scheduler.init_noise_sigma

  for t in scheduler.timesteps:
   _UpperCamelCase      :	Union[str, Any]					      =  scheduler.scale_model_input(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	List[Any]					      =  model(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	Optional[int]					      =  scheduler.step(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   )
   _UpperCamelCase      :	List[Any]					      =  output.prev_sample

  _UpperCamelCase      :	Dict					      =  torch.sum(torch.abs(lowerCAmelCase__   )   )
  _UpperCamelCase      :	Optional[Any]					      =  torch.mean(torch.abs(lowerCAmelCase__   )   )

  if str(lowerCAmelCase__   ).startswith("cpu"   ):
   # The following sum varies between 148 and 156 on mps. Why?
   assert abs(result_sum.item() - 0.1233   ) < 1E-2
   assert abs(result_mean.item() - 0.0002   ) < 1E-3
  elif str(lowerCAmelCase__   ).startswith("mps"   ):
   # Larger tolerance on mps
   assert abs(result_mean.item() - 0.0002   ) < 1E-2
  else:
   # CUDA
   assert abs(result_sum.item() - 0.1233   ) < 1E-2
   assert abs(result_mean.item() - 0.0002   ) < 1E-3






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Tuple					      =  self.scheduler_classes[0]
  _UpperCamelCase      :	List[Any]					      =  self.get_scheduler_config()
  _UpperCamelCase      :	Union[str, Any]					      =  scheduler_class(**lowerCAmelCase__					,     use_karras_sigmas=lowerCAmelCase__   )

  scheduler.set_timesteps(self.num_inference_steps					,     device=lowerCAmelCase__   )

  _UpperCamelCase      :	List[Any]					      =  self.dummy_model()
  _UpperCamelCase      :	List[str]					      =  self.dummy_sample_deter.to(lowerCAmelCase__   ) * scheduler.init_noise_sigma
  _UpperCamelCase      :	Optional[Any]					      =  sample.to(lowerCAmelCase__   )

  for t in scheduler.timesteps:
   _UpperCamelCase      :	Optional[Any]					      =  scheduler.scale_model_input(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	Optional[Any]					      =  model(lowerCAmelCase__					,     lowerCAmelCase__   )

   _UpperCamelCase      :	List[str]					      =  scheduler.step(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   )
   _UpperCamelCase      :	Optional[Any]					      =  output.prev_sample

  _UpperCamelCase      :	int					      =  torch.sum(torch.abs(lowerCAmelCase__   )   )
  _UpperCamelCase      :	List[Any]					      =  torch.mean(torch.abs(lowerCAmelCase__   )   )

  assert abs(result_sum.item() - 0.0_0015   ) < 1E-2
  assert abs(result_mean.item() - 1.98_69_55_45_35_03_46_95E-07   ) < 1E-2









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



 '''simple docstring'''





 def __init__(self					,     lowerCAmelCase__					,     lowerCAmelCase__=1.0   ):

  '''simple docstring'''







  super().__init__()

  _UpperCamelCase      :	Optional[Any]					      =  cut_size
  _UpperCamelCase      :	List[str]					      =  cut_power






 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__   ):

  '''simple docstring'''







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






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		str  ,	__lowerCAmelCase			:		Optional[int]						)    ->       Union[str, Any]:
 _UpperCamelCase      :	List[str]					      =  F.normalize(__lowerCAmelCase  ,	dim=-1						)
 _UpperCamelCase      :	Optional[Any]					      =  F.normalize(__lowerCAmelCase  ,	dim=-1						)
 return (x - y).norm(dim=-1						).div(2						).arcsin().pow(2						).mul(2						)






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		str						)    ->       List[str]:
 for param in model.parameters():
  _UpperCamelCase      :	Any					      =  value







class  _SCREAMING_SNAKE_CASE       (     SCREAMING_SNAKE_CASE__				):



 '''simple docstring'''





 def __init__(self					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     ):

  '''simple docstring'''







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

  _UpperCamelCase      :	Union[str, Any]					      =  transforms.Normalize(mean=feature_extractor.image_mean					,     std=feature_extractor.image_std   )
  _UpperCamelCase      :	Optional[int]					      =  (
      feature_extractor.size
      if isinstance(feature_extractor.size					,     lowerCAmelCase__   )
      else feature_extractor.size["shortest_edge"]
  )
  _UpperCamelCase      :	Dict					      =  MakeCutouts(self.cut_out_size   )

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






 def lowercase_      (self					,     lowerCAmelCase__ = "auto"   ):

  '''simple docstring'''







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






 def lowercase_      (self   ):

  '''simple docstring'''







  self.enable_attention_slicing(lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.vae					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.vae					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.unet					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.unet					,     lowerCAmelCase__   )






 @torch.enable_grad()
 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__=True					,     ):

  '''simple docstring'''







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

  _UpperCamelCase      :	Optional[Any]					      =  self.scheduler.scale_model_input(lowerCAmelCase__					,     lowerCAmelCase__   )

  # predict the noise residual
  _UpperCamelCase      :	Dict					      =  self.unet(lowerCAmelCase__					,     lowerCAmelCase__					,     encoder_hidden_states=lowerCAmelCase__   ).sample

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

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

  _UpperCamelCase      :	List[Any]					      =  1 / self.vae.config.scaling_factor * sample
  _UpperCamelCase      :	Dict					      =  self.vae.decode(lowerCAmelCase__   ).sample
  _UpperCamelCase      :	Union[str, Any]					      =  (image / 2 + 0.5).clamp(0					,     1   )

  if use_cutouts:
   _UpperCamelCase      :	Dict					      =  self.make_cutouts(lowerCAmelCase__					,     lowerCAmelCase__   )
  else:
   _UpperCamelCase      :	List[str]					      =  transforms.Resize(self.cut_out_size   )(lowerCAmelCase__   )
  _UpperCamelCase      :	str					      =  self.normalize(lowerCAmelCase__   ).to(latents.dtype   )

  _UpperCamelCase      :	List[Any]					      =  self.clip_model.get_image_features(lowerCAmelCase__   )
  _UpperCamelCase      :	Dict					      =  image_embeddings_clip / image_embeddings_clip.norm(p=2					,     dim=-1					,     keepdim=lowerCAmelCase__   )

  if use_cutouts:
   _UpperCamelCase      :	str					      =  spherical_dist_loss(lowerCAmelCase__					,     lowerCAmelCase__   )
   _UpperCamelCase      :	str					      =  dists.view([num_cutouts, sample.shape[0], -1]   )
   _UpperCamelCase      :	Optional[Any]					      =  dists.sum(2   ).mean(0   ).sum() * clip_guidance_scale
  else:
   _UpperCamelCase      :	Any					      =  spherical_dist_loss(lowerCAmelCase__					,     lowerCAmelCase__   ).mean() * clip_guidance_scale

  _UpperCamelCase      :	List[str]					      =  -torch.autograd.grad(lowerCAmelCase__					,     lowerCAmelCase__   )[0]

  if isinstance(self.scheduler					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	Tuple					      =  latents.detach() + grads * (sigma**2)
   _UpperCamelCase      :	int					      =  noise_pred_original
  else:
   _UpperCamelCase      :	Union[str, Any]					      =  noise_pred_original - torch.sqrt(lowerCAmelCase__   ) * grads
  return noise_pred, latents






 @torch.no_grad()
 def __call__(self					,     lowerCAmelCase__					,     lowerCAmelCase__ = 5_12					,     lowerCAmelCase__ = 5_12					,     lowerCAmelCase__ = 50					,     lowerCAmelCase__ = 7.5					,     lowerCAmelCase__ = 1					,     lowerCAmelCase__ = 0.0					,     lowerCAmelCase__ = 1_00					,     lowerCAmelCase__ = None					,     lowerCAmelCase__ = 4					,     lowerCAmelCase__ = True					,     lowerCAmelCase__ = None					,     lowerCAmelCase__ = None					,     lowerCAmelCase__ = "pil"					,     lowerCAmelCase__ = True					,     ):

  '''simple docstring'''







  if isinstance(lowerCAmelCase__					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	str					      =  1
  elif isinstance(lowerCAmelCase__					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	Dict					      =  len(lowerCAmelCase__   )
  else:
   raise ValueError(F"`prompt` has to be of type `str` or `list` but is {type(lowerCAmelCase__   )}"   )

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

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

  # set timesteps
  _UpperCamelCase      :	List[str]					      =  "offset" in set(inspect.signature(self.scheduler.set_timesteps   ).parameters.keys()   )
  _UpperCamelCase      :	str					      =  {}
  if accepts_offset:
   _UpperCamelCase      :	Tuple					      =  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
  _UpperCamelCase      :	str					      =  self.scheduler.timesteps.to(self.device   )

  # scale the initial noise by the standard deviation required by the scheduler
  _UpperCamelCase      :	Tuple					      =  latents * self.scheduler.init_noise_sigma

  # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
  # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
  # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
  # and should be between [0, 1]
  _UpperCamelCase      :	str					      =  "eta" in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
  _UpperCamelCase      :	Optional[int]					      =  {}
  if accepts_eta:
   _UpperCamelCase      :	Union[str, Any]					      =  eta

  # check if the scheduler accepts generator
  _UpperCamelCase      :	Optional[int]					      =  "generator" in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
  if accepts_generator:
   _UpperCamelCase      :	List[str]					      =  generator

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

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

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

   # perform clip guidance
   if clip_guidance_scale > 0:
    _UpperCamelCase      :	List[str]					      =  (
        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__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     )

   # compute the previous noisy sample x_t -> x_t-1
   _UpperCamelCase      :	Tuple					      =  self.scheduler.step(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     **lowerCAmelCase__   ).prev_sample

  # scale and decode the image latents with vae
  _UpperCamelCase      :	Tuple					      =  1 / self.vae.config.scaling_factor * latents
  _UpperCamelCase      :	Union[str, Any]					      =  self.vae.decode(lowerCAmelCase__   ).sample

  _UpperCamelCase      :	List[Any]					      =  (image / 2 + 0.5).clamp(0					,     1   )
  _UpperCamelCase      :	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__   )



