




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_							(				_lowercase						:     List[str]      ,      _lowercase						:     List[Any]      ,      _lowercase						:     Dict							) ->			str:
      if isinstance(lowerCAmelCase_      ,      torch.Tensor							):
            return image
      elif isinstance(lowerCAmelCase_      ,      PIL.Image.Image							):
            SCREAMING_SNAKE_CASE_     						=      [image]

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

def  UpperCAmelCase_							(				_lowercase						:     str      ,      _lowercase						:     List[str]      ,      _lowercase						:     Tuple      ,      _lowercase						:     Optional[int]=0.9_9_9_5							) ->			Optional[int]:
      if not isinstance(lowerCAmelCase_      ,      np.ndarray							):
            SCREAMING_SNAKE_CASE_     						=      True
            SCREAMING_SNAKE_CASE_     						=      va.device
            SCREAMING_SNAKE_CASE_     						=      va.cpu().numpy()
            SCREAMING_SNAKE_CASE_     						=      va.cpu().numpy()

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

      if inputs_are_torch:
            SCREAMING_SNAKE_CASE_     						=      torch.from_numpy(lowerCAmelCase_							).to(lowerCAmelCase_							)

      return va

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

def  UpperCAmelCase_							(				_lowercase						:     int      ,      _lowercase						:     Optional[Any]							) ->			Union[str, Any]:
      for param in model.parameters():
            SCREAMING_SNAKE_CASE_     						=      value





