





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


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

UpperCAmelCase__						:			Optional[int]									='''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

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





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase=8				)						->  List[str]:
				lowerCamelCase	 =height // scale_factor**2
				if height % scale_factor**2 != 0:
								new_height += 1
				lowerCamelCase	 =width // scale_factor**2
				if width % scale_factor**2 != 0:
								new_width += 1
				return new_height * scale_factor, new_width * scale_factor




class       __A     (						a__    ):


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

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


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
								if latents is None:
												lowerCamelCase	 =randn_tensor(lowerCAmelCase__   , generator=lowerCAmelCase__   , device=lowerCAmelCase__   , dtype=lowerCAmelCase__							)
								else:
												if latents.shape != shape:
																raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""							)
												lowerCamelCase	 =latents.to(lowerCAmelCase__							)

								lowerCamelCase	 =latents * scheduler.init_noise_sigma
								return latents


				def _snake_case       (		self   , UpperCAmelCase_=0							):
								if is_accelerate_available():
												from accelerate import cpu_offload
								else:
												raise ImportError("""Please install accelerate via `pip install accelerate`"""							)

								lowerCamelCase	 =torch.device(f"""cuda:{gpu_id}"""							)

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


				def _snake_case       (		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."""							)

								lowerCamelCase	 =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)

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

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


				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def _snake_case       (		self							):
								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   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 512   , UpperCAmelCase_ = 512   , UpperCAmelCase_ = 100   , UpperCAmelCase_ = 4.0   , UpperCAmelCase_ = 1   , UpperCAmelCase_ = None   , UpperCAmelCase_ = None   , UpperCAmelCase_ = "pil"   , UpperCAmelCase_ = True   , ):
								lowerCamelCase	 =self._execution_device

								lowerCamelCase	 =guidance_scale > 1.0

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

								if do_classifier_free_guidance:
												lowerCamelCase	 =image_embeds.repeat_interleave(lowerCAmelCase__   , dim=0							)
												lowerCamelCase	 =negative_image_embeds.repeat_interleave(lowerCAmelCase__   , dim=0							)

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

								self.scheduler.set_timesteps(lowerCAmelCase__   , device=lowerCAmelCase__							)
								lowerCamelCase	 =self.scheduler.timesteps

								lowerCamelCase	 =self.unet.config.in_channels

								lowerCamelCase	 =downscale_height_and_width(lowerCAmelCase__   , lowerCAmelCase__   , self.movq_scale_factor							)

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

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

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

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

												# compute the previous noisy sample x_t -> x_t-1
												lowerCamelCase	 =self.scheduler.step(
												    lowerCAmelCase__   , lowerCAmelCase__   , lowerCAmelCase__   , generator=lowerCAmelCase__   , )[0]
								# post-processing
								lowerCamelCase	 =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"]:
												lowerCamelCase	 =image * 0.5 + 0.5
												lowerCamelCase	 =image.clamp(0   , 1							)
												lowerCamelCase	 =image.cpu().permute(0   , 2   , 3   , 1							).float().numpy()

								if output_type == "pil":
												lowerCamelCase	 =self.numpy_to_pil(lowerCAmelCase__							)

								if not return_dict:
												return (image,)

								return ImagePipelineOutput(images=lowerCAmelCase__							)












from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
							if not (is_transformers_available() and is_torch_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from ...utils.dummy_torch_and_transformers_objects import KandinskyPipeline, KandinskyPriorPipeline
else:
							from .pipeline_kandinsky import KandinskyPipeline
							from .pipeline_kandinsky_imgaimg import KandinskyImgaImgPipeline
							from .pipeline_kandinsky_inpaint import KandinskyInpaintPipeline
							from .pipeline_kandinsky_prior import KandinskyPriorPipeline, KandinskyPriorPipelineOutput
							from .text_encoder import MultilingualCLIP





