





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


_snake_case  :   List[Any]						   =  logging.get_logger(__name__)  # pylint: disable=invalid-name



class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  def __init__(    self			:		Dict ,  lowerCAmelCase_			:		AutoencoderKL ,  lowerCAmelCase_			:		CLIPTextModel ,  lowerCAmelCase_			:		CLIPTokenizer ,  lowerCAmelCase_			:		UNetaDConditionModel ,  lowerCAmelCase_			:		Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] ,  lowerCAmelCase_			:		StableDiffusionSafetyChecker ,  lowerCAmelCase_			:		CLIPImageProcessor ,  )							->						Optional[int]:
     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			:		Union[str, Any] ,  lowerCAmelCase_			:		Optional[Union[str, int]] = "auto"					)							->						int:
     if slice_size == "auto":
        # half the attention head size is usually a good trade-off between
        # speed and memory
        __lowerCAmelCase							     =     self.unet.config.attention_head_dim // 2
     self.unet.set_attention_slice(lowerCAmelCase_					)






  def     lowercase		(    self			:		List[Any]					)							->						List[Any]:
     self.enable_attention_slicing(lowerCAmelCase_					)






  @torch.no_grad()
  def __call__(    self			:		str ,  lowerCAmelCase_			:		Union[str, List[str]] ,  lowerCAmelCase_			:		int = 5_1_2 ,  lowerCAmelCase_			:		int = 5_1_2 ,  lowerCAmelCase_			:		int = 5_0 ,  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_			:		List[Any] ,  )							->						Optional[int]:

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

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

     if text_embeddings is None:
        __lowerCAmelCase							     =     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							     =     text_embeddings.shape
     __lowerCAmelCase							     =     text_embeddings.repeat(1 ,  lowerCAmelCase_ ,  1					)
     __lowerCAmelCase							     =     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							     =     guidance_scale > 1.0
     # get unconditional embeddings for classifier free guidance
     if do_classifier_free_guidance:
        __lowerCAmelCase							     =     42
        if negative_prompt is None:
           __lowerCAmelCase							     =     ['']
        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							     =     [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							     =     negative_prompt

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

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

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

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

        # predict the noise residual
        __lowerCAmelCase							     =     self.unet(lowerCAmelCase_ ,  lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_					).sample

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

        # compute the previous noisy sample x_t -> x_t-1
        __lowerCAmelCase							     =     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							     =     1 / 0.1_82_15 * latents
     __lowerCAmelCase							     =     self.vae.decode(lowerCAmelCase_					).sample

     __lowerCAmelCase							     =     (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							     =     image.cpu().permute(0 ,  2 ,  3 ,  1					).float().numpy()

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

     if output_type == "pil":
        __lowerCAmelCase							     =     self.numpy_to_pil(lowerCAmelCase_					)

     if not return_dict:
        return (image, has_nsfw_concept)

     return StableDiffusionPipelineOutput(images=lowerCAmelCase_ ,  nsfw_content_detected=lowerCAmelCase_					)








from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block



@dataclass
class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			42





class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase   ):

  """simple docstring"""


  @register_to_config
  def __init__(    self			:		List[Any] ,  lowerCAmelCase_			:		int = 6_5_5_3_6 ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		int = 2 ,  lowerCAmelCase_			:		int = 2 ,  lowerCAmelCase_			:		int = 0 ,  lowerCAmelCase_			:		str = "fourier" ,  lowerCAmelCase_			:		bool = True ,  lowerCAmelCase_			:		bool = False ,  lowerCAmelCase_			:		float = 0.0 ,  lowerCAmelCase_			:		Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D") ,  lowerCAmelCase_			:		Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip") ,  lowerCAmelCase_			:		Tuple[str] = "UNetMidBlock1D" ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		Tuple[int] = (3_2, 3_2, 6_4) ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		int = 8 ,  lowerCAmelCase_			:		int = 1 ,  lowerCAmelCase_			:		bool = False ,  )							->						Optional[int]:
     super().__init__()
     __lowerCAmelCase							     =     sample_size

     # time
     if time_embedding_type == "fourier":
        __lowerCAmelCase							     =     GaussianFourierProjection(
            embedding_size=8 ,  set_W_to_weight=lowerCAmelCase_ ,  log=lowerCAmelCase_ ,  flip_sin_to_cos=lowerCAmelCase_					)
        __lowerCAmelCase							     =     2 * block_out_channels[0]
     elif time_embedding_type == "positional":
        __lowerCAmelCase							     =     Timesteps(
            block_out_channels[0] ,  flip_sin_to_cos=lowerCAmelCase_ ,  downscale_freq_shift=lowerCAmelCase_					)
        __lowerCAmelCase							     =     block_out_channels[0]

     if use_timestep_embedding:
        __lowerCAmelCase							     =     block_out_channels[0] * 4
        __lowerCAmelCase							     =     TimestepEmbedding(
            in_channels=lowerCAmelCase_ ,  time_embed_dim=lowerCAmelCase_ ,  act_fn=lowerCAmelCase_ ,  out_dim=block_out_channels[0] ,  )

     __lowerCAmelCase							     =     nn.ModuleList([]					)
     __lowerCAmelCase							     =     None
     __lowerCAmelCase							     =     nn.ModuleList([]					)
     __lowerCAmelCase							     =     None

     # down
     __lowerCAmelCase							     =     in_channels
     for i, down_block_type in enumerate(lowerCAmelCase_					):
        __lowerCAmelCase							     =     output_channel
        __lowerCAmelCase							     =     block_out_channels[i]

        if i == 0:
           input_channel += extra_in_channels

        __lowerCAmelCase							     =     i == len(lowerCAmelCase_					) - 1

        __lowerCAmelCase							     =     get_down_block(
            lowerCAmelCase_ ,  num_layers=lowerCAmelCase_ ,  in_channels=lowerCAmelCase_ ,  out_channels=lowerCAmelCase_ ,  temb_channels=block_out_channels[0] ,  add_downsample=not is_final_block or downsample_each_block ,  )
        self.down_blocks.append(lowerCAmelCase_					)

     # mid
     __lowerCAmelCase							     =     get_mid_block(
         lowerCAmelCase_ ,  in_channels=block_out_channels[-1] ,  mid_channels=block_out_channels[-1] ,  out_channels=block_out_channels[-1] ,  embed_dim=block_out_channels[0] ,  num_layers=lowerCAmelCase_ ,  add_downsample=lowerCAmelCase_ ,  )

     # up
     __lowerCAmelCase							     =     list(reversed(lowerCAmelCase_					)					)
     __lowerCAmelCase							     =     reversed_block_out_channels[0]
     if out_block_type is None:
        __lowerCAmelCase							     =     out_channels
     else:
        __lowerCAmelCase							     =     block_out_channels[0]

     for i, up_block_type in enumerate(lowerCAmelCase_					):
        __lowerCAmelCase							     =     output_channel
        __lowerCAmelCase							     =     (
            reversed_block_out_channels[i + 1] if i < len(lowerCAmelCase_					) - 1 else final_upsample_channels
        )

        __lowerCAmelCase							     =     i == len(lowerCAmelCase_					) - 1

        __lowerCAmelCase							     =     get_up_block(
            lowerCAmelCase_ ,  num_layers=lowerCAmelCase_ ,  in_channels=lowerCAmelCase_ ,  out_channels=lowerCAmelCase_ ,  temb_channels=block_out_channels[0] ,  add_upsample=not is_final_block ,  )
        self.up_blocks.append(lowerCAmelCase_					)
        __lowerCAmelCase							     =     output_channel

     # out
     __lowerCAmelCase							     =     norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4 ,  3_2					)
     __lowerCAmelCase							     =     get_out_block(
         out_block_type=lowerCAmelCase_ ,  num_groups_out=lowerCAmelCase_ ,  embed_dim=block_out_channels[0] ,  out_channels=lowerCAmelCase_ ,  act_fn=lowerCAmelCase_ ,  fc_dim=block_out_channels[-1] // 4 ,  )






  def     lowercase		(    self			:		Optional[Any] ,  lowerCAmelCase_			:		torch.FloatTensor ,  lowerCAmelCase_			:		Union[torch.Tensor, float, int] ,  lowerCAmelCase_			:		bool = True ,  )							->						Union[UNetaDOutput, Tuple]:
     __lowerCAmelCase							     =     timestep
     if not torch.is_tensor(lowerCAmelCase_					):
        __lowerCAmelCase							     =     torch.tensor([timesteps] ,  dtype=torch.long ,  device=sample.device					)
     elif torch.is_tensor(lowerCAmelCase_					) and len(timesteps.shape					) == 0:
        __lowerCAmelCase							     =     timesteps[None].to(sample.device					)

     __lowerCAmelCase							     =     self.time_proj(lowerCAmelCase_					)
     if self.config.use_timestep_embedding:
        __lowerCAmelCase							     =     self.time_mlp(lowerCAmelCase_					)
     else:
        __lowerCAmelCase							     =     timestep_embed[..., None]
        __lowerCAmelCase							     =     timestep_embed.repeat([1, 1, sample.shape[2]]					).to(sample.dtype					)
        __lowerCAmelCase							     =     timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])					)

     # 2. down
     __lowerCAmelCase							     =     ()
     for downsample_block in self.down_blocks:
        __lowerCAmelCase     ,     __lowerCAmelCase							     =     downsample_block(hidden_states=lowerCAmelCase_ ,  temb=lowerCAmelCase_					)
        down_block_res_samples += res_samples

     # 3. mid
     if self.mid_block:
        __lowerCAmelCase							     =     self.mid_block(lowerCAmelCase_ ,  lowerCAmelCase_					)

     # 4. up
     for i, upsample_block in enumerate(self.up_blocks					):
        __lowerCAmelCase							     =     down_block_res_samples[-1:]
        __lowerCAmelCase							     =     down_block_res_samples[:-1]
        __lowerCAmelCase							     =     upsample_block(lowerCAmelCase_ ,  res_hidden_states_tuple=lowerCAmelCase_ ,  temb=lowerCAmelCase_					)

     # 5. post-process
     if self.out_block:
        __lowerCAmelCase							     =     self.out_block(lowerCAmelCase_ ,  lowerCAmelCase_					)

     if not return_dict:
        return (sample,)

     return UNetaDOutput(sample=lowerCAmelCase_					)

