


"""simple docstring"""







from queue import Queue
from typing import TYPE_CHECKING, Optional


if TYPE_CHECKING:
		from ..models.auto import AutoTokenizer




class 			UpperCamelCase_     :

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    Optional[int]      ,      UpperCAmelCase__      :    List[Any]       )	->     List[str]:
						raise NotImplementedError()




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     List[Any]:
						raise NotImplementedError()




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def __init__(			self      :    int      ,      UpperCAmelCase__      :    "AutoTokenizer"      ,      UpperCAmelCase__      :    bool = False      ,      **UpperCAmelCase__      :    int       )	->     Any:
						__SCREAMING_SNAKE_CASE						       =    tokenizer
						__SCREAMING_SNAKE_CASE						       =    skip_prompt
						__SCREAMING_SNAKE_CASE						       =    decode_kwargs

						# variables used in the streaming process
						__SCREAMING_SNAKE_CASE						       =    []
						__SCREAMING_SNAKE_CASE						       =    0
						__SCREAMING_SNAKE_CASE						       =    True




	def   UpperCAmelCase_				(			self      :    List[str]      ,      UpperCAmelCase__      :    Dict       )	->     Optional[Any]:
						if len(value.shape       ) > 1 and value.shape[0] > 1:
											raise ValueError("TextStreamer only supports batch size 1"       )
						elif len(value.shape       ) > 1:
											__SCREAMING_SNAKE_CASE						       =    value[0]

						if self.skip_prompt and self.next_tokens_are_prompt:
											__SCREAMING_SNAKE_CASE						       =    False
											return

						# Add the new token to the cache and decodes the entire thing.
						self.token_cache.extend(value.tolist()       )
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer.decode(self.token_cache      ,      **self.decode_kwargs       )

						# After the symbol for a new line, we flush the cache.
						if text.endswith("\n"       ):
											__SCREAMING_SNAKE_CASE						       =    text[self.print_len :]
											__SCREAMING_SNAKE_CASE						       =    []
											__SCREAMING_SNAKE_CASE						       =    0
						# If the last token is a CJK character, we print the characters.
						elif len(UpperCAmelCase__       ) > 0 and self._is_chinese_char(ord(text[-1]       )       ):
											__SCREAMING_SNAKE_CASE						       =    text[self.print_len :]
											self.print_len += len(UpperCAmelCase__       )
						# Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
						# which may change with the subsequent token -- there are probably smarter ways to do this!)
						else:
											__SCREAMING_SNAKE_CASE						       =    text[self.print_len : text.rfind(" "       ) + 1]
											self.print_len += len(UpperCAmelCase__       )

						self.on_finalized_text(UpperCAmelCase__       )




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     Tuple:
						# Flush the cache, if it exists
						if len(self.token_cache       ) > 0:
											__SCREAMING_SNAKE_CASE						       =    self.tokenizer.decode(self.token_cache      ,      **self.decode_kwargs       )
											__SCREAMING_SNAKE_CASE						       =    text[self.print_len :]
											__SCREAMING_SNAKE_CASE						       =    []
											__SCREAMING_SNAKE_CASE						       =    0
						else:
											__SCREAMING_SNAKE_CASE						       =    ""

						__SCREAMING_SNAKE_CASE						       =    True
						self.on_finalized_text(UpperCAmelCase__      ,      stream_end=UpperCAmelCase__       )




	def   UpperCAmelCase_				(			self      :    Any      ,      UpperCAmelCase__      :    str      ,      UpperCAmelCase__      :    bool = False       )	->     List[str]:
						print(UpperCAmelCase__      ,      flush=UpperCAmelCase__      ,      end="" if not stream_end else None       )




	def   UpperCAmelCase_				(			self      :    List[Any]      ,      UpperCAmelCase__      :    Optional[int]       )	->     List[Any]:
						# This defines a "chinese character" as anything in the CJK Unicode block:
						#   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
						#
						# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
						# despite its name. The modern Korean Hangul alphabet is a different block,
						# as is Japanese Hiragana and Katakana. Those alphabets are used to write
						# space-separated words, so they are not treated specially and handled
						# like the all of the other languages.
						if (
						    (cp >= 0x4_e_0_0 and cp <= 0x9_f_f_f)
						    or (cp >= 0x3_4_0_0 and cp <= 0x4_d_b_f)  #
						    or (cp >= 0x2_0_0_0_0 and cp <= 0x2_a_6_d_f)  #
						    or (cp >= 0x2_a_7_0_0 and cp <= 0x2_b_7_3_f)  #
						    or (cp >= 0x2_b_7_4_0 and cp <= 0x2_b_8_1_f)  #
						    or (cp >= 0x2_b_8_2_0 and cp <= 0x2_c_e_a_f)  #
						    or (cp >= 0xf_9_0_0 and cp <= 0xf_a_f_f)
						    or (cp >= 0x2_f_8_0_0 and cp <= 0x2_f_a_1_f)  #
						):  #
											return True

						return False




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def __init__(			self      :    Any      ,      UpperCAmelCase__      :    "AutoTokenizer"      ,      UpperCAmelCase__      :    bool = False      ,      UpperCAmelCase__      :    Optional[float] = None      ,      **UpperCAmelCase__      :    int       )	->     Optional[Any]:
						super().__init__(UpperCAmelCase__      ,      UpperCAmelCase__      ,      **UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    Queue()
						__SCREAMING_SNAKE_CASE						       =    None
						__SCREAMING_SNAKE_CASE						       =    timeout




	def   UpperCAmelCase_				(			self      :    Optional[Any]      ,      UpperCAmelCase__      :    str      ,      UpperCAmelCase__      :    bool = False       )	->     Tuple:
						self.text_queue.put(UpperCAmelCase__      ,      timeout=self.timeout       )
						if stream_end:
											self.text_queue.put(self.stop_signal      ,      timeout=self.timeout       )




	def __iter__(			self      :    Optional[Any]       )	->     Tuple:
						return self




	def   UpperCAmelCase_				(			self      :    str       )	->     Any:
						__SCREAMING_SNAKE_CASE						       =    self.text_queue.get(timeout=self.timeout       )
						if value == self.stop_signal:
											raise StopIteration()
						else:
											return value




"""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


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




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def __init__(			self      :    List[str]      ,      UpperCAmelCase__      :    WhisperForConditionalGeneration      ,      UpperCAmelCase__      :    WhisperProcessor      ,      UpperCAmelCase__      :    AutoencoderKL      ,      UpperCAmelCase__      :    CLIPTextModel      ,      UpperCAmelCase__      :    CLIPTokenizer      ,      UpperCAmelCase__      :    UNetaDConditionModel      ,      UpperCAmelCase__      :    Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]      ,      UpperCAmelCase__      :    StableDiffusionSafetyChecker      ,      UpperCAmelCase__      :    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=UpperCAmelCase__      ,      speech_processor=UpperCAmelCase__      ,      vae=UpperCAmelCase__      ,      text_encoder=UpperCAmelCase__      ,      tokenizer=UpperCAmelCase__      ,      unet=UpperCAmelCase__      ,      scheduler=UpperCAmelCase__      ,      feature_extractor=UpperCAmelCase__      ,      )




	def   UpperCAmelCase_				(			self      :    Optional[Any]      ,      UpperCAmelCase__      :    Optional[Union[str, int]] = "auto"       )	->     str:
						if slice_size == "auto":
											__SCREAMING_SNAKE_CASE						       =    self.unet.config.attention_head_dim // 2
						self.unet.set_attention_slice(UpperCAmelCase__       )




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




	@torch.no_grad()
	def __call__(			self      :    str      ,      UpperCAmelCase__      :    Optional[int]      ,      UpperCAmelCase__      :    str=1_6_0_0_0      ,      UpperCAmelCase__      :    int = 5_1_2      ,      UpperCAmelCase__      :    int = 5_1_2      ,      UpperCAmelCase__      :    int = 5_0      ,      UpperCAmelCase__      :    float = 7.5      ,      UpperCAmelCase__      :    Optional[Union[str, List[str]]] = None      ,      UpperCAmelCase__      :    Optional[int] = 1      ,      UpperCAmelCase__      :    float = 0.0      ,      UpperCAmelCase__      :    Optional[torch.Generator] = None      ,      UpperCAmelCase__      :    Optional[torch.FloatTensor] = None      ,      UpperCAmelCase__      :    Optional[str] = "pil"      ,      UpperCAmelCase__      :    bool = True      ,      UpperCAmelCase__      :    Optional[Callable[[int, int, torch.FloatTensor], None]] = None      ,      UpperCAmelCase__      :    int = 1      ,      **UpperCAmelCase__      :    Dict      ,      )	->     Any:
						__SCREAMING_SNAKE_CASE						       =    self.speech_processor.feature_extractor(
						    UpperCAmelCase__      ,      return_tensors="pt"      ,      sampling_rate=UpperCAmelCase__       ).input_features.to(self.device       )
						__SCREAMING_SNAKE_CASE						       =    self.speech_model.generate(UpperCAmelCase__      ,      max_length=4_8_0_0_0_0       )

						__SCREAMING_SNAKE_CASE						       =    self.speech_processor.tokenizer.batch_decode(UpperCAmelCase__      ,      skip_special_tokens=UpperCAmelCase__      ,      normalize=UpperCAmelCase__       )[
						    0
						]

						if isinstance(UpperCAmelCase__      ,      UpperCAmelCase__       ):
											__SCREAMING_SNAKE_CASE						       =    1
						elif isinstance(UpperCAmelCase__      ,      UpperCAmelCase__       ):
											__SCREAMING_SNAKE_CASE						       =    len(UpperCAmelCase__       )
						else:
											raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(UpperCAmelCase__       )}"""       )

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

						# get prompt text embeddings
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(
						    UpperCAmelCase__      ,      padding="max_length"      ,      max_length=self.tokenizer.model_max_length      ,      return_tensors="pt"      ,      )
						__SCREAMING_SNAKE_CASE						       =    text_inputs.input_ids

						if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
											__SCREAMING_SNAKE_CASE						       =    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}"""       )
											__SCREAMING_SNAKE_CASE						       =    text_input_ids[:, : self.tokenizer.model_max_length]
						__SCREAMING_SNAKE_CASE						       =    self.text_encoder(text_input_ids.to(self.device       )       )[0]

						# duplicate text embeddings for each generation per prompt, using mps friendly method
						__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    text_embeddings.shape
						__SCREAMING_SNAKE_CASE						       =    text_embeddings.repeat(1      ,      UpperCAmelCase__      ,      1       )
						__SCREAMING_SNAKE_CASE						       =    text_embeddings.view(bs_embed * num_images_per_prompt      ,      UpperCAmelCase__      ,      -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.
						__SCREAMING_SNAKE_CASE						       =    guidance_scale > 1.0
						# get unconditional embeddings for classifier free guidance
						if do_classifier_free_guidance:
											__SCREAMING_SNAKE_CASE						       =    42
											if negative_prompt is None:
																__SCREAMING_SNAKE_CASE						       =    [""] * batch_size
											elif type(UpperCAmelCase__       ) is not type(UpperCAmelCase__       ):
																raise TypeError(
																    F"""`negative_prompt` should be the same type to `prompt`, but got {type(UpperCAmelCase__       )} !="""
																    F""" {type(UpperCAmelCase__       )}."""       )
											elif isinstance(UpperCAmelCase__      ,      UpperCAmelCase__       ):
																__SCREAMING_SNAKE_CASE						       =    [negative_prompt]
											elif batch_size != len(UpperCAmelCase__       ):
																raise ValueError(
																    F"""`negative_prompt`: {negative_prompt} has batch size {len(UpperCAmelCase__       )}, but `prompt`:"""
																    F""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
																    " the batch size of `prompt`."       )
											else:
																__SCREAMING_SNAKE_CASE						       =    negative_prompt

											__SCREAMING_SNAKE_CASE						       =    text_input_ids.shape[-1]
											__SCREAMING_SNAKE_CASE						       =    self.tokenizer(
											    UpperCAmelCase__      ,      padding="max_length"      ,      max_length=UpperCAmelCase__      ,      truncation=UpperCAmelCase__      ,      return_tensors="pt"      ,      )
											__SCREAMING_SNAKE_CASE						       =    self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]

											# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
											__SCREAMING_SNAKE_CASE						       =    uncond_embeddings.shape[1]
											__SCREAMING_SNAKE_CASE						       =    uncond_embeddings.repeat(1      ,      UpperCAmelCase__      ,      1       )
											__SCREAMING_SNAKE_CASE						       =    uncond_embeddings.view(batch_size * num_images_per_prompt      ,      UpperCAmelCase__      ,      -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
											__SCREAMING_SNAKE_CASE						       =    torch.cat([uncond_embeddings, text_embeddings]       )

						# get the initial random noise unless the user supplied it

						# Unlike in other pipelines, latents need to be generated in the target device
						# for 1-to-1 results reproducibility with the CompVis implementation.
						# However this currently doesn't work in `mps`.
						__SCREAMING_SNAKE_CASE						       =    (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
						__SCREAMING_SNAKE_CASE						       =    text_embeddings.dtype
						if latents is None:
											if self.device.type == "mps":
																# randn does not exist on mps
																__SCREAMING_SNAKE_CASE						       =    torch.randn(UpperCAmelCase__      ,      generator=UpperCAmelCase__      ,      device="cpu"      ,      dtype=UpperCAmelCase__       ).to(
																    self.device       )
											else:
																__SCREAMING_SNAKE_CASE						       =    torch.randn(UpperCAmelCase__      ,      generator=UpperCAmelCase__      ,      device=self.device      ,      dtype=UpperCAmelCase__       )
						else:
											if latents.shape != latents_shape:
																raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""       )
											__SCREAMING_SNAKE_CASE						       =    latents.to(self.device       )

						# set timesteps
						self.scheduler.set_timesteps(UpperCAmelCase__       )

						# Some schedulers like PNDM have timesteps as arrays
						# It's more optimized to move all timesteps to correct device beforehand
						__SCREAMING_SNAKE_CASE						       =    self.scheduler.timesteps.to(self.device       )

						# scale the initial noise by the standard deviation required by the scheduler
						__SCREAMING_SNAKE_CASE						       =    latents * self.scheduler.init_noise_sigma

						# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
						# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
						# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
						# and should be between [0, 1]
						__SCREAMING_SNAKE_CASE						       =    "eta" in set(inspect.signature(self.scheduler.step       ).parameters.keys()       )
						__SCREAMING_SNAKE_CASE						       =    {}
						if accepts_eta:
											__SCREAMING_SNAKE_CASE						       =    eta

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

											# predict the noise residual
											__SCREAMING_SNAKE_CASE						       =    self.unet(UpperCAmelCase__      ,      UpperCAmelCase__      ,      encoder_hidden_states=UpperCAmelCase__       ).sample

											# perform guidance
											if do_classifier_free_guidance:
																__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    noise_pred.chunk(2       )
																__SCREAMING_SNAKE_CASE						       =    noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

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

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

						__SCREAMING_SNAKE_CASE						       =    1 / 0.18_215 * latents
						__SCREAMING_SNAKE_CASE						       =    self.vae.decode(UpperCAmelCase__       ).sample

						__SCREAMING_SNAKE_CASE						       =    (image / 2 + 0.5).clamp(0      ,      1       )

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

						if output_type == "pil":
											__SCREAMING_SNAKE_CASE						       =    self.numpy_to_pil(UpperCAmelCase__       )

						if not return_dict:
											return image

						return StableDiffusionPipelineOutput(images=UpperCAmelCase__      ,      nsfw_content_detected=UpperCAmelCase__       )
