



'''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      _SCREAMING_SNAKE_CASE  (	UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase      ):



    """simple docstring"""





    if isinstance(lowerCAmelCase_							,   torch.Tensor      ):
        return image
    elif isinstance(lowerCAmelCase_							,   PIL.Image.Image      ):
        lowerCAmelCase__			:  Optional[int]											= [image]

    if isinstance(image[0]							,   PIL.Image.Image      ):
        lowerCAmelCase__			:  Tuple											= [np.array(i.resize((w, h)							,   resample=PIL_INTERPOLATION["""lanczos"""]      )      )[None, :] for i in image]
        lowerCAmelCase__			:  Union[str, Any]											= np.concatenate(lowerCAmelCase_							,   axis=0      )
        lowerCAmelCase__			:  List[str]											= np.array(lowerCAmelCase_      ).astype(np.floataa      ) / 255.0
        lowerCAmelCase__			:  Union[str, Any]											= image.transpose(0							,   3							,   1							,   2      )
        lowerCAmelCase__			:  Union[str, Any]											= 2.0 * image - 1.0
        lowerCAmelCase__			:  Dict											= torch.from_numpy(lowerCAmelCase_      )
    elif isinstance(image[0]							,   torch.Tensor      ):
        lowerCAmelCase__			:  Optional[Any]											= torch.cat(lowerCAmelCase_							,   dim=0      )
    return image
def      _SCREAMING_SNAKE_CASE  (	UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase=0.9995      ):



    """simple docstring"""





    if not isinstance(lowerCAmelCase_							,   np.ndarray      ):
        lowerCAmelCase__			:  Optional[Any]											= True
        lowerCAmelCase__			:  Union[str, Any]											= va.device
        lowerCAmelCase__			:  Any											= va.cpu().numpy()
        lowerCAmelCase__			:  Optional[Any]											= va.cpu().numpy()

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

    if inputs_are_torch:
        lowerCAmelCase__			:  Any											= torch.from_numpy(lowerCAmelCase_      ).to(lowerCAmelCase_      )

    return va
def      _SCREAMING_SNAKE_CASE  (	UpperCamelCase							,   UpperCamelCase      ):



    """simple docstring"""





    lowerCAmelCase__			:  Dict											= F.normalize(lowerCAmelCase_							,   dim=-1      )
    lowerCAmelCase__			:  Union[str, Any]											= F.normalize(lowerCAmelCase_							,   dim=-1      )
    return (x - y).norm(dim=-1      ).div(2      ).arcsin().pow(2      ).mul(2      )
def      _SCREAMING_SNAKE_CASE  (	UpperCamelCase							,   UpperCamelCase      ):



    """simple docstring"""





    for param in model.parameters():
        lowerCAmelCase__			:  Optional[Any]											= value






class 	lowerCAmelCase_(    _UpperCamelCase					):
  '''simple docstring'''




  def __init__(   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase=None  ,__UpperCAmelCase=None  ,__UpperCAmelCase=None  ,)   ->		Optional[Any]:
      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  ,)
      lowerCAmelCase__			:  List[Any]											= (
          feature_extractor.size
          if isinstance(feature_extractor.size  ,_UpperCAmelCase				)
          else feature_extractor.size['shortest_edge']
      )
      lowerCAmelCase__			:  str											= 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 							UpperCAmelCase_     (   self  ,__UpperCAmelCase = "auto"				)   ->		Any:
      if slice_size == "auto":
          # half the attention head size is usually a good trade-off between
          # speed and memory
          lowerCAmelCase__			:  Optional[int]											= self.unet.config.attention_head_dim // 2
      self.unet.set_attention_slice(_UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      self.enable_attention_slicing(_UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
      set_requires_grad(self.vae  ,_UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		Any:
      set_requires_grad(self.vae  ,_UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      set_requires_grad(self.unet  ,_UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self				)   ->		str:
      set_requires_grad(self.unet  ,_UpperCAmelCase				)




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Union[str, Any]:
      # get the original timestep using init_timestep
      lowerCAmelCase__			:  Union[str, Any]											= min(int(num_inference_steps * strength				)  ,_UpperCAmelCase				)

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

      return timesteps, num_inference_steps - t_start




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase=None				)   ->		Tuple:
      if not isinstance(_UpperCAmelCase  ,torch.Tensor				):
          raise ValueError(F"""`image` has to be of type `torch.Tensor` but is {type(_UpperCAmelCase				)}"""				)

      lowerCAmelCase__			:  Tuple											= image.to(device=_UpperCAmelCase  ,dtype=_UpperCAmelCase				)

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

      # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
      lowerCAmelCase__			:  Dict											= 0.1_8_2_1_5 * init_latents
      lowerCAmelCase__			:  Optional[int]											= init_latents.repeat_interleave(_UpperCAmelCase  ,dim=0				)

      lowerCAmelCase__			:  Optional[Any]											= randn_tensor(init_latents.shape  ,generator=_UpperCAmelCase  ,device=_UpperCAmelCase  ,dtype=_UpperCAmelCase				)

      # get latents
      lowerCAmelCase__			:  Optional[Any]											= self.scheduler.add_noise(_UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase				)
      lowerCAmelCase__			:  List[str]											= init_latents

      return latents




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




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Dict:
      lowerCAmelCase__			:  str											= self.feature_extractor.preprocess(_UpperCAmelCase				)
      lowerCAmelCase__			:  Optional[Any]											= torch.from_numpy(clip_image_input["""pixel_values"""][0]				).unsqueeze(0				).to(self.device				).half()
      lowerCAmelCase__			:  List[Any]											= self.clip_model.get_image_features(_UpperCAmelCase				)
      lowerCAmelCase__			:  Any											= image_embeddings_clip / image_embeddings_clip.norm(p=2  ,dim=-1  ,keepdim=_UpperCAmelCase				)
      lowerCAmelCase__			:  str											= image_embeddings_clip.repeat_interleave(_UpperCAmelCase  ,dim=0				)
      return image_embeddings_clip




  @torch.enable_grad()
  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,)   ->		List[str]:
      lowerCAmelCase__			:  List[str]											= latents.detach().requires_grad_()

      lowerCAmelCase__			:  Dict											= self.scheduler.scale_model_input(_UpperCAmelCase  ,_UpperCAmelCase				)

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

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

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

      lowerCAmelCase__			:  Optional[Any]											= transforms.Resize(self.feature_extractor_size				)(_UpperCAmelCase				)
      lowerCAmelCase__			:  Union[str, Any]											= self.normalize(_UpperCAmelCase				).to(latents.dtype				)

      lowerCAmelCase__			:  Tuple											= self.clip_model.get_image_features(_UpperCAmelCase				)
      lowerCAmelCase__			:  Dict											= image_embeddings_clip / image_embeddings_clip.norm(p=2  ,dim=-1  ,keepdim=_UpperCAmelCase				)

      lowerCAmelCase__			:  Dict											= spherical_dist_loss(_UpperCAmelCase  ,_UpperCAmelCase				).mean() * clip_guidance_scale

      lowerCAmelCase__			:  Optional[int]											= -torch.autograd.grad(_UpperCAmelCase  ,_UpperCAmelCase				)[0]

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






  @torch.no_grad()
  def __call__(   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase = None  ,__UpperCAmelCase = None  ,__UpperCAmelCase = 512  ,__UpperCAmelCase = 512  ,__UpperCAmelCase = 0.6  ,__UpperCAmelCase = 50  ,__UpperCAmelCase = 7.5  ,__UpperCAmelCase = 1  ,__UpperCAmelCase = 0.0  ,__UpperCAmelCase = 100  ,__UpperCAmelCase = None  ,__UpperCAmelCase = "pil"  ,__UpperCAmelCase = True  ,__UpperCAmelCase = 0.8  ,__UpperCAmelCase = 0.1  ,__UpperCAmelCase = 0.1  ,)   ->		int:
      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:
          lowerCAmelCase__			:  Optional[int]											= [generator] + [None] * (batch_size - 1)

      lowerCAmelCase__			:  Optional[Any]											= [
          ('model', self.coca_model is None),
          ('tokenizer', self.coca_tokenizer is None),
          ('transform', self.coca_transform is None),
      ]
      lowerCAmelCase__			:  List[Any]											= [x[0] for x in coca_is_none if x[1]]
      lowerCAmelCase__			:  int											= ', '.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."""				)
          lowerCAmelCase__			:  Tuple											= 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."""				)
          lowerCAmelCase__			:  Union[str, Any]											= self.get_image_description(_UpperCAmelCase				)

      # get prompt text embeddings for content and style
      lowerCAmelCase__			:  Optional[int]											= self.tokenizer(
          _UpperCAmelCase  ,padding="""max_length"""  ,max_length=self.tokenizer.model_max_length  ,truncation=_UpperCAmelCase  ,return_tensors="""pt"""  ,)
      lowerCAmelCase__			:  List[str]											= self.text_encoder(content_text_input.input_ids.to(self.device				)				)[0]

      lowerCAmelCase__			:  Optional[int]											= self.tokenizer(
          _UpperCAmelCase  ,padding="""max_length"""  ,max_length=self.tokenizer.model_max_length  ,truncation=_UpperCAmelCase  ,return_tensors="""pt"""  ,)
      lowerCAmelCase__			:  str											= self.text_encoder(style_text_input.input_ids.to(self.device				)				)[0]

      lowerCAmelCase__			:  Any											= slerp(_UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase				)

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

      # set timesteps
      lowerCAmelCase__			:  int											= 'offset' in set(inspect.signature(self.scheduler.set_timesteps				).parameters.keys()				)
      lowerCAmelCase__			:  List[Any]											= {}
      if accepts_offset:
          lowerCAmelCase__			:  Union[str, Any]											= 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				)

      lowerCAmelCase__			:  List[Any]											= self.get_timesteps(_UpperCAmelCase  ,_UpperCAmelCase  ,self.device				)
      lowerCAmelCase__			:  List[Any]											= timesteps[:1].repeat(_UpperCAmelCase				)

      # Preprocess image
      lowerCAmelCase__			:  str											= preprocess(_UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase				)
      lowerCAmelCase__			:  Union[str, Any]											= self.prepare_latents(
          _UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase  ,text_embeddings.dtype  ,self.device  ,_UpperCAmelCase				)

      lowerCAmelCase__			:  str											= preprocess(_UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase				)
      lowerCAmelCase__			:  Optional[int]											= self.prepare_latents(
          _UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase  ,text_embeddings.dtype  ,self.device  ,_UpperCAmelCase				)

      lowerCAmelCase__			:  str											= slerp(_UpperCAmelCase  ,_UpperCAmelCase  ,_UpperCAmelCase				)

      if clip_guidance_scale > 0:
          lowerCAmelCase__			:  Optional[Any]											= self.get_clip_image_embeddings(_UpperCAmelCase  ,_UpperCAmelCase				)
          lowerCAmelCase__			:  Optional[Any]											= self.get_clip_image_embeddings(_UpperCAmelCase  ,_UpperCAmelCase				)
          lowerCAmelCase__			:  List[str]											= 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.
      lowerCAmelCase__			:  str											= guidance_scale > 1.0
      # get unconditional embeddings for classifier free guidance
      if do_classifier_free_guidance:
          lowerCAmelCase__			:  Union[str, Any]											= content_text_input.input_ids.shape[-1]
          lowerCAmelCase__			:  Tuple											= self.tokenizer([""""""]  ,padding="""max_length"""  ,max_length=_UpperCAmelCase  ,return_tensors="""pt"""				)
          lowerCAmelCase__			:  Dict											= self.text_encoder(uncond_input.input_ids.to(self.device				)				)[0]
          # duplicate unconditional embeddings for each generation per prompt
          lowerCAmelCase__			:  Optional[int]											= 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
          lowerCAmelCase__			:  Any											= torch.cat([uncond_embeddings, text_embeddings]				)

      # get the initial random noise unless the user supplied it

      # Unlike in other pipelines, latents need to be generated in the target device
      # for 1-to-1 results reproducibility with the CompVis implementation.
      # However this currently doesn't work in `mps`.
      lowerCAmelCase__			:  Dict											= (batch_size, self.unet.config.in_channels, height // 8, width // 8)
      lowerCAmelCase__			:  Union[str, Any]											= text_embeddings.dtype
      if latents is None:
          if self.device.type == "mps":
              # randn does not work reproducibly on mps
              lowerCAmelCase__			:  Optional[int]											= torch.randn(_UpperCAmelCase  ,generator=_UpperCAmelCase  ,device="""cpu"""  ,dtype=_UpperCAmelCase				).to(
                  self.device				)
          else:
              lowerCAmelCase__			:  Any											= 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}"""				)
          lowerCAmelCase__			:  List[str]											= latents.to(self.device				)

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

      # check if the scheduler accepts generator
      lowerCAmelCase__			:  Union[str, Any]											= 'generator' in set(inspect.signature(self.scheduler.step				).parameters.keys()				)
      if accepts_generator:
          lowerCAmelCase__			:  Union[str, Any]											= generator

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

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

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

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

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

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

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

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

      if not return_dict:
          return (image, None)

      return StableDiffusionPipelineOutput(images=_UpperCAmelCase  ,nsfw_content_detected=_UpperCAmelCase				)







'''simple docstring'''





import json
import os
from typing import Optional, Tuple

import regex as re

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


_lowerCAmelCase										=     logging.get_logger(__name__)

_lowerCAmelCase										=     {
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
}

_lowerCAmelCase										=     {
    '''vocab_file''': {'''ctrl''': '''https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-vocab.json'''},
    '''merges_file''': {'''ctrl''': '''https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-merges.txt'''},
}

_lowerCAmelCase										=     {
    '''ctrl''': 256,
}

_lowerCAmelCase										=     {
    '''Pregnancy''': 16_8629,
    '''Christianity''': 7675,
    '''Explain''': 10_6423,
    '''Fitness''': 6_3440,
    '''Saving''': 6_3163,
    '''Ask''': 2_7171,
    '''Ass''': 9_5985,
    '''Joke''': 16_3509,
    '''Questions''': 4_5622,
    '''Thoughts''': 4_9605,
    '''Retail''': 5_2342,
    '''Feminism''': 16_4338,
    '''Writing''': 1_1992,
    '''Atheism''': 19_2263,
    '''Netflix''': 4_8616,
    '''Computing''': 3_9639,
    '''Opinion''': 4_3213,
    '''Alone''': 4_4967,
    '''Funny''': 5_8917,
    '''Gaming''': 4_0358,
    '''Human''': 4088,
    '''India''': 1331,
    '''Joker''': 7_7138,
    '''Diet''': 3_6206,
    '''Legal''': 1_1859,
    '''Norman''': 4939,
    '''Tip''': 7_2689,
    '''Weight''': 5_2343,
    '''Movies''': 4_6273,
    '''Running''': 2_3425,
    '''Science''': 2090,
    '''Horror''': 3_7793,
    '''Confession''': 6_0572,
    '''Finance''': 1_2250,
    '''Politics''': 1_6360,
    '''Scary''': 19_1985,
    '''Support''': 1_2654,
    '''Technologies''': 3_2516,
    '''Teenage''': 6_6160,
    '''Event''': 3_2769,
    '''Learned''': 6_7460,
    '''Notion''': 18_2770,
    '''Wikipedia''': 3_7583,
    '''Books''': 6665,
    '''Extract''': 7_6050,
    '''Confessions''': 10_2701,
    '''Conspiracy''': 7_5932,
    '''Links''': 6_3674,
    '''Narcissus''': 15_0425,
    '''Relationship''': 5_4766,
    '''Relationships''': 13_4796,
    '''Reviews''': 4_1671,
    '''News''': 4256,
    '''Translation''': 2_6820,
    '''multilingual''': 12_8406,
}
def        _SCREAMING_SNAKE_CASE  (	UpperCamelCase      ):



    """simple docstring"""





    lowerCAmelCase__			:  Dict											= set()
    lowerCAmelCase__			:  List[str]											= word[0]
    for char in word[1:]:
        pairs.add((prev_char, char)      )
        lowerCAmelCase__			:  str											= char

    lowerCAmelCase__			:  int											= set(UpperCamelCase      )
    return pairs






class 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_					):
  '''simple docstring'''




  __lowercase							:				Optional[int]           =				VOCAB_FILES_NAMES
  __lowercase							:				Union[str, Any]           =				PRETRAINED_VOCAB_FILES_MAP
  __lowercase							:				List[Any]           =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  __lowercase							:				str           =				CONTROL_CODES




  def __init__(   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase="<unk>"  ,**__UpperCAmelCase				)   ->		Optional[Any]:
      super().__init__(unk_token=__UpperCAmelCase  ,**__UpperCAmelCase				)

      with open(__UpperCAmelCase  ,encoding="""utf-8"""				) as vocab_handle:
          lowerCAmelCase__			:  List[Any]											= json.load(__UpperCAmelCase				)
      lowerCAmelCase__			:  str											= {v: k for k, v in self.encoder.items()}
      with open(__UpperCAmelCase  ,encoding="""utf-8"""				) as merges_handle:
          lowerCAmelCase__			:  Any											= merges_handle.read().split("""\n"""				)[1:-1]
      lowerCAmelCase__			:  Optional[Any]											= [tuple(merge.split()				) for merge in merges]
      lowerCAmelCase__			:  Tuple											= dict(zip(__UpperCAmelCase  ,range(len(__UpperCAmelCase				)				)				)				)
      lowerCAmelCase__			:  int											= {}




  @property
  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      return len(self.encoder				)




  def 							UpperCAmelCase_     (   self				)   ->		Optional[int]:
      return dict(self.encoder  ,**self.added_tokens_encoder				)




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase				)   ->		Optional[Any]:
      if token in self.cache:
          return self.cache[token]
      lowerCAmelCase__			:  int											= tuple(__UpperCAmelCase				)
      lowerCAmelCase__			:  str											= tuple(list(word[:-1]				) + [word[-1] + """</w>"""]				)
      lowerCAmelCase__			:  Tuple											= get_pairs(__UpperCAmelCase				)

      if not pairs:
          return token

      while True:
          lowerCAmelCase__			:  Tuple											= min(__UpperCAmelCase  ,key=lambda __UpperCAmelCase				: self.bpe_ranks.get(__UpperCAmelCase  ,float("""inf"""				)				)				)
          if bigram not in self.bpe_ranks:
              break
          lowerCAmelCase__				,       lowerCAmelCase__			:  Optional[Any]											= bigram
          lowerCAmelCase__			:  Dict											= []
          lowerCAmelCase__			:  int											= 0
          while i < len(__UpperCAmelCase				):
              try:
                  lowerCAmelCase__			:  Any											= word.index(__UpperCAmelCase  ,__UpperCAmelCase				)
              except ValueError:
                  new_word.extend(word[i:]				)
                  break
              else:
                  new_word.extend(word[i:j]				)
                  lowerCAmelCase__			:  Any											= j

              if word[i] == first and i < len(__UpperCAmelCase				) - 1 and word[i + 1] == second:
                  new_word.append(first + second				)
                  i += 2
              else:
                  new_word.append(word[i]				)
                  i += 1
          lowerCAmelCase__			:  Any											= tuple(__UpperCAmelCase				)
          lowerCAmelCase__			:  Tuple											= new_word
          if len(__UpperCAmelCase				) == 1:
              break
          else:
              lowerCAmelCase__			:  List[str]											= get_pairs(__UpperCAmelCase				)
      lowerCAmelCase__			:  List[Any]											= """@@ """.join(__UpperCAmelCase				)
      lowerCAmelCase__			:  int											= word[:-4]
      lowerCAmelCase__			:  Optional[Any]											= word
      return word




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase				)   ->		Union[str, Any]:
      lowerCAmelCase__			:  List[Any]											= []

      lowerCAmelCase__			:  Optional[int]											= re.findall(R"""\S+\n?"""  ,__UpperCAmelCase				)

      for token in words:
          split_tokens.extend(list(self.bpe(__UpperCAmelCase				).split(""" """				)				)				)
      return split_tokens




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase				)   ->		Tuple:
      return self.encoder.get(__UpperCAmelCase  ,self.encoder.get(self.unk_token				)				)




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase				)   ->		Optional[int]:
      return self.decoder.get(__UpperCAmelCase  ,self.unk_token				)




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase				)   ->		Dict:
      lowerCAmelCase__			:  Optional[int]											= """ """.join(__UpperCAmelCase				).replace("""@@ """  ,""""""				).strip()
      return out_string






  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase = None				)   ->		Tuple[str]:
      if not os.path.isdir(__UpperCAmelCase				):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""				)
          return
      lowerCAmelCase__			:  Tuple											= os.path.join(
          __UpperCAmelCase  ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]				)
      lowerCAmelCase__			:  str											= os.path.join(
          __UpperCAmelCase  ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]				)

      with open(__UpperCAmelCase  ,"""w"""  ,encoding="""utf-8"""				) as f:
          f.write(json.dumps(self.encoder  ,indent=2  ,sort_keys=__UpperCAmelCase  ,ensure_ascii=__UpperCAmelCase				) + """\n"""				)

      lowerCAmelCase__			:  Optional[int]											= 0
      with open(__UpperCAmelCase  ,"""w"""  ,encoding="""utf-8"""				) as writer:
          writer.write("""#version: 0.2\n"""				)
          for bpe_tokens, token_index in sorted(self.bpe_ranks.items()  ,key=lambda __UpperCAmelCase				: kv[1]				):
              if index != token_index:
                  logger.warning(
                      F"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
                      """ Please check that the tokenizer is not corrupted!"""				)
                  lowerCAmelCase__			:  List[str]											= token_index
              writer.write(""" """.join(__UpperCAmelCase				) + """\n"""				)
              index += 1

      return vocab_file, merge_file

  # def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
  #     filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
  #     tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
  #     tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
  #     return ''.join(tokens_generated_so_far)


