





"""simple docstring"""


import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextModel,
    CLIPTokenizer,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import logging


lowerCamelCase__						= logging.get_logger(__name__)  # pylint: disable=invalid-name



class   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase ):







  '''simple docstring'''


  def __init__(		self :							Any       , __a :							WhisperForConditionalGeneration       , __a :							WhisperProcessor       , __a :							AutoencoderKL       , __a :							CLIPTextModel       , __a :							CLIPTokenizer       , __a :							UNetaDConditionModel       , __a :							Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]       , __a :							StableDiffusionSafetyChecker       , __a :							CLIPImageProcessor       , )		->					Optional[int]:
        super().__init__()

        if safety_checker is None:
              logger.warning(
                  F'''You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure'''
                  " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
                  " results in services or applications open to the public. Both the diffusers team and Hugging Face"
                  " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
                  " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
                  " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."						)

        self.register_modules(
            speech_model=__a       , speech_processor=__a       , vae=__a       , text_encoder=__a       , tokenizer=__a       , unet=__a       , scheduler=__a       , feature_extractor=__a       , )

  def  __SCREAMING_SNAKE_CASE					(		self :							Dict       , __a :							Optional[Union[str, int]] = "auto"						)		->					List[str]:
        if slice_size == "auto":
              _UpperCamelCase   :						List[str]    =							self.unet.config.attention_head_dim // 2
        self.unet.set_attention_slice(__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					Tuple:
        self.enable_attention_slicing(__a						)


  @torch.no_grad()
  def __call__(		self :							Union[str, Any]       , __a :							Dict       , __a :							Any=1_6000       , __a :							int = 512       , __a :							int = 512       , __a :							int = 50       , __a :							float = 7.5       , __a :							Optional[Union[str, List[str]]] = None       , __a :							Optional[int] = 1       , __a :							float = 0.0       , __a :							Optional[torch.Generator] = None       , __a :							Optional[torch.FloatTensor] = None       , __a :							Optional[str] = "pil"       , __a :							bool = True       , __a :							Optional[Callable[[int, int, torch.FloatTensor], None]] = None       , __a :							int = 1       , **__a :							str       , )		->					Optional[Any]:
        _UpperCamelCase   :						List[str]    =							self.speech_processor.feature_extractor(
            __a       , return_tensors="pt"       , sampling_rate=__a						).input_features.to(self.device						)
        _UpperCamelCase   :						Optional[int]    =							self.speech_model.generate(__a       , max_length=48_0000						)

        _UpperCamelCase   :						Tuple    =							self.speech_processor.tokenizer.batch_decode(__a       , skip_special_tokens=__a       , normalize=__a						)[
            0
        ]

        if isinstance(__a       , __a						):
              _UpperCamelCase   :						Tuple    =							1
        elif isinstance(__a       , __a						):
              _UpperCamelCase   :						Union[str, Any]    =							len(__a						)
        else:
              raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(__a						)}'''						)

        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 (callback_steps is None) or (
            callback_steps is not None and (not isinstance(__a       , __a						) or callback_steps <= 0)
        ):
              raise ValueError(
                  F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
                  F''' {type(__a						)}.'''						)

        # get prompt text embeddings
        _UpperCamelCase   :						List[Any]    =							self.tokenizer(
            __a       , padding="max_length"       , max_length=self.tokenizer.model_max_length       , return_tensors="pt"       , )
        _UpperCamelCase   :						Optional[int]    =							text_inputs.input_ids

        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
              _UpperCamelCase   :						Optional[Any]    =							self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]						)
              logger.warning(
                  "The following part of your input was truncated because CLIP can only handle sequences up to"
                  F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''						)
              _UpperCamelCase   :						Optional[Any]    =							text_input_ids[:, : self.tokenizer.model_max_length]
        _UpperCamelCase   :						List[Any]    =							self.text_encoder(text_input_ids.to(self.device						)						)[0]

        # duplicate text embeddings for each generation per prompt, using mps friendly method
        _UpperCamelCase,	_UpperCamelCase,	_UpperCamelCase   :						str    =							text_embeddings.shape
        _UpperCamelCase   :						List[str]    =							text_embeddings.repeat(1       , __a       , 1						)
        _UpperCamelCase   :						str    =							text_embeddings.view(bs_embed * num_images_per_prompt       , __a       , -1						)

        # 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   :						Union[str, Any]    =							guidance_scale > 1.0
        # get unconditional embeddings for classifier free guidance
        if do_classifier_free_guidance:
              _UpperCamelCase   :						List[str]
              if negative_prompt is None:
                    _UpperCamelCase   :						List[Any]    =							[""] * batch_size
              elif type(__a						) is not type(__a						):
                    raise TypeError(
                        F'''`negative_prompt` should be the same type to `prompt`, but got {type(__a						)} !='''
                        F''' {type(__a						)}.'''						)
              elif isinstance(__a       , __a						):
                    _UpperCamelCase   :						Tuple    =							[negative_prompt]
              elif batch_size != len(__a						):
                    raise ValueError(
                        F'''`negative_prompt`: {negative_prompt} has batch size {len(__a						)}, but `prompt`:'''
                        F''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
                        " the batch size of `prompt`."						)
              else:
                    _UpperCamelCase   :						Any    =							negative_prompt

              _UpperCamelCase   :						Optional[int]    =							text_input_ids.shape[-1]
              _UpperCamelCase   :						str    =							self.tokenizer(
                  __a       , padding="max_length"       , max_length=__a       , truncation=__a       , return_tensors="pt"       , )
              _UpperCamelCase   :						List[Any]    =							self.text_encoder(uncond_input.input_ids.to(self.device						)						)[0]

              # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
              _UpperCamelCase   :						Optional[Any]    =							uncond_embeddings.shape[1]
              _UpperCamelCase   :						Dict    =							uncond_embeddings.repeat(1       , __a       , 1						)
              _UpperCamelCase   :						int    =							uncond_embeddings.view(batch_size * num_images_per_prompt       , __a       , -1						)

              # 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   :						Optional[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`.
        _UpperCamelCase   :						int    =							(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
        _UpperCamelCase   :						Union[str, Any]    =							text_embeddings.dtype
        if latents is None:
              if self.device.type == "mps":
                    # randn does not exist on mps
                    _UpperCamelCase   :						Dict    =							torch.randn(__a       , generator=__a       , device="cpu"       , dtype=__a						).to(
                        self.device						)
              else:
                    _UpperCamelCase   :						Optional[int]    =							torch.randn(__a       , generator=__a       , device=self.device       , dtype=__a						)
        else:
              if latents.shape != latents_shape:
                    raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''						)
              _UpperCamelCase   :						Union[str, Any]    =							latents.to(self.device						)

        # set timesteps
        self.scheduler.set_timesteps(__a						)

        # Some schedulers like PNDM have timesteps as arrays
        # It's more optimized to move all timesteps to correct device beforehand
        _UpperCamelCase   :						Optional[int]    =							self.scheduler.timesteps.to(self.device						)

        # scale the initial noise by the standard deviation required by the scheduler
        _UpperCamelCase   :						Optional[Any]    =							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   :						Any    =							{}
        if accepts_eta:
              _UpperCamelCase   :						Any    =							eta

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

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

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

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

              # call the callback, if provided
              if callback is not None and i % callback_steps == 0:
                    callback(__a       , __a       , __a						)

        _UpperCamelCase   :						Tuple    =							1 / 0.1_82_15 * latents
        _UpperCamelCase   :						List[str]    =							self.vae.decode(__a						).sample

        _UpperCamelCase   :						Dict    =							(image / 2 + 0.5).clamp(0       , 1						)

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
        _UpperCamelCase   :						Optional[Any]    =							image.cpu().permute(0       , 2       , 3       , 1						).float().numpy()

        if output_type == "pil":
              _UpperCamelCase   :						Union[str, Any]    =							self.numpy_to_pil(__a						)

        if not return_dict:
              return image

        return StableDiffusionPipelineOutput(images=__a       , nsfw_content_detected=__a						)







"""simple docstring"""


from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


lowerCamelCase__						= logging.get_logger(__name__)


if is_vision_available():
  import PIL



class   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Optional[int]						  =       ["pixel_values"]

  def __init__(		self :							Tuple       , __a :							bool = True       , __a :							Dict[str, int] = None       , __a :							PILImageResampling = PILImageResampling.BICUBIC       , __a :							bool = True       , __a :							Dict[str, int] = None       , __a :							bool = True       , __a :							Union[int, float] = 1 / 255       , __a :							bool = True       , __a :							Optional[Union[float, List[float]]] = None       , __a :							Optional[Union[float, List[float]]] = None       , __a :							bool = True       , **__a :							List[Any]       , )		->					None:
        super().__init__(**__a						)
        _UpperCamelCase   :						Optional[int]    =							size if size is not None else {"shortest_edge": 224}
        _UpperCamelCase   :						Optional[int]    =							get_size_dict(__a       , default_to_square=__a						)
        _UpperCamelCase   :						Union[str, Any]    =							crop_size if crop_size is not None else {"height": 224, "width": 224}
        _UpperCamelCase   :						Tuple    =							get_size_dict(__a       , default_to_square=__a       , param_name="crop_size"						)

        _UpperCamelCase   :						Optional[Any]    =							do_resize
        _UpperCamelCase   :						Dict    =							size
        _UpperCamelCase   :						Any    =							resample
        _UpperCamelCase   :						Tuple    =							do_center_crop
        _UpperCamelCase   :						str    =							crop_size
        _UpperCamelCase   :						Any    =							do_rescale
        _UpperCamelCase   :						Dict    =							rescale_factor
        _UpperCamelCase   :						int    =							do_normalize
        _UpperCamelCase   :						Dict    =							image_mean if image_mean is not None else OPENAI_CLIP_MEAN
        _UpperCamelCase   :						Tuple    =							image_std if image_std is not None else OPENAI_CLIP_STD
        _UpperCamelCase   :						List[Any]    =							do_convert_rgb

  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[int]       , __a :							np.ndarray       , __a :							Dict[str, int]       , __a :							PILImageResampling = PILImageResampling.BICUBIC       , __a :							Optional[Union[str, ChannelDimension]] = None       , **__a :							str       , )		->					np.ndarray:
        _UpperCamelCase   :						Any    =							get_size_dict(__a       , default_to_square=__a						)
        if "shortest_edge" not in size:
              raise ValueError(F'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''						)
        _UpperCamelCase   :						Dict    =							get_resize_output_image_size(__a       , size=size["shortest_edge"]       , default_to_square=__a						)
        return resize(__a       , size=__a       , resample=__a       , data_format=__a       , **__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]       , __a :							np.ndarray       , __a :							Dict[str, int]       , __a :							Optional[Union[str, ChannelDimension]] = None       , **__a :							Tuple       , )		->					np.ndarray:
        _UpperCamelCase   :						List[str]    =							get_size_dict(__a						)
        if "height" not in size or "width" not in size:
              raise ValueError(F'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''						)
        return center_crop(__a       , size=(size["height"], size["width"])       , data_format=__a       , **__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]       , __a :							np.ndarray       , __a :							Union[int, float]       , __a :							Optional[Union[str, ChannelDimension]] = None       , **__a :							Optional[Any]       , )		->					List[str]:
        return rescale(__a       , scale=__a       , data_format=__a       , **__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Tuple       , __a :							np.ndarray       , __a :							Union[float, List[float]]       , __a :							Union[float, List[float]]       , __a :							Optional[Union[str, ChannelDimension]] = None       , **__a :							Tuple       , )		->					np.ndarray:
        return normalize(__a       , mean=__a       , std=__a       , data_format=__a       , **__a						)


  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[int]       , __a :							ImageInput       , __a :							bool = None       , __a :							Dict[str, int] = None       , __a :							PILImageResampling = None       , __a :							bool = None       , __a :							int = None       , __a :							bool = None       , __a :							float = None       , __a :							bool = None       , __a :							Optional[Union[float, List[float]]] = None       , __a :							Optional[Union[float, List[float]]] = None       , __a :							bool = None       , __a :							Optional[Union[str, TensorType]] = None       , __a :							Optional[ChannelDimension] = ChannelDimension.FIRST       , **__a :							Union[str, Any]       , )		->					PIL.Image.Image:
        _UpperCamelCase   :						Tuple    =							do_resize if do_resize is not None else self.do_resize
        _UpperCamelCase   :						int    =							size if size is not None else self.size
        _UpperCamelCase   :						int    =							get_size_dict(__a       , param_name="size"       , default_to_square=__a						)
        _UpperCamelCase   :						List[Any]    =							resample if resample is not None else self.resample
        _UpperCamelCase   :						List[Any]    =							do_center_crop if do_center_crop is not None else self.do_center_crop
        _UpperCamelCase   :						Union[str, Any]    =							crop_size if crop_size is not None else self.crop_size
        _UpperCamelCase   :						List[Any]    =							get_size_dict(__a       , param_name="crop_size"       , default_to_square=__a						)
        _UpperCamelCase   :						Dict    =							do_rescale if do_rescale is not None else self.do_rescale
        _UpperCamelCase   :						Union[str, Any]    =							rescale_factor if rescale_factor is not None else self.rescale_factor
        _UpperCamelCase   :						Union[str, Any]    =							do_normalize if do_normalize is not None else self.do_normalize
        _UpperCamelCase   :						List[Any]    =							image_mean if image_mean is not None else self.image_mean
        _UpperCamelCase   :						List[Any]    =							image_std if image_std is not None else self.image_std
        _UpperCamelCase   :						List[str]    =							do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

        _UpperCamelCase   :						Any    =							make_list_of_images(__a						)

        if not valid_images(__a						):
              raise ValueError(
                  "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                  "torch.Tensor, tf.Tensor or jax.ndarray."						)

        if do_resize and size is None:
              raise ValueError("Size must be specified if do_resize is True."						)

        if do_center_crop and crop_size is None:
              raise ValueError("Crop size must be specified if do_center_crop is True."						)

        if do_rescale and rescale_factor is None:
              raise ValueError("Rescale factor must be specified if do_rescale is True."						)

        if do_normalize and (image_mean is None or image_std is None):
              raise ValueError("Image mean and std must be specified if do_normalize is True."						)

        # PIL RGBA images are converted to RGB
        if do_convert_rgb:
              _UpperCamelCase   :						int    =							[convert_to_rgb(__a						) for image in images]

        # All transformations expect numpy arrays.
        _UpperCamelCase   :						List[Any]    =							[to_numpy_array(__a						) for image in images]

        if do_resize:
              _UpperCamelCase   :						Union[str, Any]    =							[self.resize(image=__a       , size=__a       , resample=__a						) for image in images]

        if do_center_crop:
              _UpperCamelCase   :						Optional[Any]    =							[self.center_crop(image=__a       , size=__a						) for image in images]

        if do_rescale:
              _UpperCamelCase   :						Dict    =							[self.rescale(image=__a       , scale=__a						) for image in images]

        if do_normalize:
              _UpperCamelCase   :						List[Any]    =							[self.normalize(image=__a       , mean=__a       , std=__a						) for image in images]

        _UpperCamelCase   :						List[Any]    =							[to_channel_dimension_format(__a       , __a						) for image in images]

        _UpperCamelCase   :						str    =							{"pixel_values": images}
        return BatchFeature(data=__a       , tensor_type=__a						)
