


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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   if isinstance(__lowercase			,       torch.Tensor       ):
      return image
   elif isinstance(__lowercase			,       PIL.Image.Image       ):
      UpperCAmelCase_							:      List[Any]						=   [image]

   if isinstance(image[0]			,       PIL.Image.Image       ):
      UpperCAmelCase_							:      Tuple						=   [np.array(i.resize((w, h)			,       resample=PIL_INTERPOLATION['''lanczos''']       )       )[None, :] for i in image]
      UpperCAmelCase_							:      List[str]						=   np.concatenate(__lowercase			,       axis=0       )
      UpperCAmelCase_							:      Dict						=   np.array(__lowercase       ).astype(np.floataa       ) / 2_5_5.0
      UpperCAmelCase_							:      Optional[int]						=   image.transpose(0			,       3			,       1			,       2       )
      UpperCAmelCase_							:      Optional[Any]						=   2.0 * image - 1.0
      UpperCAmelCase_							:      int						=   torch.from_numpy(__lowercase       )
   elif isinstance(image[0]			,       torch.Tensor       ):
      UpperCAmelCase_							:      Optional[int]						=   torch.cat(__lowercase			,       dim=0       )
   return image
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase			,       __lowercase=0.9_9_9_5       ):
   if not isinstance(__lowercase			,       np.ndarray       ):
      UpperCAmelCase_							:      Tuple						=   True
      UpperCAmelCase_							:      Union[str, Any]						=   va.device
      UpperCAmelCase_							:      int						=   va.cpu().numpy()
      UpperCAmelCase_							:      Optional[int]						=   va.cpu().numpy()

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

   if inputs_are_torch:
      UpperCAmelCase_							:      Dict						=   torch.from_numpy(__lowercase       ).to(__lowercase       )

   return va
