

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

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


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







class        a__       (   a__       ):







		"""simple docstring"""




		def __init__(			self    : List[Any]				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : List[str]				,				)					->Optional[Any]:
								"""simple docstring"""







								super().__init__()
								self.register_modules(
								    vae=lowerCamelCase_				,				text_encoder=lowerCamelCase_				,				tokenizer=lowerCamelCase_				,				unet=lowerCamelCase_				,				scheduler=lowerCamelCase_				,				safety_checker=lowerCamelCase_				,				feature_extractor=lowerCamelCase_				,				)


		def  _lowercase      (			self    : List[str]				,				UpperCAmelCase__    : str = "auto"		)					->Any:
								"""simple docstring"""







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


		def  _lowercase      (			self    : str		)					->Optional[Any]:
								"""simple docstring"""







								self.enable_attention_slicing(lowerCamelCase_		)




		@torch.no_grad()
		def __call__(			self    : Optional[Any]				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : List[str] = 5_1_2				,				UpperCAmelCase__    : Union[str, Any] = 5_1_2				,				UpperCAmelCase__    : str = 5_0				,				UpperCAmelCase__    : Optional[Any] = 7.5				,				UpperCAmelCase__    : int = None				,				UpperCAmelCase__    : List[Any] = 1				,				UpperCAmelCase__    : Optional[Any] = 0.0				,				UpperCAmelCase__    : Any = None				,				UpperCAmelCase__    : Optional[int] = None				,				UpperCAmelCase__    : str = "pil"				,				UpperCAmelCase__    : Union[str, Any] = True				,				UpperCAmelCase__    : List[Any] = None				,				UpperCAmelCase__    : Optional[Any] = 1				,				UpperCAmelCase__    : Optional[int] = None				,				**UpperCAmelCase__    : List[str]				,				)					->Any:
								"""simple docstring"""







								if isinstance(lowerCamelCase_				,				lowerCamelCase_		):
														SCREAMING_SNAKE_CASE :      int       =   1
								elif isinstance(lowerCamelCase_				,				lowerCamelCase_		):
														SCREAMING_SNAKE_CASE :      Optional[int]       =   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
								SCREAMING_SNAKE_CASE :      List[Any]       =   self.tokenizer(
								    lowerCamelCase_				,				padding="""max_length"""				,				max_length=self.tokenizer.model_max_length				,				return_tensors="""pt"""				,				)
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   text_inputs.input_ids

								if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
														SCREAMING_SNAKE_CASE :      List[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}"		)
														SCREAMING_SNAKE_CASE :      List[str]       =   text_input_ids[:, : self.tokenizer.model_max_length]

								if text_embeddings is None:
														SCREAMING_SNAKE_CASE :      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
								SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      Dict       =   text_embeddings.shape
								SCREAMING_SNAKE_CASE :      List[str]       =   text_embeddings.repeat(1				,				lowerCamelCase_				,				1		)
								SCREAMING_SNAKE_CASE :      int       =   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.
								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   guidance_scale > 1.0
								# get unconditional embeddings for classifier free guidance
								if do_classifier_free_guidance:
														SCREAMING_SNAKE_CASE :      Optional[int]       =   4_2
														if negative_prompt is None:
																				SCREAMING_SNAKE_CASE :      str       =   [""""""]
														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_		):
																				SCREAMING_SNAKE_CASE :      Tuple       =   [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:
																				SCREAMING_SNAKE_CASE :      str       =   negative_prompt

														SCREAMING_SNAKE_CASE :      List[Any]       =   text_input_ids.shape[-1]
														SCREAMING_SNAKE_CASE :      Dict       =   self.tokenizer(
														    lowerCamelCase_				,				padding="""max_length"""				,				max_length=lowerCamelCase_				,				truncation=lowerCamelCase_				,				return_tensors="""pt"""				,				)
														SCREAMING_SNAKE_CASE :      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
														SCREAMING_SNAKE_CASE :      Optional[int]       =   uncond_embeddings.shape[1]
														SCREAMING_SNAKE_CASE :      List[str]       =   uncond_embeddings.repeat(lowerCamelCase_				,				lowerCamelCase_				,				1		)
														SCREAMING_SNAKE_CASE :      Dict       =   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
														SCREAMING_SNAKE_CASE :      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`.
								SCREAMING_SNAKE_CASE :      Any       =   (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   (batch_size * num_images_per_prompt, self.unet.config.in_channels, 6_4, 6_4)
								SCREAMING_SNAKE_CASE :      str       =   text_embeddings.dtype
								if latents is None:
														if self.device.type == "mps":
																				# randn does not exist on mps
																				SCREAMING_SNAKE_CASE :      Any       =   torch.randn(
																				    lowerCamelCase_				,				generator=lowerCamelCase_				,				device="""cpu"""				,				dtype=lowerCamelCase_		).to(self.device		)
																				SCREAMING_SNAKE_CASE :      List[str]       =   torch.randn(lowerCamelCase_				,				generator=lowerCamelCase_				,				device="""cpu"""				,				dtype=lowerCamelCase_		).to(
																				    self.device		)
														else:
																				SCREAMING_SNAKE_CASE :      Tuple       =   torch.randn(
																				    lowerCamelCase_				,				generator=lowerCamelCase_				,				device=self.device				,				dtype=lowerCamelCase_		)
																				SCREAMING_SNAKE_CASE :      List[str]       =   torch.randn(lowerCamelCase_				,				generator=lowerCamelCase_				,				device=self.device				,				dtype=lowerCamelCase_		)
								else:
														if latents_reference.shape != latents_shape:
																				raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"		)
														SCREAMING_SNAKE_CASE :      List[Any]       =   latents_reference.to(self.device		)
														SCREAMING_SNAKE_CASE :      Union[str, Any]       =   latents.to(self.device		)

								# This is the key part of the pipeline where we
								# try to ensure that the generated images w/ the same seed
								# but different sizes actually result in similar images
								SCREAMING_SNAKE_CASE :      int       =   (latents_shape[3] - latents_shape_reference[3]) // 2
								SCREAMING_SNAKE_CASE :      Optional[int]       =   (latents_shape[2] - latents_shape_reference[2]) // 2
								SCREAMING_SNAKE_CASE :      Optional[int]       =   latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
								SCREAMING_SNAKE_CASE :      List[str]       =   0 if dx < 0 else dx
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   0 if dy < 0 else dy
								SCREAMING_SNAKE_CASE :      Dict       =   max(-dx				,				0		)
								SCREAMING_SNAKE_CASE :      str       =   max(-dy				,				0		)
								# import pdb
								# pdb.set_trace()
								SCREAMING_SNAKE_CASE :      Optional[Any]       =   latents_reference[:, :, dy : dy + h, dx : dx + w]

								# 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
								SCREAMING_SNAKE_CASE :      Optional[int]       =   self.scheduler.timesteps.to(self.device		)

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

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

														# predict the noise residual
														SCREAMING_SNAKE_CASE :      Optional[int]       =   self.unet(lowerCamelCase_				,				lowerCamelCase_				,				encoder_hidden_states=lowerCamelCase_		).sample

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

														# compute the previous noisy sample x_t -> x_t-1
														SCREAMING_SNAKE_CASE :      List[Any]       =   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_		)

								SCREAMING_SNAKE_CASE :      Any       =   1 / 0.1_82_15 * latents
								SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.vae.decode(lowerCamelCase_		).sample

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

								if self.safety_checker is not None:
														SCREAMING_SNAKE_CASE :      Optional[int]       =   self.feature_extractor(self.numpy_to_pil(lowerCamelCase_		)				,				return_tensors="""pt"""		).to(
														    self.device		)
														SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      str       =   self.safety_checker(
														    images=lowerCamelCase_				,				clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype		)		)
								else:
														SCREAMING_SNAKE_CASE :      List[str]       =   None

								if output_type == "pil":
														SCREAMING_SNAKE_CASE :      Optional[int]       =   self.numpy_to_pil(lowerCamelCase_		)

								if not return_dict:
														return (image, has_nsfw_concept)

								return StableDiffusionPipelineOutput(images=lowerCamelCase_				,				nsfw_content_detected=lowerCamelCase_		)







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


UpperCAmelCase__ :     List[Any]				    =    {"""configuration_vit""": ["""VIT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViTConfig""", """ViTOnnxConfig"""]}

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     Union[str, Any]				    =    ["""ViTFeatureExtractor"""]
  UpperCAmelCase__ :     List[Any]				    =    ["""ViTImageProcessor"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     List[Any]				    =    [
      """VIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """ViTForImageClassification""",
      """ViTForMaskedImageModeling""",
      """ViTModel""",
      """ViTPreTrainedModel""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     Tuple				    =    [
      """TFViTForImageClassification""",
      """TFViTModel""",
      """TFViTPreTrainedModel""",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     str				    =    [
      """FlaxViTForImageClassification""",
      """FlaxViTModel""",
      """FlaxViTPreTrainedModel""",
  ]

if TYPE_CHECKING:
  from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig, ViTOnnxConfig

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_vit import ViTFeatureExtractor
    from .image_processing_vit import ViTImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_vit import (
        VIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        ViTForImageClassification,
        ViTForMaskedImageModeling,
        ViTModel,
        ViTPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_vit import TFViTForImageClassification, TFViTModel, TFViTPreTrainedModel

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel, FlaxViTPreTrainedModel


else:
  import sys

  UpperCAmelCase__ :     Optional[Any]				    =    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