class 	lowerCamelCase_							(				_UpperCamelCase ):


 '''simple docstring'''







 def __init__(							self  :		Tuple    ,    _lowerCAmelCase  :		AutoencoderKL    ,    _lowerCAmelCase  :		CLIPTextModel    ,    _lowerCAmelCase  :		CLIPModel    ,    _lowerCAmelCase  :		CLIPTokenizer    ,    _lowerCAmelCase  :		UNetaDConditionModel    ,    _lowerCAmelCase  :		Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]    ,    _lowerCAmelCase  :		CLIPFeatureExtractor    ,    _lowerCAmelCase  :		List[str]=None    ,    _lowerCAmelCase  :		str=None    ,    _lowerCAmelCase  :		List[Any]=None    ,    ):
       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    ,    )
       SCREAMING_SNAKE_CASE_     						=      (
           feature_extractor.size
           if isinstance(feature_extractor.size    ,    _UpperCAmelCase						)
           else feature_extractor.size['shortest_edge']
       )
       SCREAMING_SNAKE_CASE_     						=      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 				lowerCAmelCase_  (							self  :		Tuple    ,    _lowerCAmelCase  :		Optional[Union[str, int]] = "auto"						):
       if slice_size == "auto":
             # half the attention head size is usually a good trade-off between
             # speed and memory
             SCREAMING_SNAKE_CASE_     						=      self.unet.config.attention_head_dim // 2
       self.unet.set_attention_slice(_UpperCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		int						):
       self.enable_attention_slicing(_UpperCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       set_requires_grad(self.vae    ,    _UpperCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Any						):
       set_requires_grad(self.vae    ,    _UpperCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		List[Any]						):
       set_requires_grad(self.unet    ,    _UpperCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       set_requires_grad(self.unet    ,    _UpperCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Any    ,    _lowerCAmelCase  :		Union[str, Any]    ,    _lowerCAmelCase  :		Any    ,    _lowerCAmelCase  :		Dict						):
       # get the original timestep using init_timestep
       SCREAMING_SNAKE_CASE_     						=      min(int(num_inference_steps * strength						)    ,    _UpperCAmelCase						)

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

       return timesteps, num_inference_steps - t_start



 def 				lowerCAmelCase_  (							self  :		Optional[int]    ,    _lowerCAmelCase  :		Any    ,    _lowerCAmelCase  :		Dict    ,    _lowerCAmelCase  :		int    ,    _lowerCAmelCase  :		Optional[Any]    ,    _lowerCAmelCase  :		Optional[Any]    ,    _lowerCAmelCase  :		Optional[int]=None						):
       if not isinstance(_UpperCAmelCase    ,    torch.Tensor						):
             raise ValueError(F"`image` has to be of type `torch.Tensor` but is {type(_UpperCAmelCase						)}"						)

       SCREAMING_SNAKE_CASE_     						=      image.to(device=_UpperCAmelCase    ,    dtype=_UpperCAmelCase						)

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

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

       SCREAMING_SNAKE_CASE_     						=      randn_tensor(init_latents.shape    ,    generator=_UpperCAmelCase    ,    device=_UpperCAmelCase    ,    dtype=_UpperCAmelCase						)

       # get latents
       SCREAMING_SNAKE_CASE_     						=      self.scheduler.add_noise(_UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase						)
       SCREAMING_SNAKE_CASE_     						=      init_latents

       return latents



 def 				lowerCAmelCase_  (							self  :		int    ,    _lowerCAmelCase  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_     						=      self.coca_transform(_UpperCAmelCase						).unsqueeze(0						)
       with torch.no_grad(), torch.cuda.amp.autocast():
             SCREAMING_SNAKE_CASE_     						=      self.coca_model.generate(transformed_image.to(device=self.device    ,    dtype=self.coca_model.dtype						)						)
       SCREAMING_SNAKE_CASE_     						=      self.coca_tokenizer.decode(generated[0].cpu().numpy()						)
       return generated.split('<end_of_text>'						)[0].replace('<start_of_text>'    ,    ''						).rstrip(' .,'						)



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



 @torch.enable_grad()
 def 				lowerCAmelCase_  (							self  :		Optional[Any]    ,    _lowerCAmelCase  :		int    ,    _lowerCAmelCase  :		List[str]    ,    _lowerCAmelCase  :		Dict    ,    _lowerCAmelCase  :		Optional[int]    ,    _lowerCAmelCase  :		int    ,    _lowerCAmelCase  :		Optional[Any]    ,    _lowerCAmelCase  :		Any    ,    ):
       SCREAMING_SNAKE_CASE_     						=      latents.detach().requires_grad_()

       SCREAMING_SNAKE_CASE_     						=      self.scheduler.scale_model_input(_UpperCAmelCase    ,    _UpperCAmelCase						)

       # predict the noise residual
       SCREAMING_SNAKE_CASE_     						=      self.unet(_UpperCAmelCase    ,    _UpperCAmelCase    ,    encoder_hidden_states=_UpperCAmelCase						).sample

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

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

       SCREAMING_SNAKE_CASE_     						=      transforms.Resize(self.feature_extractor_size						)(_UpperCAmelCase						)
       SCREAMING_SNAKE_CASE_     						=      self.normalize(_UpperCAmelCase						).to(latents.dtype						)

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

       SCREAMING_SNAKE_CASE_     						=      spherical_dist_loss(_UpperCAmelCase    ,    _UpperCAmelCase						).mean() * clip_guidance_scale

       SCREAMING_SNAKE_CASE_     						=      -torch.autograd.grad(_UpperCAmelCase    ,    _UpperCAmelCase						)[0]

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




 @torch.no_grad()
 def __call__(							self  :		Union[str, Any]    ,    _lowerCAmelCase  :		Union[torch.FloatTensor, PIL.Image.Image]    ,    _lowerCAmelCase  :		Union[torch.FloatTensor, PIL.Image.Image]    ,    _lowerCAmelCase  :		Optional[str] = None    ,    _lowerCAmelCase  :		Optional[str] = None    ,    _lowerCAmelCase  :		Optional[int] = 512    ,    _lowerCAmelCase  :		Optional[int] = 512    ,    _lowerCAmelCase  :		float = 0.6    ,    _lowerCAmelCase  :		Optional[int] = 50    ,    _lowerCAmelCase  :		Optional[float] = 7.5    ,    _lowerCAmelCase  :		Optional[int] = 1    ,    _lowerCAmelCase  :		float = 0.0    ,    _lowerCAmelCase  :		Optional[float] = 100    ,    _lowerCAmelCase  :		Optional[torch.Generator] = None    ,    _lowerCAmelCase  :		Optional[str] = "pil"    ,    _lowerCAmelCase  :		bool = True    ,    _lowerCAmelCase  :		float = 0.8    ,    _lowerCAmelCase  :		float = 0.1    ,    _lowerCAmelCase  :		float = 0.1    ,    ):
       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:
             SCREAMING_SNAKE_CASE_     						=      [generator] + [None] * (batch_size - 1)

       SCREAMING_SNAKE_CASE_     						=      [
           ('model', self.coca_model is None),
           ('tokenizer', self.coca_tokenizer is None),
           ('transform', self.coca_transform is None),
       ]
       SCREAMING_SNAKE_CASE_     						=      [x[0] for x in coca_is_none if x[1]]
       SCREAMING_SNAKE_CASE_     						=      ', '.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."						)
             SCREAMING_SNAKE_CASE_     						=      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."						)
             SCREAMING_SNAKE_CASE_     						=      self.get_image_description(_UpperCAmelCase						)

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

       SCREAMING_SNAKE_CASE_     						=      self.tokenizer(
           _UpperCAmelCase    ,    padding='max_length'    ,    max_length=self.tokenizer.model_max_length    ,    truncation=_UpperCAmelCase    ,    return_tensors='pt'    ,    )
       SCREAMING_SNAKE_CASE_     						=      self.text_encoder(style_text_input.input_ids.to(self.device						)						)[0]

       SCREAMING_SNAKE_CASE_     						=      slerp(_UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase						)

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

       # set timesteps
       SCREAMING_SNAKE_CASE_     						=      'offset' in set(inspect.signature(self.scheduler.set_timesteps						).parameters.keys()						)
       SCREAMING_SNAKE_CASE_     						=      {}
       if accepts_offset:
             SCREAMING_SNAKE_CASE_     						=      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						)

       SCREAMING_SNAKE_CASE_     						=      self.get_timesteps(_UpperCAmelCase    ,    _UpperCAmelCase    ,    self.device						)
       SCREAMING_SNAKE_CASE_     						=      timesteps[:1].repeat(_UpperCAmelCase						)

       # Preprocess image
       SCREAMING_SNAKE_CASE_     						=      preprocess(_UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase						)
       SCREAMING_SNAKE_CASE_     						=      self.prepare_latents(
           _UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase    ,    text_embeddings.dtype    ,    self.device    ,    _UpperCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      preprocess(_UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase						)
       SCREAMING_SNAKE_CASE_     						=      self.prepare_latents(
           _UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase    ,    text_embeddings.dtype    ,    self.device    ,    _UpperCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      slerp(_UpperCAmelCase    ,    _UpperCAmelCase    ,    _UpperCAmelCase						)

       if clip_guidance_scale > 0:
             SCREAMING_SNAKE_CASE_     						=      self.get_clip_image_embeddings(_UpperCAmelCase    ,    _UpperCAmelCase						)
             SCREAMING_SNAKE_CASE_     						=      self.get_clip_image_embeddings(_UpperCAmelCase    ,    _UpperCAmelCase						)
             SCREAMING_SNAKE_CASE_     						=      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.
       SCREAMING_SNAKE_CASE_     						=      guidance_scale > 1.0
       # get unconditional embeddings for classifier free guidance
       if do_classifier_free_guidance:
             SCREAMING_SNAKE_CASE_     						=      content_text_input.input_ids.shape[-1]
             SCREAMING_SNAKE_CASE_     						=      self.tokenizer(['']    ,    padding='max_length'    ,    max_length=_UpperCAmelCase    ,    return_tensors='pt'						)
             SCREAMING_SNAKE_CASE_     						=      self.text_encoder(uncond_input.input_ids.to(self.device						)						)[0]
             # duplicate unconditional embeddings for each generation per prompt
             SCREAMING_SNAKE_CASE_     						=      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
             SCREAMING_SNAKE_CASE_     						=      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`.
       SCREAMING_SNAKE_CASE_     						=      (batch_size, self.unet.config.in_channels, height // 8, width // 8)
       SCREAMING_SNAKE_CASE_     						=      text_embeddings.dtype
       if latents is None:
             if self.device.type == "mps":
                   # randn does not work reproducibly on mps
                   SCREAMING_SNAKE_CASE_     						=      torch.randn(_UpperCAmelCase    ,    generator=_UpperCAmelCase    ,    device='cpu'    ,    dtype=_UpperCAmelCase						).to(
                       self.device						)
             else:
                   SCREAMING_SNAKE_CASE_     						=      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}"						)
             SCREAMING_SNAKE_CASE_     						=      latents.to(self.device						)

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

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

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

                   # predict the noise residual
                   SCREAMING_SNAKE_CASE_     						=      self.unet(_UpperCAmelCase    ,    _UpperCAmelCase    ,    encoder_hidden_states=_UpperCAmelCase						).sample

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

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

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

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

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

       if output_type == "pil":
             SCREAMING_SNAKE_CASE_     						=      self.numpy_to_pil(_UpperCAmelCase						)

       if not return_dict:
             return (image, None)

       return StableDiffusionPipelineOutput(images=_UpperCAmelCase    ,    nsfw_content_detected=_UpperCAmelCase						)





from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase__			:       Dict         =							logging.get_logger(__name__)

lowerCamelCase__			:       Any         =							{
    'facebook/s2t-small-librispeech-asr': (
        'https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/config.json'
    ),
    # See all Speech2Text models at https://huggingface.co/models?filter=speech_to_text
}





class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE ):


 '''simple docstring'''







 lowercase_       =   "speech_to_text"
 lowercase_       =   ["past_key_values"]
 lowercase_       =   {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}



 def __init__(							self  :		Union[str, Any]    ,    _lowerCAmelCase  :		Optional[int]=10_000    ,    _lowerCAmelCase  :		str=12    ,    _lowerCAmelCase  :		Tuple=2_048    ,    _lowerCAmelCase  :		str=4    ,    _lowerCAmelCase  :		Tuple=6    ,    _lowerCAmelCase  :		Optional[int]=2_048    ,    _lowerCAmelCase  :		Dict=4    ,    _lowerCAmelCase  :		List[Any]=0.0    ,    _lowerCAmelCase  :		int=0.0    ,    _lowerCAmelCase  :		Optional[int]=True    ,    _lowerCAmelCase  :		Union[str, Any]=True    ,    _lowerCAmelCase  :		Any="relu"    ,    _lowerCAmelCase  :		Any=256    ,    _lowerCAmelCase  :		List[str]=0.1    ,    _lowerCAmelCase  :		Optional[Any]=0.0    ,    _lowerCAmelCase  :		Optional[int]=0.0    ,    _lowerCAmelCase  :		Tuple=0.02    ,    _lowerCAmelCase  :		List[str]=2    ,    _lowerCAmelCase  :		Tuple=True    ,    _lowerCAmelCase  :		List[str]=1    ,    _lowerCAmelCase  :		Dict=0    ,    _lowerCAmelCase  :		Tuple=2    ,    _lowerCAmelCase  :		Optional[int]=6_000    ,    _lowerCAmelCase  :		Tuple=1_024    ,    _lowerCAmelCase  :		Any=2    ,    _lowerCAmelCase  :		str=(5, 5)    ,    _lowerCAmelCase  :		Optional[int]=1_024    ,    _lowerCAmelCase  :		List[Any]=80    ,    _lowerCAmelCase  :		List[Any]=1    ,    **_lowerCAmelCase  :		List[Any]    ,    ):
       SCREAMING_SNAKE_CASE_     						=      vocab_size
       SCREAMING_SNAKE_CASE_     						=      d_model
       SCREAMING_SNAKE_CASE_     						=      encoder_ffn_dim
       SCREAMING_SNAKE_CASE_     						=      encoder_layers
       SCREAMING_SNAKE_CASE_     						=      encoder_attention_heads
       SCREAMING_SNAKE_CASE_     						=      decoder_ffn_dim
       SCREAMING_SNAKE_CASE_     						=      decoder_layers
       SCREAMING_SNAKE_CASE_     						=      decoder_attention_heads
       SCREAMING_SNAKE_CASE_     						=      dropout
       SCREAMING_SNAKE_CASE_     						=      attention_dropout
       SCREAMING_SNAKE_CASE_     						=      activation_dropout
       SCREAMING_SNAKE_CASE_     						=      activation_function
       SCREAMING_SNAKE_CASE_     						=      init_std
       SCREAMING_SNAKE_CASE_     						=      encoder_layerdrop
       SCREAMING_SNAKE_CASE_     						=      decoder_layerdrop
       SCREAMING_SNAKE_CASE_     						=      use_cache
       SCREAMING_SNAKE_CASE_     						=      encoder_layers
       SCREAMING_SNAKE_CASE_     						=      scale_embedding  # scale factor will be sqrt(d_model) if True
       SCREAMING_SNAKE_CASE_     						=      max_source_positions
       SCREAMING_SNAKE_CASE_     						=      max_target_positions
       SCREAMING_SNAKE_CASE_     						=      num_conv_layers
       SCREAMING_SNAKE_CASE_     						=      list(_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_     						=      conv_channels
       SCREAMING_SNAKE_CASE_     						=      input_feat_per_channel
       SCREAMING_SNAKE_CASE_     						=      input_channels

       if len(self.conv_kernel_sizes						) != self.num_conv_layers:
             raise ValueError(
                 'Configuration for convolutional module is incorrect. '
                 'It is required that `len(config.conv_kernel_sizes)` == `config.num_conv_layers` '
                 F"but is `len(config.conv_kernel_sizes) = {len(self.conv_kernel_sizes						)}`, "
                 F"`config.num_conv_layers = {self.num_conv_layers}`."						)

       super().__init__(
           pad_token_id=_lowerCAmelCase    ,    bos_token_id=_lowerCAmelCase    ,    eos_token_id=_lowerCAmelCase    ,    is_encoder_decoder=_lowerCAmelCase    ,    decoder_start_token_id=_lowerCAmelCase    ,    **_lowerCAmelCase    ,    )