






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


A_		:		Union[str, Any]		     =			logging.get_logger(__name__)  # pylint: disable=invalid-name




class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      ):
         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=_lowerCamelCase  ,      speech_processor=_lowerCamelCase  ,      vae=_lowerCamelCase  ,      text_encoder=_lowerCamelCase  ,      tokenizer=_lowerCamelCase  ,      unet=_lowerCamelCase  ,      scheduler=_lowerCamelCase  ,      feature_extractor=_lowerCamelCase  ,      )





       def       _a   ( self  ,      _lowerCamelCase = "auto"							):
         if slice_size == "auto":
           UpperCamelCase_: Optional[Any]   				=    self.unet.config.attention_head_dim // 2
         self.unet.set_attention_slice(_lowerCamelCase							)





       def       _a   ( self							):
         self.enable_attention_slicing(_lowerCamelCase							)





       @torch.no_grad()
       def __call__( self  ,      _lowerCamelCase  ,      _lowerCamelCase=1_6_0_0_0  ,      _lowerCamelCase = 5_1_2  ,      _lowerCamelCase = 5_1_2  ,      _lowerCamelCase = 5_0  ,      _lowerCamelCase = 7.5  ,      _lowerCamelCase = None  ,      _lowerCamelCase = 1  ,      _lowerCamelCase = 0.0  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = "pil"  ,      _lowerCamelCase = True  ,      _lowerCamelCase = None  ,      _lowerCamelCase = 1  ,      **_lowerCamelCase  ,      ):
         UpperCamelCase_: List[str]   				=    self.speech_processor.feature_extractor(
             _lowerCamelCase  ,      return_tensors='pt'  ,      sampling_rate=_lowerCamelCase							).input_features.to(self.device							)
         UpperCamelCase_: Union[str, Any]   				=    self.speech_model.generate(_lowerCamelCase  ,      max_length=4_8_0_0_0_0							)

         UpperCamelCase_: int   				=    self.speech_processor.tokenizer.batch_decode(_lowerCamelCase  ,      skip_special_tokens=_lowerCamelCase  ,      normalize=_lowerCamelCase							)[
             0
         ]

         if isinstance(_lowerCamelCase  ,      _lowerCamelCase							):
           UpperCamelCase_: List[Any]   				=    1
         elif isinstance(_lowerCamelCase  ,      _lowerCamelCase							):
           UpperCamelCase_: Any   				=    len(_lowerCamelCase							)
         else:
           raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(_lowerCamelCase							)}'''							)

         if height % 8 != 0 or width % 8 != 0:
           raise ValueError(f'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''							)

         if (callback_steps is None) or (
             callback_steps is not None and (not isinstance(_lowerCamelCase  ,      _lowerCamelCase							) or callback_steps <= 0)
         ):
           raise ValueError(
               f'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
               f''' {type(_lowerCamelCase							)}.'''							)

         # get prompt text embeddings
         UpperCamelCase_: int   				=    self.tokenizer(
             _lowerCamelCase  ,      padding='max_length'  ,      max_length=self.tokenizer.model_max_length  ,      return_tensors='pt'  ,      )
         UpperCamelCase_: Any   				=    text_inputs.input_ids

         if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
           UpperCamelCase_: List[str]   				=    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_: Union[str, 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_: int   				=    text_embeddings.shape
         UpperCamelCase_: Tuple   				=    text_embeddings.repeat(1  ,      _lowerCamelCase  ,      1							)
         UpperCamelCase_: Any   				=    text_embeddings.view(bs_embed * num_images_per_prompt  ,      _lowerCamelCase  ,      -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_: int   				=    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[str]   				=    [''] * batch_size
           elif type(_lowerCamelCase							) is not type(_lowerCamelCase							):
             raise TypeError(
                 f'''`negative_prompt` should be the same type to `prompt`, but got {type(_lowerCamelCase							)} !='''
                 f''' {type(_lowerCamelCase							)}.'''							)
           elif isinstance(_lowerCamelCase  ,      _lowerCamelCase							):
             UpperCamelCase_: List[str]   				=    [negative_prompt]
           elif batch_size != len(_lowerCamelCase							):
             raise ValueError(
                 f'''`negative_prompt`: {negative_prompt} has batch size {len(_lowerCamelCase							)}, 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(
               _lowerCamelCase  ,      padding='max_length'  ,      max_length=_lowerCamelCase  ,      truncation=_lowerCamelCase  ,      return_tensors='pt'  ,      )
           UpperCamelCase_: Tuple   				=    self.text_encoder(uncond_input.input_ids.to(self.device							)							)[0]

           # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
           UpperCamelCase_: Dict   				=    uncond_embeddings.shape[1]
           UpperCamelCase_: List[Any]   				=    uncond_embeddings.repeat(1  ,      _lowerCamelCase  ,      1							)
           UpperCamelCase_: List[str]   				=    uncond_embeddings.view(batch_size * num_images_per_prompt  ,      _lowerCamelCase  ,      -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_: Union[str, 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_: Optional[int]   				=    (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
         UpperCamelCase_: Tuple   				=    text_embeddings.dtype
         if latents is None:
           if self.device.type == "mps":
             # randn does not exist on mps
             UpperCamelCase_: int   				=    torch.randn(_lowerCamelCase  ,      generator=_lowerCamelCase  ,      device='cpu'  ,      dtype=_lowerCamelCase							).to(
                 self.device							)
           else:
             UpperCamelCase_: List[str]   				=    torch.randn(_lowerCamelCase  ,      generator=_lowerCamelCase  ,      device=self.device  ,      dtype=_lowerCamelCase							)
         else:
           if latents.shape != latents_shape:
             raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''							)
           UpperCamelCase_: Optional[int]   				=    latents.to(self.device							)

         # set timesteps
         self.scheduler.set_timesteps(_lowerCamelCase							)

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

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

         for i, t in enumerate(self.progress_bar(_lowerCamelCase							)							):
           # expand the latents if we are doing classifier free guidance
           UpperCamelCase_: Dict   				=    torch.cat([latents] * 2							) if do_classifier_free_guidance else latents
           UpperCamelCase_: Union[str, Any]   				=    self.scheduler.scale_model_input(_lowerCamelCase  ,      _lowerCamelCase							)

           # predict the noise residual
           UpperCamelCase_: List[Any]   				=    self.unet(_lowerCamelCase  ,      _lowerCamelCase  ,      encoder_hidden_states=_lowerCamelCase							).sample

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

           # compute the previous noisy sample x_t -> x_t-1
           UpperCamelCase_: str   				=    self.scheduler.step(_lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      **_lowerCamelCase							).prev_sample

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

         UpperCamelCase_: List[str]   				=    1 / 0.1_8_2_1_5 * latents
         UpperCamelCase_: List[Any]   				=    self.vae.decode(_lowerCamelCase							).sample

         UpperCamelCase_: Optional[Any]   				=    (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_: str   				=    image.cpu().permute(0  ,      2  ,      3  ,      1							).float().numpy()

         if output_type == "pil":
           UpperCamelCase_: int   				=    self.numpy_to_pil(_lowerCamelCase							)

         if not return_dict:
           return image

         return StableDiffusionPipelineOutput(images=_lowerCamelCase  ,      nsfw_content_detected=_lowerCamelCase							)







from abc import ABC, abstractmethod
from argparse import ArgumentParser




class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       @staticmethod
       @abstractmethod
       def       _a   ( _lowerCamelCase							):
         raise NotImplementedError()





       @abstractmethod
       def       _a   ( self							):
         raise NotImplementedError()