



from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

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


UpperCAmelCase_							:      Union[str, Any]      =  logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase_							:      Tuple      =  "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")\n        >>> pipe.to(\"cuda\")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save(\"cat.png\")\n        ```\n"
def  UpperCamelCase						(       _A			:				int			,							_A			:				int			,							_A			:				List[str]=8    )->       int:





						"""simple docstring"""





						A__							=       h // scale_factor**2
						if h % scale_factor**2 != 0:
												new_h += 1
						A__							=       w // scale_factor**2
						if w % scale_factor**2 != 0:
												new_w += 1
						return new_h * scale_factor, new_w * scale_factor






class      UpperCamelCase    (						_UpperCAmelCase     ):



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

							self.register_modules(
							    text_encoder=UpperCAmelCase__			,	tokenizer=UpperCAmelCase__			,	unet=UpperCAmelCase__			,	scheduler=UpperCAmelCase__			,	movq=UpperCAmelCase__			,	)
							A__							=       2 ** (len(self.movq.config.block_out_channels		) - 1)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							if latents is None:
													A__							=       randn_tensor(UpperCAmelCase__			,	generator=UpperCAmelCase__			,	device=UpperCAmelCase__			,	dtype=UpperCAmelCase__		)
							else:
													if latents.shape != shape:
																			raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""		)
													A__							=       latents.to(UpperCAmelCase__		)

							A__							=       latents * scheduler.init_noise_sigma
							return latents



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__=None			,	):
							A__							=       len(UpperCAmelCase__		) if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		) else 1
							# get prompt text embeddings
							A__							=       self.tokenizer(
							    UpperCAmelCase__			,	padding="max_length"			,	truncation=UpperCAmelCase__			,	max_length=77			,	return_attention_mask=UpperCAmelCase__			,	add_special_tokens=UpperCAmelCase__			,	return_tensors="pt"			,	)

							A__							=       text_inputs.input_ids
							A__							=       self.tokenizer(UpperCAmelCase__			,	padding="longest"			,	return_tensors="pt"		).input_ids

							if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(UpperCAmelCase__			,	UpperCAmelCase__		):
													A__							=       self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]		)
													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.to(UpperCAmelCase__		)
							A__							=       text_inputs.attention_mask.to(UpperCAmelCase__		)

							A__						,       A__							=       self.text_encoder(
							    input_ids=UpperCAmelCase__			,	attention_mask=UpperCAmelCase__		)

							A__							=       prompt_embeds.repeat_interleave(UpperCAmelCase__			,	dim=0		)
							A__							=       text_encoder_hidden_states.repeat_interleave(UpperCAmelCase__			,	dim=0		)
							A__							=       text_mask.repeat_interleave(UpperCAmelCase__			,	dim=0		)

							if do_classifier_free_guidance:
													A__							=       42
													if negative_prompt is None:
																			A__							=       [""] * batch_size
													elif type(UpperCAmelCase__		) is not type(UpperCAmelCase__		):
																			raise TypeError(
																			    F"""`negative_prompt` should be the same type to `prompt`, but got {type(UpperCAmelCase__		)} !="""
																			    F""" {type(UpperCAmelCase__		)}."""		)
													elif isinstance(UpperCAmelCase__			,	UpperCAmelCase__		):
																			A__							=       [negative_prompt]
													elif batch_size != len(UpperCAmelCase__		):
																			raise ValueError(
																			    F"""`negative_prompt`: {negative_prompt} has batch size {len(UpperCAmelCase__		)}, 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__							=       self.tokenizer(
													    UpperCAmelCase__			,	padding="max_length"			,	max_length=77			,	truncation=UpperCAmelCase__			,	return_attention_mask=UpperCAmelCase__			,	add_special_tokens=UpperCAmelCase__			,	return_tensors="pt"			,	)
													A__							=       uncond_input.input_ids.to(UpperCAmelCase__		)
													A__							=       uncond_input.attention_mask.to(UpperCAmelCase__		)

													A__						,       A__							=       self.text_encoder(
													    input_ids=UpperCAmelCase__			,	attention_mask=UpperCAmelCase__		)

													# duplicate unconditional embeddings for each generation per prompt, using mps friendly method

													A__							=       negative_prompt_embeds.shape[1]
													A__							=       negative_prompt_embeds.repeat(1			,	UpperCAmelCase__		)
													A__							=       negative_prompt_embeds.view(batch_size * num_images_per_prompt			,	UpperCAmelCase__		)

													A__							=       uncond_text_encoder_hidden_states.shape[1]
													A__							=       uncond_text_encoder_hidden_states.repeat(1			,	UpperCAmelCase__			,	1		)
													A__							=       uncond_text_encoder_hidden_states.view(
													    batch_size * num_images_per_prompt			,	UpperCAmelCase__			,	-1		)
													A__							=       uncond_text_mask.repeat_interleave(UpperCAmelCase__			,	dim=0		)

													# done duplicates

													# 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([negative_prompt_embeds, prompt_embeds]		)
													A__							=       torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]		)

													A__							=       torch.cat([uncond_text_mask, text_mask]		)

							return prompt_embeds, text_encoder_hidden_states, text_mask



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

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

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



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

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

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

							A__							=       None
							for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
													A__						,       A__							=       cpu_offload_with_hook(UpperCAmelCase__			,	UpperCAmelCase__			,	prev_module_hook=UpperCAmelCase__		)

							if self.safety_checker is not None:
													A__						,       A__							=       cpu_offload_with_hook(self.safety_checker			,	UpperCAmelCase__			,	prev_module_hook=UpperCAmelCase__		)

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



	@property
	# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
	def 			__A							(       self		):
							if not hasattr(self.unet			,	"_hf_hook"		):
													return self.device
							for module in self.unet.modules():
													if (
													    hasattr(UpperCAmelCase__			,	"_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(UpperCAmelCase__		)
	def __call__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = 512			,	UpperCAmelCase__ = 512			,	UpperCAmelCase__ = 100			,	UpperCAmelCase__ = 4.0			,	UpperCAmelCase__ = 1			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = "pil"			,	UpperCAmelCase__ = True			,	):

							if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		):
													A__							=       1
							elif isinstance(UpperCAmelCase__			,	UpperCAmelCase__		):
													A__							=       len(UpperCAmelCase__		)
							else:
													raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(UpperCAmelCase__		)}"""		)

							A__							=       self._execution_device

							A__							=       batch_size * num_images_per_prompt
							A__							=       guidance_scale > 1.0

							A__						,       A__						,       A__							=       self._encode_prompt(
							    UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		)

							if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		):
													A__							=       torch.cat(UpperCAmelCase__			,	dim=0		)
							if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		):
													A__							=       torch.cat(UpperCAmelCase__			,	dim=0		)

							if do_classifier_free_guidance:
													A__							=       image_embeds.repeat_interleave(UpperCAmelCase__			,	dim=0		)
													A__							=       negative_image_embeds.repeat_interleave(UpperCAmelCase__			,	dim=0		)

							A__							=       torch.cat([negative_image_embeds, image_embeds]			,	dim=0		).to(
							    dtype=prompt_embeds.dtype			,	device=UpperCAmelCase__		)

							self.scheduler.set_timesteps(UpperCAmelCase__			,	device=UpperCAmelCase__		)
							A__							=       self.scheduler.timesteps

							A__							=       self.unet.config.in_channels

							A__						,       A__							=       get_new_h_w(UpperCAmelCase__			,	UpperCAmelCase__			,	self.movq_scale_factor		)

							# create initial latent
							A__							=       self.prepare_latents(
							    (batch_size, num_channels_latents, height, width)			,	text_encoder_hidden_states.dtype			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	self.scheduler			,	)

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

													A__							=       {"text_embeds": prompt_embeds, "image_embeds": image_embeds}
													A__							=       self.unet(
													    sample=UpperCAmelCase__			,	timestep=UpperCAmelCase__			,	encoder_hidden_states=UpperCAmelCase__			,	added_cond_kwargs=UpperCAmelCase__			,	return_dict=UpperCAmelCase__			,	)[0]

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

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

							if output_type == "pil":
													A__							=       self.numpy_to_pil(UpperCAmelCase__		)

							if not return_dict:
													return (image,)

							return ImagePipelineOutput(images=UpperCAmelCase__		)










from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

UpperCAmelCase_							:      List[Any]      =  [
    "python",
    "tqdm",
    "regex",
    "requests",
    "packaging",
    "filelock",
    "numpy",
    "tokenizers",
    "huggingface-hub",
    "safetensors",
    "accelerate",
    "pyyaml",
]

for pkg in pkgs_to_check_at_runtime:
						if pkg in deps:
												if pkg == "tokenizers":
																		# must be loaded here, or else tqdm check may fail
																		from .utils import is_tokenizers_available

																		if not is_tokenizers_available():
																								continue  # not required, check version only if installed
												elif pkg == "accelerate":
																		# must be loaded here, or else tqdm check may fail
																		from .utils import is_accelerate_available

																		# Maybe switch to is_torch_available in the future here so that Accelerate is hard dep of
																		# Transformers with PyTorch
																		if not is_accelerate_available():
																								continue  # not required, check version only if installed

												require_version_core(deps[pkg])
						else:
												raise ValueError(F'''can\'t find {pkg} in {deps.keys()}, check dependency_versions_table.py''')
def  UpperCamelCase						(       _A			:				List[Any]			,							_A			:				int=None    )->       Optional[int]:





						"""simple docstring"""





						require_version(deps[pkg]			,							_A    )





