





from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

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,
)


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

__lowercase       :str       =  "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline\n        >>> from diffusers.utils import load_image\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"A red cartoon frog, 4k\"\n        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)\n\n        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-decoder\", torch_dtype=torch.float16\n        ... )\n        >>> pipe.to(\"cuda\")\n\n        >>> init_image = load_image(\n        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"\n        ...     \"/kandinsky/frog.png\"\n        ... )\n\n        >>> image = pipe(\n        ...     image=init_image,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ...     strength=0.2,\n        ... ).images\n\n        >>> image[0].save(\"red_frog.png\")\n        ```\n"







def       UpperCAmelCase							(     _lowerCamelCase      :  int			,   _lowerCamelCase      :  Optional[Any]			,   _lowerCamelCase      :  Any=8     ):





							'''simple docstring'''


							SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					height // scale_factor**2
							if height % scale_factor**2 != 0:
														new_height += 1
							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					width // scale_factor**2
							if width % scale_factor**2 != 0:
														new_width += 1
							return new_height * scale_factor, new_width * scale_factor







def       UpperCAmelCase							(     _lowerCamelCase      :  Any			,   _lowerCamelCase      :  Tuple=512			,   _lowerCamelCase      :  Union[str, Any]=512     ):





							'''simple docstring'''


							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					pil_image.resize((w, h)			,   resample=Image.BICUBIC			,   reducing_gap=1     )
							SCREAMING_SNAKE_CASE__     :				Tuple   =					np.array(pil_image.convert("RGB"     )     )
							SCREAMING_SNAKE_CASE__     :				List[str]   =					arr.astype(np.floataa     ) / 127.5 - 1
							SCREAMING_SNAKE_CASE__     :				int   =					np.transpose(__lowerCAmelCase			,   [2, 0, 1]     )
							SCREAMING_SNAKE_CASE__     :				str   =					torch.from_numpy(__lowerCAmelCase     ).unsqueeze(0     )
							return image