def      snake_case_						( __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      str						=   F.normalize(__lowercase			,       dim=-1       )
   UpperCAmelCase_							:      Optional[int]						=   F.normalize(__lowercase			,       dim=-1       )
   return (x - y).norm(dim=-1       ).div(2       ).arcsin().pow(2       ).mul(2       )
def      snake_case_						( __lowercase			,       __lowercase       ):
   for param in model.parameters():
      UpperCAmelCase_							:      List[Any]						=   value





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    def __init__(	self				:   List[Any]       ,   __snake_case				:   AutoencoderKL       ,   __snake_case				:   CLIPTextModel       ,   __snake_case				:   CLIPModel       ,   __snake_case				:   CLIPTokenizer       ,   __snake_case				:   UNetaDConditionModel       ,   __snake_case				:   Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]       ,   __snake_case				:   CLIPFeatureExtractor       ,   __snake_case				:   int=None       ,   __snake_case				:   str=None       ,   __snake_case				:   Any=None       ,   ):





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





    def 					_lowerCamelCase (	self				:   Union[str, Any]       ,   __snake_case				:   Optional[Union[str, int]] = "auto"	):





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





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       self.enable_attention_slicing(__snake_case	)





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       set_requires_grad(self.vae       ,   __snake_case	)





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       set_requires_grad(self.vae       ,   __snake_case	)





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       set_requires_grad(self.unet       ,   __snake_case	)





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       set_requires_grad(self.unet       ,   __snake_case	)





    def 					_lowerCamelCase (	self				:   List[str]       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   str	):





       '''simple docstring'''
       # get the original timestep using init_timestep
       UpperCAmelCase_							:      Optional[int]						=   min(int(num_inference_steps * strength	)       ,   __snake_case	)

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

       return timesteps, num_inference_steps - t_start





    def 					_lowerCamelCase (	self				:   Dict       ,   __snake_case				:   str       ,   __snake_case				:   str       ,   __snake_case				:   Optional[Any]       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   Any       ,   __snake_case				:   Optional[Any]=None	):





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

       UpperCAmelCase_							:      str						=   image.to(device=__snake_case       ,   dtype=__snake_case	)

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

       # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
       UpperCAmelCase_							:      str						=   0.18_215 * init_latents
       UpperCAmelCase_							:      Tuple						=   init_latents.repeat_interleave(__snake_case       ,   dim=0	)

       UpperCAmelCase_							:      Optional[int]						=   randn_tensor(init_latents.shape       ,   generator=__snake_case       ,   device=__snake_case       ,   dtype=__snake_case	)

       # get latents
       UpperCAmelCase_							:      Union[str, Any]						=   self.scheduler.add_noise(__snake_case       ,   __snake_case       ,   __snake_case	)
       UpperCAmelCase_							:      List[Any]						=   init_latents

       return latents





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   int	):





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





    def 					_lowerCamelCase (	self				:   Optional[int]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   Any	):





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





    @torch.enable_grad()
    def 					_lowerCamelCase (	self				:   int       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   Dict       ,   __snake_case				:   Dict       ,   __snake_case				:   Optional[int]       ,   __snake_case				:   int       ,   __snake_case				:   str       ,   __snake_case				:   List[str]       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   latents.detach().requires_grad_()

       UpperCAmelCase_							:      List[str]						=   self.scheduler.scale_model_input(__snake_case       ,   __snake_case	)

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

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

          UpperCAmelCase_							:      Optional[int]						=   torch.sqrt(__snake_case	)
          UpperCAmelCase_							:      List[str]						=   pred_original_sample * (fac) + latents * (1 - fac)
       elif isinstance(self.scheduler       ,   __snake_case	):
          UpperCAmelCase_							:      Optional[int]						=   self.scheduler.sigmas[index]
          UpperCAmelCase_							:      Tuple						=   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_							:      Union[str, Any]						=   1 / 0.18_215 * sample
       UpperCAmelCase_							:      Optional[Any]						=   self.vae.decode(__snake_case	).sample
       UpperCAmelCase_							:      List[str]						=   (image / 2 + 0.5).clamp(0       ,   1	)

       UpperCAmelCase_							:      Optional[Any]						=   transforms.Resize(self.feature_extractor_size	)(__snake_case	)
       UpperCAmelCase_							:      int						=   self.normalize(__snake_case	).to(latents.dtype	)

       UpperCAmelCase_							:      Optional[int]						=   self.clip_model.get_image_features(__snake_case	)
       UpperCAmelCase_							:      str						=   image_embeddings_clip / image_embeddings_clip.norm(p=2       ,   dim=-1       ,   keepdim=__snake_case	)

       UpperCAmelCase_							:      Optional[int]						=   spherical_dist_loss(__snake_case       ,   __snake_case	).mean() * clip_guidance_scale

       UpperCAmelCase_							:      Dict						=   -torch.autograd.grad(__snake_case       ,   __snake_case	)[0]

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





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





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

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

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

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

       UpperCAmelCase_							:      List[Any]						=   slerp(__snake_case       ,   __snake_case       ,   __snake_case	)

       # duplicate text embeddings for each generation per prompt
       UpperCAmelCase_							:      List[str]						=   text_embeddings.repeat_interleave(__snake_case       ,   dim=0	)

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

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

       # Preprocess image
       UpperCAmelCase_							:      Optional[Any]						=   preprocess(__snake_case       ,   __snake_case       ,   __snake_case	)
       UpperCAmelCase_							:      int						=   self.prepare_latents(
           __snake_case       ,   __snake_case       ,   __snake_case       ,   text_embeddings.dtype       ,   self.device       ,   __snake_case	)

       UpperCAmelCase_							:      Union[str, Any]						=   preprocess(__snake_case       ,   __snake_case       ,   __snake_case	)
       UpperCAmelCase_							:      Tuple						=   self.prepare_latents(
           __snake_case       ,   __snake_case       ,   __snake_case       ,   text_embeddings.dtype       ,   self.device       ,   __snake_case	)

       UpperCAmelCase_							:      List[str]						=   slerp(__snake_case       ,   __snake_case       ,   __snake_case	)

       if clip_guidance_scale > 0:
          UpperCAmelCase_							:      Any						=   self.get_clip_image_embeddings(__snake_case       ,   __snake_case	)
          UpperCAmelCase_							:      Tuple						=   self.get_clip_image_embeddings(__snake_case       ,   __snake_case	)
          UpperCAmelCase_							:      Union[str, Any]						=   slerp(
              __snake_case       ,   __snake_case       ,   __snake_case	)

       # 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_							:      int						=   content_text_input.input_ids.shape[-1]
          UpperCAmelCase_							:      Optional[Any]						=   self.tokenizer(['''''']       ,   padding='''max_length'''       ,   max_length=__snake_case       ,   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_							:      int						=   uncond_embeddings.repeat_interleave(__snake_case       ,   dim=0	)

          # For classifier free guidance, we need to do two forward passes.
          # Here we concatenate the unconditional and text embeddings into a single batch
          # to avoid doing two forward passes
          UpperCAmelCase_							:      List[str]						=   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_							:      Tuple						=   (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_							:      int						=   torch.randn(__snake_case       ,   generator=__snake_case       ,   device='''cpu'''       ,   dtype=__snake_case	).to(
                 self.device	)
          else:
             UpperCAmelCase_							:      Dict						=   torch.randn(__snake_case       ,   generator=__snake_case       ,   device=self.device       ,   dtype=__snake_case	)
       else:
          if latents.shape != latents_shape:
             raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''	)
          UpperCAmelCase_							:      Dict						=   latents.to(self.device	)

       # scale the initial noise by the standard deviation required by the scheduler
       UpperCAmelCase_							:      str						=   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_							:      List[str]						=   {}
       if accepts_eta:
          UpperCAmelCase_							:      Tuple						=   eta

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

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

             # predict the noise residual
             UpperCAmelCase_							:      Optional[Any]						=   self.unet(__snake_case       ,   __snake_case       ,   encoder_hidden_states=__snake_case	).sample

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

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

             # compute the previous noisy sample x_t -> x_t-1
             UpperCAmelCase_							:      Union[str, Any]						=   self.scheduler.step(__snake_case       ,   __snake_case       ,   __snake_case       ,   **__snake_case	).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
       UpperCAmelCase_							:      Union[str, Any]						=   1 / 0.18_215 * latents
       UpperCAmelCase_							:      Optional[Any]						=   self.vae.decode(__snake_case	).sample

       UpperCAmelCase_							:      Tuple						=   (image / 2 + 0.5).clamp(0       ,   1	)
       UpperCAmelCase_							:      Optional[int]						=   image.cpu().permute(0       ,   2       ,   3       ,   1	).numpy()

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

       if not return_dict:
          return (image, None)

       return StableDiffusionPipelineOutput(images=__snake_case       ,   nsfw_content_detected=__snake_case	)



# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__UpperCamelCase				:  List[Any]      					=							{
    'configuration_vivit': ['VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'VivitConfig'],
}
try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase				:  Optional[Any]      					=							['VivitImageProcessor']


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase				:  List[Any]      					=							[
           'VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
           'VivitModel',
           'VivitPreTrainedModel',
           'VivitForVideoClassification',
       ]


if TYPE_CHECKING:
       from .configuration_vivit import VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP, VivitConfig

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .image_processing_vivit import VivitImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_vivit import (
                  VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  VivitForVideoClassification,
                  VivitModel,
                  VivitPreTrainedModel,
              )


else:
       import sys

       __UpperCamelCase				:  str      					=							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)