





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       : Union[str, Any]			     =       logging.get_logger(__name__)  # pylint: disable=invalid-name

class      lowerCAmelCase							(      a__    ):
 def __init__(   self						:   int    ,  UpperCAmelCase						:   Union[str, Any]    ,  UpperCAmelCase						:   Union[str, Any]    ,  UpperCAmelCase						:   Union[str, Any]    ,  UpperCAmelCase						:   int    ,  UpperCAmelCase						:   Optional[Any]    ,  UpperCAmelCase						:   Dict    ,  UpperCAmelCase						:   Any    ,  )    ->			Optional[Any]:
     super().__init__()
     self.register_modules(
         vae=lowerCamelCase_    ,  text_encoder=lowerCamelCase_    ,  tokenizer=lowerCamelCase_    ,  unet=lowerCamelCase_    ,  scheduler=lowerCamelCase_    ,  safety_checker=lowerCamelCase_    ,  feature_extractor=lowerCamelCase_    ,  )
 def   A_			(   self						:   Optional[int]    ,  UpperCAmelCase						:   List[Any] = "auto"			)    ->			Any:
     if slice_size == "auto":
         # half the attention head size is usually a good trade-off between
         # speed and memory
         lowerCamelCase__					:   Tuple      				=	self.unet.config.attention_head_dim // 2
     self.unet.set_attention_slice(lowerCamelCase_			)
 def   A_			(   self						:   Any			)    ->			Optional[Any]:
     self.enable_attention_slicing(lowerCamelCase_			)







 @torch.no_grad()
 def __call__(   self						:   Any    ,  UpperCAmelCase						:   Optional[int]    ,  UpperCAmelCase						:   Any = 512    ,  UpperCAmelCase						:   Optional[int] = 512    ,  UpperCAmelCase						:   Optional[int] = 50    ,  UpperCAmelCase						:   List[str] = 7.5    ,  UpperCAmelCase						:   Tuple = None    ,  UpperCAmelCase						:   str = 1    ,  UpperCAmelCase						:   Any = 0.0    ,  UpperCAmelCase						:   Dict = None    ,  UpperCAmelCase						:   List[str] = None    ,  UpperCAmelCase						:   Tuple = "pil"    ,  UpperCAmelCase						:   Any = True    ,  UpperCAmelCase						:   str = None    ,  UpperCAmelCase						:   str = 1    ,  UpperCAmelCase						:   List[Any] = None    ,  **UpperCAmelCase						:   Optional[Any]    ,  )    ->			Any:

     if isinstance(lowerCamelCase_    ,  lowerCamelCase_			):
         lowerCamelCase__					:   Any      				=	1
     elif isinstance(lowerCamelCase_    ,  lowerCamelCase_			):
         lowerCamelCase__					:   Union[str, 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
     lowerCamelCase__					:   Tuple      				=	self.tokenizer(
         lowerCamelCase_    ,  padding='max_length'    ,  max_length=self.tokenizer.model_max_length    ,  return_tensors='pt'    ,  )
     lowerCamelCase__					:   List[str]      				=	text_inputs.input_ids

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

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

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

         lowerCamelCase__					:   Optional[Any]      				=	text_input_ids.shape[-1]
         lowerCamelCase__					:   Any      				=	self.tokenizer(
             lowerCamelCase_    ,  padding='max_length'    ,  max_length=lowerCamelCase_    ,  truncation=lowerCamelCase_    ,  return_tensors='pt'    ,  )
         lowerCamelCase__					:   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
         lowerCamelCase__					:   List[str]      				=	uncond_embeddings.shape[1]
         lowerCamelCase__					:   Tuple      				=	uncond_embeddings.repeat(lowerCamelCase_    ,  lowerCamelCase_    ,  1			)
         lowerCamelCase__					:   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
         lowerCamelCase__					:   Dict      				=	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`.
     lowerCamelCase__					:   int      				=	(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
     lowerCamelCase__					:   Union[str, Any]      				=	(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
     lowerCamelCase__					:   Dict      				=	text_embeddings.dtype
     if latents is None:
         if self.device.type == "mps":
             # randn does not exist on mps
             lowerCamelCase__					:   int      				=	torch.randn(
                 lowerCamelCase_    ,  generator=lowerCamelCase_    ,  device='cpu'    ,  dtype=lowerCamelCase_			).to(self.device			)
             lowerCamelCase__					:   str      				=	torch.randn(lowerCamelCase_    ,  generator=lowerCamelCase_    ,  device='cpu'    ,  dtype=lowerCamelCase_			).to(
                 self.device			)
         else:
             lowerCamelCase__					:   Optional[int]      				=	torch.randn(
                 lowerCamelCase_    ,  generator=lowerCamelCase_    ,  device=self.device    ,  dtype=lowerCamelCase_			)
             lowerCamelCase__					:   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}"""			)
         lowerCamelCase__					:   Dict      				=	latents_reference.to(self.device			)
         lowerCamelCase__					:   Optional[int]      				=	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
     lowerCamelCase__					:   Optional[int]      				=	(latents_shape[3] - latents_shape_reference[3]) // 2
     lowerCamelCase__					:   Optional[int]      				=	(latents_shape[2] - latents_shape_reference[2]) // 2
     lowerCamelCase__					:   str      				=	latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
     lowerCamelCase__					:   Dict      				=	latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
     lowerCamelCase__					:   Union[str, Any]      				=	0 if dx < 0 else dx
     lowerCamelCase__					:   Dict      				=	0 if dy < 0 else dy
     lowerCamelCase__					:   str      				=	max(-dx    ,  0			)
     lowerCamelCase__					:   Tuple      				=	max(-dy    ,  0			)
     # import pdb
     # pdb.set_trace()
     lowerCamelCase__					:   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
     lowerCamelCase__					:   Tuple      				=	self.scheduler.timesteps.to(self.device			)

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

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

         # predict the noise residual
         lowerCamelCase__					:   Tuple      				=	self.unet(lowerCamelCase_    ,  lowerCamelCase_    ,  encoder_hidden_states=lowerCamelCase_			).sample

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

         # compute the previous noisy sample x_t -> x_t-1
         lowerCamelCase__					:   Union[str, 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_			)

     lowerCamelCase__					:   Any      				=	1 / 0.1_8_2_1_5 * latents
     lowerCamelCase__					:   Optional[Any]      				=	self.vae.decode(lowerCamelCase_			).sample

     lowerCamelCase__					:   List[str]      				=	(image / 2 + 0.5).clamp(0    ,  1			)

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

     if self.safety_checker is not None:
         lowerCamelCase__					:   int      				=	self.feature_extractor(self.numpy_to_pil(lowerCamelCase_			)    ,  return_tensors='pt'			).to(
             self.device			)
         lowerCamelCase__				,							lowerCamelCase__					:   Tuple      				=	self.safety_checker(
             images=lowerCamelCase_    ,  clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype			)			)
     else:
         lowerCamelCase__					:   List[Any]      				=	None

     if output_type == "pil":
         lowerCamelCase__					:   Tuple      				=	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

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCAmelCase       : List[str]			     =       {
    """configuration_autoformer""": [
        """AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """AutoformerConfig""",
    ],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _UpperCAmelCase       : List[Any]			     =       [
        """AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """AutoformerForPrediction""",
        """AutoformerModel""",
        """AutoformerPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_autoformer import (
        AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
        AutoformerConfig,
    )

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_autoformer import (
            AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            AutoformerForPrediction,
            AutoformerModel,
            AutoformerPreTrainedModel,
        )

else:
    import sys

    _UpperCAmelCase       : Dict			     =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