class    _a     (				UpperCAmelCase__							):






							"""simple docstring"""

							def __init__(      self				:     int							,	a				:     UNetaDConditionModel							,	a				:     DDPMScheduler							,	a				:     VQModel							,	)	->Union[str, Any]:
														super().__init__()

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

							def      A_  (      self				:     Optional[Any]							,	a				:     int							,	a				:     Tuple							,	a				:     Union[str, Any]    )	->Union[str, Any]:
														SCREAMING_SNAKE_CASE__     :				int   =					min(int(num_inference_steps * strength    )							,	lowerCamelCase__    )

														SCREAMING_SNAKE_CASE__     :				Optional[int]   =					max(num_inference_steps - init_timestep							,	0    )
														SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					self.scheduler.timesteps[t_start:]

														return timesteps, num_inference_steps - t_start

							def      A_  (      self				:     List[Any]							,	a				:     Tuple							,	a				:     Dict							,	a				:     Optional[Any]							,	a				:     Dict							,	a				:     Union[str, Any]							,	a				:     List[str]							,	a				:     Dict=None    )	->Any:







														if not isinstance(lowerCamelCase__							,	(torch.Tensor, PIL.Image.Image, list)    ):
																					raise ValueError(
																					    f"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(lowerCamelCase__    )}"""    )

														SCREAMING_SNAKE_CASE__     :				str   =					image.to(device=lowerCamelCase__							,	dtype=lowerCamelCase__    )

														SCREAMING_SNAKE_CASE__     :				Tuple   =					batch_size * num_images_per_prompt

														if image.shape[1] == 4:
																					SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					image

														else:
																					if isinstance(lowerCamelCase__							,	lowerCamelCase__    ) and len(lowerCamelCase__    ) != batch_size:
																												raise ValueError(
																												    f"""You have passed a list of generators of length {len(lowerCamelCase__    )}, but requested an effective batch"""
																												    f""" size of {batch_size}. Make sure the batch size matches the length of the generators."""    )

																					elif isinstance(lowerCamelCase__							,	lowerCamelCase__    ):
																												SCREAMING_SNAKE_CASE__     :				List[Any]   =					[
																												    self.movq.encode(image[i : i + 1]    ).latent_dist.sample(generator[i]    ) for i in range(lowerCamelCase__    )
																												]
																												SCREAMING_SNAKE_CASE__     :				str   =					torch.cat(lowerCamelCase__							,	dim=0    )
																					else:
																												SCREAMING_SNAKE_CASE__     :				List[str]   =					self.movq.encode(lowerCamelCase__    ).latent_dist.sample(lowerCamelCase__    )

																					SCREAMING_SNAKE_CASE__     :				List[str]   =					self.movq.config.scaling_factor * init_latents

														SCREAMING_SNAKE_CASE__     :				List[Any]   =					torch.cat([init_latents]							,	dim=0    )

														SCREAMING_SNAKE_CASE__     :				int   =					init_latents.shape
														SCREAMING_SNAKE_CASE__     :				List[Any]   =					randn_tensor(lowerCamelCase__							,	generator=lowerCamelCase__							,	device=lowerCamelCase__							,	dtype=lowerCamelCase__    )

														# get latents
														SCREAMING_SNAKE_CASE__     :				Optional[int]   =					self.scheduler.add_noise(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__    )

														SCREAMING_SNAKE_CASE__     :				Dict   =					init_latents

														return latents

							def      A_  (      self				:     Any							,	a				:     List[Any]=0    )	->str:







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

														SCREAMING_SNAKE_CASE__     :				List[str]   =					torch.device(f"""cuda:{gpu_id}"""    )

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

							def      A_  (      self				:     Any							,	a				:     Tuple=0    )	->List[Any]:







														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__     :				Dict   =					torch.device(f"""cuda:{gpu_id}"""    )

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

														SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					None
														for cpu_offloaded_model in [self.unet, self.movq]:
																					SCREAMING_SNAKE_CASE__     :				Tuple   =					cpu_offload_with_hook(lowerCamelCase__							,	lowerCamelCase__							,	prev_module_hook=lowerCamelCase__    )

														# We'll offload the last model manually.
														SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					hook

							@property
							# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
							def      A_  (      self				:     Any    )	->Optional[Any]:







														if not hasattr(self.unet							,	"_hf_hook"    ):
																					return self.device
														for module in self.unet.modules():
																					if (
																					    hasattr(lowerCamelCase__							,	"_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(lowerCamelCase__    )
							def __call__(      self				:     str							,	a				:     Union[torch.FloatTensor, List[torch.FloatTensor]]							,	a				:     Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]							,	a				:     Union[torch.FloatTensor, List[torch.FloatTensor]]							,	a				:     int = 5_12							,	a				:     int = 5_12							,	a				:     int = 1_00							,	a				:     float = 4.0							,	a				:     float = 0.3							,	a				:     int = 1							,	a				:     Optional[Union[torch.Generator, List[torch.Generator]]] = None							,	a				:     Optional[str] = "pil"							,	a				:     bool = True							,	)	->str:
														SCREAMING_SNAKE_CASE__     :				Any   =					self._execution_device

														SCREAMING_SNAKE_CASE__     :				str   =					guidance_scale > 1.0

														if isinstance(lowerCamelCase__							,	lowerCamelCase__    ):
																					SCREAMING_SNAKE_CASE__     :				List[str]   =					torch.cat(lowerCamelCase__							,	dim=0    )
														SCREAMING_SNAKE_CASE__     :				Any   =					image_embeds.shape[0]
														if isinstance(lowerCamelCase__							,	lowerCamelCase__    ):
																					SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					torch.cat(lowerCamelCase__							,	dim=0    )

														if do_classifier_free_guidance:
																					SCREAMING_SNAKE_CASE__     :				str   =					image_embeds.repeat_interleave(lowerCamelCase__							,	dim=0    )
																					SCREAMING_SNAKE_CASE__     :				Optional[int]   =					negative_image_embeds.repeat_interleave(lowerCamelCase__							,	dim=0    )

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

														if not isinstance(lowerCamelCase__							,	lowerCamelCase__    ):
																					SCREAMING_SNAKE_CASE__     :				List[Any]   =					[image]
														if not all(isinstance(lowerCamelCase__							,	(PIL.Image.Image, torch.Tensor)    ) for i in image    ):
																					raise ValueError(
																					    f"""Input is in incorrect format: {[type(lowerCamelCase__    ) for i in image]}. Currently, we only support  PIL image and pytorch tensor"""    )

														SCREAMING_SNAKE_CASE__     :				Dict   =					torch.cat([prepare_image(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__    ) for i in image]							,	dim=0    )
														SCREAMING_SNAKE_CASE__     :				int   =					image.to(dtype=image_embeds.dtype							,	device=lowerCamelCase__    )

														SCREAMING_SNAKE_CASE__     :				Dict   =					self.movq.encode(lowerCamelCase__    )["latents"]
														SCREAMING_SNAKE_CASE__     :				List[str]   =					latents.repeat_interleave(lowerCamelCase__							,	dim=0    )
														self.scheduler.set_timesteps(lowerCamelCase__							,	device=lowerCamelCase__    )
														SCREAMING_SNAKE_CASE__     :				Any   =					self.get_timesteps(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__    )
														SCREAMING_SNAKE_CASE__     :				List[str]   =					timesteps[:1].repeat(batch_size * num_images_per_prompt    )
														SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					downscale_height_and_width(lowerCamelCase__							,	lowerCamelCase__							,	self.movq_scale_factor    )
														SCREAMING_SNAKE_CASE__     :				List[Any]   =					self.prepare_latents(
														    lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	image_embeds.dtype							,	lowerCamelCase__							,	lowerCamelCase__    )
														for i, t in enumerate(self.progress_bar(lowerCamelCase__    )    ):
																					# expand the latents if we are doing classifier free guidance
																					SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents

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

																					if do_classifier_free_guidance:
																												SCREAMING_SNAKE_CASE__     :				List[str]   =					noise_pred.split(latents.shape[1]							,	dim=1    )
																												SCREAMING_SNAKE_CASE__     :				Optional[int]   =					noise_pred.chunk(2    )
																												SCREAMING_SNAKE_CASE__     :				Tuple   =					variance_pred.chunk(2    )
																												SCREAMING_SNAKE_CASE__     :				str   =					noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																												SCREAMING_SNAKE_CASE__     :				Dict   =					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__     :				int   =					noise_pred.split(latents.shape[1]							,	dim=1    )

																					# compute the previous noisy sample x_t -> x_t-1
																					SCREAMING_SNAKE_CASE__     :				int   =					self.scheduler.step(
																					    lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	generator=lowerCamelCase__							,	)[0]

														# post-processing
														SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					self.movq.decode(lowerCamelCase__							,	force_not_quantize=lowerCamelCase__    )["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__     :				Tuple   =					image * 0.5 + 0.5
																					SCREAMING_SNAKE_CASE__     :				List[Any]   =					image.clamp(0							,	1    )
																					SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					image.cpu().permute(0							,	2							,	3							,	1    ).float().numpy()

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

														if not return_dict:
																					return (image,)

														return ImagePipelineOutput(images=lowerCamelCase__    )






# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import doctest
import sys
import warnings
from os.path import abspath, dirname, join

import _pytest

from transformers.testing_utils import HfDoctestModule, HfDocTestParser


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
__lowercase       :List[Any]       =  abspath(join(dirname(__file__), "src"))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action="ignore", category=FutureWarning)







def       UpperCAmelCase							(     _lowerCamelCase      :  int     ):





						'''simple docstring'''


						config.addinivalue_line(
						    "markers"			,   "is_pt_tf_cross_test: mark test to run only when PT and TF interactions are tested"     )
						config.addinivalue_line(
						    "markers"			,   "is_pt_flax_cross_test: mark test to run only when PT and FLAX interactions are tested"     )
						config.addinivalue_line("markers"			,   "is_pipeline_test: mark test to run only when pipelines are tested"     )
						config.addinivalue_line("markers"			,   "is_staging_test: mark test to run only in the staging environment"     )
						config.addinivalue_line("markers"			,   "accelerate_tests: mark test that require accelerate"     )
						config.addinivalue_line("markers"			,   "tool_tests: mark the tool tests that are run on their specific schedule"     )







def       UpperCAmelCase							(     _lowerCamelCase      :  str     ):





						'''simple docstring'''


						from transformers.testing_utils import pytest_addoption_shared

						pytest_addoption_shared(_lowerCamelCase     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Tuple     ):





						'''simple docstring'''


						from transformers.testing_utils import pytest_terminal_summary_main

						SCREAMING_SNAKE_CASE__     :				List[str]   =					terminalreporter.config.getoption("--make-reports"     )
						if make_reports:
												pytest_terminal_summary_main(_lowerCamelCase			,   id=_lowerCamelCase     )







def       UpperCAmelCase							(     _lowerCamelCase      :  List[Any]			,   _lowerCamelCase      :  Dict     ):





						'''simple docstring'''


						if exitstatus == 5:
												SCREAMING_SNAKE_CASE__     :				List[str]   =					0


# Doctest custom flag to ignore output.
__lowercase       :Optional[Any]       =  doctest.register_optionflag("IGNORE_RESULT")

__lowercase       :Dict       =  doctest.OutputChecker



class    _a     (				lowercase__							):






							"""simple docstring"""

							def      A_  (      self				:     Dict							,	a				:     List[str]							,	a				:     Dict							,	a				:     int    )	->Optional[Any]:
													if IGNORE_RESULT & optionflags:
																			return True
													return OutputChecker.check_output(self							,	a							,	a							,	a    )







__lowercase       :Any       =  CustomOutputChecker
__lowercase       :Any       =  HfDoctestModule



__lowercase       :int       =  HfDocTestParser