





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


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

_SCREAMING_SNAKE_CASE 		=						'\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 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a=8				):
	snake_case_						:     Union[str, Any]     =		height // scale_factor**2
	if height % scale_factor**2 != 0:
		new_height += 1
	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 							SCREAMING_SNAKE_CASE__  (      __a				,	__a=5_12				,	__a=5_12				):
	snake_case_						:     Any     =		pil_image.resize((w, h)				,	resample=Image.BICUBIC				,	reducing_gap=1				)
	snake_case_						:     Optional[Any]     =		np.array(pil_image.convert('RGB'				)				)
	snake_case_						:     Dict     =		arr.astype(np.floataa				) / 127.5 - 1
	snake_case_						:     str     =		np.transpose(lowerCAmelCase__				,	[2, 0, 1]				)
	snake_case_						:     Optional[Any]     =		torch.from_numpy(lowerCAmelCase__				).unsqueeze(0				)
	return image


class SCREAMING_SNAKE_CASE_  (   _lowerCAmelCase   ):






						def __init__( self : List[str]						,				_A : int						,				_A : int						,				_A : Any						,				)							->      List[Any]:







							"""simple docstring"""


							super().__init__()

							self.register_modules(
							    unet=SCREAMING_SNAKE_CASE_						,				scheduler=SCREAMING_SNAKE_CASE_						,				movq=SCREAMING_SNAKE_CASE_						,				)
							snake_case_						:     List[Any]     =		2 ** (len(self.movq.config.block_out_channels		) - 1)






						def UpperCAmelCase_			( self : Tuple						,				_A : Optional[Any]						,				_A : Tuple						,				_A : List[Any]		)							->      List[Any]:







							"""simple docstring"""


							snake_case_						:     str     =		min(int(num_inference_steps * strength		)						,				SCREAMING_SNAKE_CASE_		)

							snake_case_						:     List[str]     =		max(num_inference_steps - init_timestep						,				0		)
							snake_case_						:     List[Any]     =		self.scheduler.timesteps[t_start:]

							return timesteps, num_inference_steps - t_start






						def UpperCAmelCase_			( self : Optional[int]						,				_A : int						,				_A : Dict						,				_A : Optional[Any]						,				_A : List[Any]						,				_A : Tuple						,				_A : int						,				_A : Any=None		)							->      List[str]:







							"""simple docstring"""


							if not isinstance(SCREAMING_SNAKE_CASE_						,				(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(SCREAMING_SNAKE_CASE_		)}"""		)

							snake_case_						:     str     =		image.to(device=SCREAMING_SNAKE_CASE_						,				dtype=SCREAMING_SNAKE_CASE_		)

							snake_case_						:     List[Any]     =		batch_size * num_images_per_prompt

							if image.shape[1] == 4:
								snake_case_						:     Optional[int]     =		image

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

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

								snake_case_						:     Dict     =		self.movq.config.scaling_factor * init_latents

							snake_case_						:     List[str]     =		torch.cat([init_latents]						,				dim=0		)

							snake_case_						:     Any     =		init_latents.shape
							snake_case_						:     int     =		randn_tensor(SCREAMING_SNAKE_CASE_						,				generator=SCREAMING_SNAKE_CASE_						,				device=SCREAMING_SNAKE_CASE_						,				dtype=SCREAMING_SNAKE_CASE_		)

							# get latents
							snake_case_						:     Optional[Any]     =		self.scheduler.add_noise(SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_		)

							snake_case_						:     str     =		init_latents

							return latents






						def UpperCAmelCase_			( self : int						,				_A : Union[str, Any]=0		)							->      str:







							"""simple docstring"""


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

							snake_case_						:     Any     =		torch.device(F"""cuda:{gpu_id}"""		)

							snake_case_						:     List[Any]     =		[
							    self.unet,
							    self.movq,
							]
							for cpu_offloaded_model in models:
								if cpu_offloaded_model is not None:
									cpu_offload(SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_		)






						def UpperCAmelCase_			( self : Any						,				_A : List[str]=0		)							->      Optional[int]:







							"""simple docstring"""


							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.'		)

							snake_case_						:     int     =		torch.device(F"""cuda:{gpu_id}"""		)

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

							snake_case_						:     Any     =		None
							for cpu_offloaded_model in [self.unet, self.movq]:
								snake_case_						:     Dict     =		cpu_offload_with_hook(SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_						,				prev_module_hook=SCREAMING_SNAKE_CASE_		)

							# We'll offload the last model manually.
							snake_case_						:     Dict     =		hook






						@property
						# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
						def UpperCAmelCase_			( self : Tuple		)							->      Any:







							"""simple docstring"""


							if not hasattr(self.unet						,				'_hf_hook'		):
								return self.device
							for module in self.unet.modules():
								if (
								    hasattr(SCREAMING_SNAKE_CASE_						,				'_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(SCREAMING_SNAKE_CASE_		)
						def __call__( self : Tuple						,				_A : Optional[Any]						,				_A : Union[str, Any]						,				_A : Dict						,				_A : Tuple = 512						,				_A : int = 512						,				_A : Optional[int] = 100						,				_A : Dict = 4.0						,				_A : Dict = 0.3						,				_A : int = 1						,				_A : int = None						,				_A : List[str] = "pil"						,				_A : Any = True						,				)							->      Optional[Any]:







							"""simple docstring"""


							snake_case_						:     int     =		self._execution_device

							snake_case_						:     Any     =		guidance_scale > 1.0

							if isinstance(SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_		):
								snake_case_						:     Optional[Any]     =		torch.cat(SCREAMING_SNAKE_CASE_						,				dim=0		)
							snake_case_						:     int     =		image_embeds.shape[0]
							if isinstance(SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_		):
								snake_case_						:     Dict     =		torch.cat(SCREAMING_SNAKE_CASE_						,				dim=0		)

							if do_classifier_free_guidance:
								snake_case_						:     int     =		image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_						,				dim=0		)
								snake_case_						:     Tuple     =		negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_						,				dim=0		)

							snake_case_						:     Optional[Any]     =		torch.cat([negative_image_embeds, image_embeds]						,				dim=0		).to(dtype=self.unet.dtype						,				device=SCREAMING_SNAKE_CASE_		)

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

							snake_case_						:     Optional[int]     =		torch.cat([prepare_image(SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_						,				SCREAMING_SNAKE_CASE_		) for i in image]						,				dim=0		)
							snake_case_						:     List[Any]     =		image.to(dtype=image_embeds.dtype						,				device=SCREAMING_SNAKE_CASE_		)

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

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

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

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

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

							if output_type == "pil":
								snake_case_						:     List[Any]     =		self.numpy_to_pil(SCREAMING_SNAKE_CASE_		)

							if not return_dict:
								return (image,)

							return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE_		)








_SCREAMING_SNAKE_CASE 		=						{
    """Pillow""": """Pillow""",
    """accelerate""": """accelerate>=0.11.0""",
    """compel""": """compel==0.1.8""",
    """black""": """black~=23.1""",
    """datasets""": """datasets""",
    """filelock""": """filelock""",
    """flax""": """flax>=0.4.1""",
    """hf-doc-builder""": """hf-doc-builder>=0.3.0""",
    """huggingface-hub""": """huggingface-hub>=0.13.2""",
    """requests-mock""": """requests-mock==1.10.0""",
    """importlib_metadata""": """importlib_metadata""",
    """invisible-watermark""": """invisible-watermark""",
    """isort""": """isort>=5.5.4""",
    """jax""": """jax>=0.2.8,!=0.3.2""",
    """jaxlib""": """jaxlib>=0.1.65""",
    """Jinja2""": """Jinja2""",
    """k-diffusion""": """k-diffusion>=0.0.12""",
    """torchsde""": """torchsde""",
    """note_seq""": """note_seq""",
    """librosa""": """librosa""",
    """numpy""": """numpy""",
    """omegaconf""": """omegaconf""",
    """parameterized""": """parameterized""",
    """protobuf""": """protobuf>=3.20.3,<4""",
    """pytest""": """pytest""",
    """pytest-timeout""": """pytest-timeout""",
    """pytest-xdist""": """pytest-xdist""",
    """ruff""": """ruff>=0.0.241""",
    """safetensors""": """safetensors""",
    """sentencepiece""": """sentencepiece>=0.1.91,!=0.1.92""",
    """scipy""": """scipy""",
    """onnx""": """onnx""",
    """regex""": """regex!=2019.12.17""",
    """requests""": """requests""",
    """tensorboard""": """tensorboard""",
    """torch""": """torch>=1.4""",
    """torchvision""": """torchvision""",
    """transformers""": """transformers>=4.25.1""",
    """urllib3""": """urllib3<=2.0.0""",
}

