



"""simple docstring"""







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       UpperCamelCase__			(   _lowerCAmelCase      ):

			"""simple docstring"""


			def __init__(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	)      ->	Union[str, Any]:
				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   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__ = "auto"		)      ->	int:
				if slice_size == "auto":
					# half the attention head size is usually a good trade-off between
					# speed and memory
					A__											=					self.unet.config.attention_head_dim // 2
				self.unet.set_attention_slice(SCREAMING_SNAKE_CASE__		)


			def   snake_case__	(      self		)      ->	Union[str, Any]:
				self.enable_attention_slicing(SCREAMING_SNAKE_CASE__		)







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

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

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

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

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

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

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

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

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

					# predict the noise residual
					A__											=					self.unet(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	encoder_hidden_states=SCREAMING_SNAKE_CASE__		).sample

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

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

				A__											=					1 / 0.1_8_2_1_5 * latents
				A__											=					self.vae.decode(SCREAMING_SNAKE_CASE__		).sample

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

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

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

				if output_type == "pil":
					A__											=					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__		)











'''simple docstring'''




import torch

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest




class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	(UnCLIPScheduler,)


					def lowerCamelCase_			(						self				:			List[str]      ,						**lowerCamelCase_				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			{
							    """num_train_timesteps""": 10_00,
							    """variance_type""": """fixed_small_log""",
							    """clip_sample""": True,
							    """clip_sample_range""": 1.0,
							    """prediction_type""": """epsilon""",
							}

							config.update(**lowerCamelCase_						)
							return config


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for timesteps in [1, 5, 1_00, 10_00]:
									self.check_over_configs(num_train_timesteps=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							for variance in ["fixed_small_log", "learned_range"]:
									self.check_over_configs(variance_type=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for clip_sample in [True, False]:
									self.check_over_configs(clip_sample=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			Union[str, Any]						):

							'''simple docstring'''



							for clip_sample_range in [1, 5, 10, 20]:
									self.check_over_configs(clip_sample_range=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for prediction_type in ["epsilon", "sample"]:
									self.check_over_configs(prediction_type=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for time_step in [0, 5_00, 9_99]:
									for prev_timestep in [None, 5, 1_00, 2_50, 5_00, 7_50]:
											if prev_timestep is not None and prev_timestep >= time_step:
													continue

											self.check_over_forward(time_step=lowerCamelCase_      ,						prev_timestep=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				int		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				Dict		=			self.get_scheduler_config(variance_type="""fixed_small_log"""						)
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			scheduler_class(**lowerCamelCase_						)

							assert torch.sum(torch.abs(scheduler._get_variance(0						) - 1.0000e-10						)						) < 1e-5
							assert torch.sum(torch.abs(scheduler._get_variance(4_87						) - 0.0_549_625						)						) < 1e-5
							assert torch.sum(torch.abs(scheduler._get_variance(9_99						) - 0.9_994_987						)						) < 1e-5


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.get_scheduler_config(variance_type="""learned_range"""						)
							SCREAMING_SNAKE_CASE :				List[str]		=			scheduler_class(**lowerCamelCase_						)

							SCREAMING_SNAKE_CASE :				List[str]		=			0.5

							assert scheduler._get_variance(1      ,						predicted_variance=lowerCamelCase_						) - -10.1_712_790 < 1e-5
							assert scheduler._get_variance(4_87      ,						predicted_variance=lowerCamelCase_						) - -5.7_998_052 < 1e-5
							assert scheduler._get_variance(9_99      ,						predicted_variance=lowerCamelCase_						) - -0.0_010_011 < 1e-5


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				int		=			self.get_scheduler_config()
							SCREAMING_SNAKE_CASE :				List[Any]		=			scheduler_class(**lowerCamelCase_						)

							SCREAMING_SNAKE_CASE :				Optional[Any]		=			scheduler.timesteps

							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.dummy_model()
							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.dummy_sample_deter
							SCREAMING_SNAKE_CASE :				str		=			torch.manual_seed(0						)

							for i, t in enumerate(lowerCamelCase_						):
									# 1. predict noise residual
									SCREAMING_SNAKE_CASE :				List[str]		=			model(lowerCamelCase_      ,						lowerCamelCase_						)

									# 2. predict previous mean of sample x_t-1
									SCREAMING_SNAKE_CASE :				Optional[int]		=			scheduler.step(lowerCamelCase_      ,						lowerCamelCase_      ,						lowerCamelCase_      ,						generator=lowerCamelCase_						).prev_sample

									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			pred_prev_sample

							SCREAMING_SNAKE_CASE :				Optional[int]		=			torch.sum(torch.abs(lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE :				int		=			torch.mean(torch.abs(lowerCamelCase_						)						)

							assert abs(result_sum.item() - 252.2_682_495						) < 1e-2
							assert abs(result_mean.item() - 0.3_284_743						) < 1e-3


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				Dict		=			self.get_scheduler_config()
							SCREAMING_SNAKE_CASE :				List[str]		=			scheduler_class(**lowerCamelCase_						)

							scheduler.set_timesteps(25						)

							SCREAMING_SNAKE_CASE :				List[str]		=			scheduler.timesteps

							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.dummy_model()
							SCREAMING_SNAKE_CASE :				Tuple		=			self.dummy_sample_deter
							SCREAMING_SNAKE_CASE :				Any		=			torch.manual_seed(0						)

							for i, t in enumerate(lowerCamelCase_						):
									# 1. predict noise residual
									SCREAMING_SNAKE_CASE :				int		=			model(lowerCamelCase_      ,						lowerCamelCase_						)

									if i + 1 == timesteps.shape[0]:
											SCREAMING_SNAKE_CASE :				Union[str, Any]		=			None
									else:
											SCREAMING_SNAKE_CASE :				int		=			timesteps[i + 1]

									# 2. predict previous mean of sample x_t-1
									SCREAMING_SNAKE_CASE :				int		=			scheduler.step(
									    lowerCamelCase_      ,						lowerCamelCase_      ,						lowerCamelCase_      ,						prev_timestep=lowerCamelCase_      ,						generator=lowerCamelCase_						).prev_sample

									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			pred_prev_sample

							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torch.sum(torch.abs(lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torch.mean(torch.abs(lowerCamelCase_						)						)

							assert abs(result_sum.item() - 258.2_044_983						) < 1e-2
							assert abs(result_mean.item() - 0.3_362_038						) < 1e-3


					def lowerCamelCase_			(						self				:			Dict						):

							'''simple docstring'''



							pass



					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							pass



