


from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


_lowerCamelCase     =  logging.get_logger(__name__)  # pylint: disable=invalid-name

_lowerCamelCase     =  '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")\n        >>> pipe_prior.to("cuda")\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")\n        >>> pipe.to("cuda")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save("cat.png")\n        ```\n'


def __lowercase     (      lowerCamelCase_  :      Optional[int]						,					lowerCamelCase_  :      Union[str, Any]						,					lowerCamelCase_  :      int=8  ):
				SCREAMING_SNAKE_CASE__     							=				height // scale_factor**2
				if height % scale_factor**2 != 0:
								new_height += 1
				SCREAMING_SNAKE_CASE__     							=				width // scale_factor**2
				if width % scale_factor**2 != 0:
								new_width += 1
				return new_height * scale_factor, new_width * scale_factor


class 			lowerCamelCase_		(			__SCREAMING_SNAKE_CASE   ):





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		):
								super().__init__()

								self.register_modules(
								    unet=_a       ,		scheduler=_a       ,		movq=_a       ,		)
								SCREAMING_SNAKE_CASE__     							=				2 ** (len(self.movq.config.block_out_channels     ) - 1)



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):




								if latents is None:
												SCREAMING_SNAKE_CASE__     							=				randn_tensor(_a       ,		generator=_a       ,		device=_a       ,		dtype=_a     )
								else:
												if latents.shape != shape:
																raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}'''     )
												SCREAMING_SNAKE_CASE__     							=				latents.to(_a     )

								SCREAMING_SNAKE_CASE__     							=				latents * scheduler.init_noise_sigma
								return latents



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__=0     ):




								if is_accelerate_available():
												from accelerate import cpu_offload
								else:
												raise ImportError("Please install accelerate via `pip install accelerate`"     )

								SCREAMING_SNAKE_CASE__     							=				torch.device(f'''cuda:{gpu_id}'''     )

								SCREAMING_SNAKE_CASE__     							=				[
								    self.unet,
								    self.movq,
								]
								for cpu_offloaded_model in models:
												if cpu_offloaded_model is not None:
																cpu_offload(_a       ,		_a     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__=0     ):




								if is_accelerate_available() and is_accelerate_version(">="       ,		"0.17.0.dev0"     ):
												from accelerate import cpu_offload_with_hook
								else:
												raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."     )

								SCREAMING_SNAKE_CASE__     							=				torch.device(f'''cuda:{gpu_id}'''     )

								if self.device.type != "cpu":
												self.to("cpu"       ,		silence_dtype_warnings=_a     )
												torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

								SCREAMING_SNAKE_CASE__     							=				None
								for cpu_offloaded_model in [self.unet, self.movq]:
												SCREAMING_SNAKE_CASE__     							=				cpu_offload_with_hook(_a       ,		_a       ,		prev_module_hook=_a     )

								# We'll offload the last model manually.
								SCREAMING_SNAKE_CASE__     							=				hook



				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def   lowerCAmelCase__		(						self     ):




								if not hasattr(self.unet       ,		"_hf_hook"     ):
												return self.device
								for module in self.unet.modules():
												if (
												    hasattr(_a       ,		"_hf_hook"     )
												    and hasattr(module._hf_hook       ,		"execution_device"     )
												    and module._hf_hook.execution_device is not None
												):
																return torch.device(module._hf_hook.execution_device     )
								return self.device






				@torch.no_grad()
				@replace_example_docstring(_a     )
				def __call__(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = 512       ,		UpperCAmelCase__ = 512       ,		UpperCAmelCase__ = 100       ,		UpperCAmelCase__ = 4.0       ,		UpperCAmelCase__ = 1       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = "pil"       ,		UpperCAmelCase__ = True       ,		):
								SCREAMING_SNAKE_CASE__     							=				self._execution_device

								SCREAMING_SNAKE_CASE__     							=				guidance_scale > 1.0

								if isinstance(_a       ,		_a     ):
												SCREAMING_SNAKE_CASE__     							=				torch.cat(_a       ,		dim=0     )
								SCREAMING_SNAKE_CASE__     							=				image_embeds.shape[0] * num_images_per_prompt
								if isinstance(_a       ,		_a     ):
												SCREAMING_SNAKE_CASE__     							=				torch.cat(_a       ,		dim=0     )

								if do_classifier_free_guidance:
												SCREAMING_SNAKE_CASE__     							=				image_embeds.repeat_interleave(_a       ,		dim=0     )
												SCREAMING_SNAKE_CASE__     							=				negative_image_embeds.repeat_interleave(_a       ,		dim=0     )

								SCREAMING_SNAKE_CASE__     							=				torch.cat([negative_image_embeds, image_embeds]       ,		dim=0     ).to(dtype=self.unet.dtype       ,		device=_a     )

								self.scheduler.set_timesteps(_a       ,		device=_a     )
								SCREAMING_SNAKE_CASE__     							=				self.scheduler.timesteps

								SCREAMING_SNAKE_CASE__     							=				self.unet.config.in_channels

								SCREAMING_SNAKE_CASE__     							=				downscale_height_and_width(_a       ,		_a       ,		self.movq_scale_factor     )

								# create initial latent
								SCREAMING_SNAKE_CASE__     							=				self.prepare_latents(
								    (batch_size, num_channels_latents, height, width)       ,		image_embeds.dtype       ,		_a       ,		_a       ,		_a       ,		self.scheduler       ,		)

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

												SCREAMING_SNAKE_CASE__     							=				{"""image_embeds""": image_embeds}
												SCREAMING_SNAKE_CASE__     							=				self.unet(
												    sample=_a       ,		timestep=_a       ,		encoder_hidden_states=_a       ,		added_cond_kwargs=_a       ,		return_dict=_a       ,		)[0]

												if do_classifier_free_guidance:
																SCREAMING_SNAKE_CASE__     							=				noise_pred.split(latents.shape[1]       ,		dim=1     )
																SCREAMING_SNAKE_CASE__     							=				noise_pred.chunk(2     )
																SCREAMING_SNAKE_CASE__     							=				variance_pred.chunk(2     )
																SCREAMING_SNAKE_CASE__     							=				noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																SCREAMING_SNAKE_CASE__     							=				torch.cat([noise_pred, variance_pred_text]       ,		dim=1     )

												if not (
												    hasattr(self.scheduler.config       ,		"variance_type"     )
												    and self.scheduler.config.variance_type in ["learned", "learned_range"]
												):
																SCREAMING_SNAKE_CASE__     							=				noise_pred.split(latents.shape[1]       ,		dim=1     )

												# compute the previous noisy sample x_t -> x_t-1
												SCREAMING_SNAKE_CASE__     							=				self.scheduler.step(
												    _a       ,		_a       ,		_a       ,		generator=_a       ,		)[0]
								# post-processing
								SCREAMING_SNAKE_CASE__     							=				self.movq.decode(_a       ,		force_not_quantize=_a     )["""sample"""]

								if output_type not in ["pt", "np", "pil"]:
												raise ValueError(f'''Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'''     )

								if output_type in ["np", "pil"]:
												SCREAMING_SNAKE_CASE__     							=				image * 0.5 + 0.5
												SCREAMING_SNAKE_CASE__     							=				image.clamp(0       ,		1     )
												SCREAMING_SNAKE_CASE__     							=				image.cpu().permute(0       ,		2       ,		3       ,		1     ).float().numpy()

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

								if not return_dict:
												return (image,)

								return ImagePipelineOutput(images=_a     )




"""simple docstring"""






# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from typing import Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import SchedulerMixin


class 			lowerCamelCase_		(			lowercase       ,  lowercase   ):





				"""simple docstring"""







				_lowerCAmelCase      :      List[Any]      							=					1



				@register_to_config
				def __init__(						self       ,		UpperCAmelCase__=2000       ,		UpperCAmelCase__=0.1       ,		UpperCAmelCase__=20       ,		UpperCAmelCase__=1e-3     ):
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				None



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None     ):
								SCREAMING_SNAKE_CASE__     							=				torch.linspace(1       ,		self.config.sampling_eps       ,		UpperCAmelCase__       ,		device=UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__=None     ):
								if self.timesteps is None:
												raise ValueError(
												    "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"     )

								# TODO(Patrick) better comments + non-PyTorch
								# postprocess model score
								SCREAMING_SNAKE_CASE__     							=				(
								    -0.25 * t**2 * (self.config.beta_max - self.config.beta_min) - 0.5 * t * self.config.beta_min
								)
								SCREAMING_SNAKE_CASE__     							=				torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff     )     )
								SCREAMING_SNAKE_CASE__     							=				std.flatten()
								while len(std.shape     ) < len(score.shape     ):
												SCREAMING_SNAKE_CASE__     							=				std.unsqueeze(-1     )
								SCREAMING_SNAKE_CASE__     							=				-score / std

								# compute
								SCREAMING_SNAKE_CASE__     							=				-1.0 / len(self.timesteps     )

								SCREAMING_SNAKE_CASE__     							=				self.config.beta_min + t * (self.config.beta_max - self.config.beta_min)
								SCREAMING_SNAKE_CASE__     							=				beta_t.flatten()
								while len(beta_t.shape     ) < len(x.shape     ):
												SCREAMING_SNAKE_CASE__     							=				beta_t.unsqueeze(-1     )
								SCREAMING_SNAKE_CASE__     							=				-0.5 * beta_t * x

								SCREAMING_SNAKE_CASE__     							=				torch.sqrt(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				drift - diffusion**2 * score
								SCREAMING_SNAKE_CASE__     							=				x + drift * dt

								# add noise
								SCREAMING_SNAKE_CASE__     							=				randn_tensor(x.shape       ,		layout=x.layout       ,		generator=UpperCAmelCase__       ,		device=x.device       ,		dtype=x.dtype     )
								SCREAMING_SNAKE_CASE__     							=				x_mean + diffusion * math.sqrt(-dt     ) * noise

								return x, x_mean






				def __len__(						self     ):
								return self.config.num_train_timesteps
