






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


A__	:			Tuple			   = logging.get_logger(__name__)  # pylint: disable=invalid-name
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				def __init__(     self						:    Tuple,  lowerCamelCase						:    AutoencoderKL,  lowerCamelCase						:    CLIPTextModel,  lowerCamelCase						:    CLIPTokenizer,  lowerCamelCase						:    UNetaDConditionModel,  lowerCamelCase						:    Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],  lowerCamelCase						:    StableDiffusionSafetyChecker,  lowerCamelCase						:    CLIPImageProcessor,  ):

								'''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						:    int,  lowerCamelCase						:    Optional[Union[str, int]] = "auto"  ):

								'''simple docstring'''
								if slice_size == "auto":
												# half the attention head size is usually a good trade-off between
												# speed and memory
												lowercase__    =			self.unet.config.attention_head_dim // 2
								self.unet.set_attention_slice(lowerCamelCase  )




				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								self.enable_attention_slicing(lowerCamelCase  )




				@torch.no_grad()
				def __call__(     self						:    Any,  lowerCamelCase						:    Union[str, List[str]],  lowerCamelCase						:    int = 512,  lowerCamelCase						:    int = 512,  lowerCamelCase						:    int = 50,  lowerCamelCase						:    float = 7.5,  lowerCamelCase						:    Optional[Union[str, List[str]]] = None,  lowerCamelCase						:    Optional[int] = 1,  lowerCamelCase						:    float = 0.0,  lowerCamelCase						:    Optional[torch.Generator] = None,  lowerCamelCase						:    Optional[torch.FloatTensor] = None,  lowerCamelCase						:    Optional[str] = "pil",  lowerCamelCase						:    bool = True,  lowerCamelCase						:    Optional[Callable[[int, int, torch.FloatTensor], None]] = None,  lowerCamelCase						:    int = 1,  lowerCamelCase						:    Optional[torch.FloatTensor] = None,  **lowerCamelCase						:    Any,  ):

								'''simple docstring'''

								if isinstance(lowerCamelCase,  lowerCamelCase  ):
												lowercase__    =			1
								elif isinstance(lowerCamelCase,  lowerCamelCase  ):
												lowercase__    =			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
								lowercase__    =			self.tokenizer(
								    lowerCamelCase,  padding='''max_length''',  max_length=self.tokenizer.model_max_length,  return_tensors='''pt''',  )
								lowercase__    =			text_inputs.input_ids

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

								if text_embeddings is None:
												lowercase__    =			self.text_encoder(text_input_ids.to(self.device  )  )[0]

								# duplicate text embeddings for each generation per prompt, using mps friendly method
								lowercase__ ,						lowercase__ ,						lowercase__    =			text_embeddings.shape
								lowercase__    =			text_embeddings.repeat(1,  lowerCamelCase,  1  )
								lowercase__    =			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.
								lowercase__    =			guidance_scale > 1.0
								# get unconditional embeddings for classifier free guidance
								if do_classifier_free_guidance:
												lowercase__    =			42
												if negative_prompt is None:
																lowercase__    =			['''''']
												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  ):
																lowercase__    =			[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:
																lowercase__    =			negative_prompt

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

												# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
												lowercase__    =			uncond_embeddings.shape[1]
												lowercase__    =			uncond_embeddings.repeat(lowerCamelCase,  lowerCamelCase,  1  )
												lowercase__    =			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
												lowercase__    =			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`.
								lowercase__    =			(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
								lowercase__    =			(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
								lowercase__    =			text_embeddings.dtype
								if latents is None:
												if self.device.type == "mps":
																# randn does not exist on mps
																lowercase__    =			torch.randn(
																    lowerCamelCase,  generator=lowerCamelCase,  device='''cpu''',  dtype=lowerCamelCase  ).to(self.device  )
																lowercase__    =			torch.randn(lowerCamelCase,  generator=lowerCamelCase,  device='''cpu''',  dtype=lowerCamelCase  ).to(
																    self.device  )
												else:
																lowercase__    =			torch.randn(
																    lowerCamelCase,  generator=lowerCamelCase,  device=self.device,  dtype=lowerCamelCase  )
																lowercase__    =			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}"""  )
												lowercase__    =			latents_reference.to(self.device  )
												lowercase__    =			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
								lowercase__    =			(latents_shape[3] - latents_shape_reference[3]) // 2
								lowercase__    =			(latents_shape[2] - latents_shape_reference[2]) // 2
								lowercase__    =			latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
								lowercase__    =			latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
								lowercase__    =			0 if dx < 0 else dx
								lowercase__    =			0 if dy < 0 else dy
								lowercase__    =			max(-dx,  0  )
								lowercase__    =			max(-dy,  0  )
								# import pdb
								# pdb.set_trace()
								lowercase__    =			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
								lowercase__    =			self.scheduler.timesteps.to(self.device  )

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

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

												# predict the noise residual
												lowercase__    =			self.unet(lowerCamelCase,  lowerCamelCase,  encoder_hidden_states=lowerCamelCase  ).sample

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

												# compute the previous noisy sample x_t -> x_t-1
												lowercase__    =			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  )

								lowercase__    =			1 / 0.18215 * latents
								lowercase__    =			self.vae.decode(lowerCamelCase  ).sample

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

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

								if self.safety_checker is not None:
												lowercase__    =			self.feature_extractor(self.numpy_to_pil(lowerCamelCase  ),  return_tensors='''pt'''  ).to(
												    self.device  )
												lowercase__ ,						lowercase__    =			self.safety_checker(
												    images=lowerCamelCase,  clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype  )  )
								else:
												lowercase__    =			None

								if output_type == "pil":
												lowercase__    =			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_sentencepiece_available


A__	:			Optional[int]			   = {}

try:
							if not is_sentencepiece_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							A__	:			Optional[int]			   = ['GPTSw3Tokenizer']


if TYPE_CHECKING:
							try:
														if not is_sentencepiece_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_gpt_swa import GPTSwaTokenizer

else:
							import sys

							A__	:			Optional[int]			   = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




