
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       =					logging.get_logger(__name__)  # pylint: disable=invalid-name







class 			A_  (     __lowerCamelCase  ):




	'''simple docstring'''


	def __init__(			self						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       snake_case						,       ):
		super().__init__()
		self.register_modules(
		    vae=SCREAMING_SNAKE_CASE__						,       text_encoder=SCREAMING_SNAKE_CASE__						,       tokenizer=SCREAMING_SNAKE_CASE__						,       unet=SCREAMING_SNAKE_CASE__						,       scheduler=SCREAMING_SNAKE_CASE__						,       safety_checker=SCREAMING_SNAKE_CASE__						,       feature_extractor=SCREAMING_SNAKE_CASE__						,       )





	def       SCREAMING_SNAKE_CASE__		(			self						,       snake_case = "auto"			):



		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(SCREAMING_SNAKE_CASE__			)





	def       SCREAMING_SNAKE_CASE__		(			self			):
		self.enable_attention_slicing(SCREAMING_SNAKE_CASE__			)





	@torch.no_grad()
	def __call__(			self						,       snake_case						,       snake_case = 512						,       snake_case = 512						,       snake_case = 50						,       snake_case = 7.5						,       snake_case = None						,       snake_case = 1						,       snake_case = 0.0						,       snake_case = None						,       snake_case = None						,       snake_case = "pil"						,       snake_case = True						,       snake_case = None						,       snake_case = 1						,       snake_case = None						,       **snake_case						,       ):



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

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

		# get prompt text embeddings
		lowercase	    =       self.tokenizer(
		    SCREAMING_SNAKE_CASE__						,       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						,       SCREAMING_SNAKE_CASE__						,       1			)
		lowercase	    =       text_embeddings.view(bs_embed * num_images_per_prompt						,       SCREAMING_SNAKE_CASE__						,       -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(SCREAMING_SNAKE_CASE__			) is not type(SCREAMING_SNAKE_CASE__			):
				raise TypeError(
				    F'''`negative_prompt` should be the same type to `prompt`, but got {type(SCREAMING_SNAKE_CASE__			)} !='''
				    F''' {type(SCREAMING_SNAKE_CASE__			)}.'''			)
			elif isinstance(SCREAMING_SNAKE_CASE__						,       SCREAMING_SNAKE_CASE__			):
				lowercase	    =       [negative_prompt]
			elif batch_size != len(SCREAMING_SNAKE_CASE__			):
				raise ValueError(
				    F'''`negative_prompt`: {negative_prompt} has batch size {len(SCREAMING_SNAKE_CASE__			)}, 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(
			    SCREAMING_SNAKE_CASE__						,       padding='max_length'						,       max_length=SCREAMING_SNAKE_CASE__						,       truncation=SCREAMING_SNAKE_CASE__						,       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(SCREAMING_SNAKE_CASE__						,       SCREAMING_SNAKE_CASE__						,       1			)
			lowercase	    =       uncond_embeddings.view(batch_size * num_images_per_prompt						,       SCREAMING_SNAKE_CASE__						,       -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(
				    SCREAMING_SNAKE_CASE__						,       generator=SCREAMING_SNAKE_CASE__						,       device='cpu'						,       dtype=SCREAMING_SNAKE_CASE__			).to(self.device			)
				lowercase	    =       torch.randn(SCREAMING_SNAKE_CASE__						,       generator=SCREAMING_SNAKE_CASE__						,       device='cpu'						,       dtype=SCREAMING_SNAKE_CASE__			).to(
				    self.device			)
			else:
				lowercase	    =       torch.randn(
				    SCREAMING_SNAKE_CASE__						,       generator=SCREAMING_SNAKE_CASE__						,       device=self.device						,       dtype=SCREAMING_SNAKE_CASE__			)
				lowercase	    =       torch.randn(SCREAMING_SNAKE_CASE__						,       generator=SCREAMING_SNAKE_CASE__						,       device=self.device						,       dtype=SCREAMING_SNAKE_CASE__			)
		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(SCREAMING_SNAKE_CASE__			)

		# 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(SCREAMING_SNAKE_CASE__			)			):
			# 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(SCREAMING_SNAKE_CASE__						,       SCREAMING_SNAKE_CASE__			)

			# predict the noise residual
			lowercase	    =       self.unet(SCREAMING_SNAKE_CASE__						,       SCREAMING_SNAKE_CASE__						,       encoder_hidden_states=SCREAMING_SNAKE_CASE__			).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(SCREAMING_SNAKE_CASE__						,       SCREAMING_SNAKE_CASE__						,       SCREAMING_SNAKE_CASE__						,       **SCREAMING_SNAKE_CASE__			).prev_sample

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

		lowercase	    =       1 / 0.18_215 * latents
		lowercase	    =       self.vae.decode(SCREAMING_SNAKE_CASE__			).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(SCREAMING_SNAKE_CASE__			)						,       return_tensors='pt'			).to(
			    self.device			)
			lowercase    ,   lowercase	    =       self.safety_checker(
			    images=SCREAMING_SNAKE_CASE__						,       clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype			)			)
		else:
			lowercase	    =       None

		if output_type == "pil":
			lowercase	    =       self.numpy_to_pil(SCREAMING_SNAKE_CASE__			)

		if not return_dict:
			return (image, has_nsfw_concept)

		return StableDiffusionPipelineOutput(images=SCREAMING_SNAKE_CASE__						,       nsfw_content_detected=SCREAMING_SNAKE_CASE__			)



from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
 import PIL


_snake_case      = logging.get_logger(__name__)

def  __lowerCamelCase					( _lowercase						,     _lowercase			)       ->      Optional[Any]:
     UpperCamelCase						    =					b.T
     UpperCamelCase						    =					np.sum(np.square(_lowercase			)						,     axis=1			)
     UpperCamelCase						    =					np.sum(np.square(_lowercase			)						,     axis=0			)
     UpperCamelCase						    =					np.matmul(_lowercase						,     _lowercase			)
     UpperCamelCase						    =					aa[:, None] - 2 * ab + ba[None, :]
     return d

def  __lowerCamelCase					( _lowercase						,     _lowercase			)       ->      List[Any]:
     UpperCamelCase						    =					x.reshape(-1						,     3			)
     UpperCamelCase						    =					squared_euclidean_distance(_lowercase						,     _lowercase			)
     return np.argmin(_lowercase						,     axis=1			)







class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					Optional[Any]       						=["pixel_values"]

 def __init__(						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Optional[Union[List[List[int]], np.ndarray]] = None					, SCREAMING_SNAKE_CASE__			:			bool = True					, SCREAMING_SNAKE_CASE__			:			Dict[str, int] = None					, SCREAMING_SNAKE_CASE__			:			PILImageResampling = PILImageResampling.BILINEAR					, SCREAMING_SNAKE_CASE__			:			bool = True					, SCREAMING_SNAKE_CASE__			:			bool = True					, **SCREAMING_SNAKE_CASE__			:			Optional[int]					, ):






      """simple docstring"""



      super().__init__(**SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					size if size is not None else {'height': 2_56, 'width': 2_56}
      UpperCamelCase						    =					get_size_dict(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					np.array(SCREAMING_SNAKE_CASE__   ) if clusters is not None else None
      UpperCamelCase						    =					do_resize
      UpperCamelCase						    =					size
      UpperCamelCase						    =					resample
      UpperCamelCase						    =					do_normalize
      UpperCamelCase						    =					do_color_quantize

 def      __lowerCAmelCase  (						self			:			int					, SCREAMING_SNAKE_CASE__			:			np.ndarray					, SCREAMING_SNAKE_CASE__			:			Dict[str, int]					, SCREAMING_SNAKE_CASE__			:			PILImageResampling = PILImageResampling.BILINEAR					, SCREAMING_SNAKE_CASE__			:			Optional[Union[str, ChannelDimension]] = None					, **SCREAMING_SNAKE_CASE__			:			List[str]					, ):






      """simple docstring"""



      UpperCamelCase						    =					get_size_dict(SCREAMING_SNAKE_CASE__   )
      if "height" not in size or "width" not in size:
           raise ValueError(F'Size dictionary must contain both height and width keys. Got {size.keys()}'   )
      return resize(
          SCREAMING_SNAKE_CASE__					, size=(size['height'], size['width'])					, resample=SCREAMING_SNAKE_CASE__					, data_format=SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			np.ndarray					, SCREAMING_SNAKE_CASE__			:			Optional[Union[str, ChannelDimension]] = None					, ):






      """simple docstring"""



      UpperCamelCase						    =					rescale(image=SCREAMING_SNAKE_CASE__					, scale=1 / 127.5					, data_format=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					image - 1
      return image






 def      __lowerCAmelCase  (						self			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			ImageInput					, SCREAMING_SNAKE_CASE__			:			bool = None					, SCREAMING_SNAKE_CASE__			:			Dict[str, int] = None					, SCREAMING_SNAKE_CASE__			:			PILImageResampling = None					, SCREAMING_SNAKE_CASE__			:			bool = None					, SCREAMING_SNAKE_CASE__			:			Optional[bool] = None					, SCREAMING_SNAKE_CASE__			:			Optional[Union[List[List[int]], np.ndarray]] = None					, SCREAMING_SNAKE_CASE__			:			Optional[Union[str, TensorType]] = None					, SCREAMING_SNAKE_CASE__			:			Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST					, **SCREAMING_SNAKE_CASE__			:			str					, ):






      """simple docstring"""



      UpperCamelCase						    =					do_resize if do_resize is not None else self.do_resize
      UpperCamelCase						    =					size if size is not None else self.size
      UpperCamelCase						    =					get_size_dict(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					resample if resample is not None else self.resample
      UpperCamelCase						    =					do_normalize if do_normalize is not None else self.do_normalize
      UpperCamelCase						    =					do_color_quantize if do_color_quantize is not None else self.do_color_quantize
      UpperCamelCase						    =					clusters if clusters is not None else self.clusters
      UpperCamelCase						    =					np.array(SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					make_list_of_images(SCREAMING_SNAKE_CASE__   )

      if not valid_images(SCREAMING_SNAKE_CASE__   ):
           raise ValueError(
               'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
               'torch.Tensor, tf.Tensor or jax.ndarray.'   )

      if do_resize and size is None or resample is None:
           raise ValueError('Size and resample must be specified if do_resize is True.'   )

      if do_color_quantize and clusters is None:
           raise ValueError('Clusters must be specified if do_color_quantize is True.'   )

      # All transformations expect numpy arrays.
      UpperCamelCase						    =					[to_numpy_array(SCREAMING_SNAKE_CASE__   ) for image in images]

      if do_resize:
           UpperCamelCase						    =					[self.resize(image=SCREAMING_SNAKE_CASE__					, size=SCREAMING_SNAKE_CASE__					, resample=SCREAMING_SNAKE_CASE__   ) for image in images]

      if do_normalize:
           UpperCamelCase						    =					[self.normalize(image=SCREAMING_SNAKE_CASE__   ) for image in images]

      if do_color_quantize:
           UpperCamelCase						    =					[to_channel_dimension_format(SCREAMING_SNAKE_CASE__					, ChannelDimension.LAST   ) for image in images]
           # color quantize from (batch_size, height, width, 3) to (batch_size, height, width)
           UpperCamelCase						    =					np.array(SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					color_quantize(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   ).reshape(images.shape[:-1]   )

           # flatten to (batch_size, height*width)
           UpperCamelCase						    =					images.shape[0]
           UpperCamelCase						    =					images.reshape(SCREAMING_SNAKE_CASE__					, -1   )

           # We need to convert back to a list of images to keep consistent behaviour across processors.
           UpperCamelCase						    =					list(SCREAMING_SNAKE_CASE__   )
      else:
           UpperCamelCase						    =					[to_channel_dimension_format(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   ) for image in images]

      UpperCamelCase						    =					{'input_ids': images}
      return BatchFeature(data=SCREAMING_SNAKE_CASE__					, tensor_type=SCREAMING_SNAKE_CASE__   )


