



'''simple docstring'''



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


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

lowerCAmelCase	:		List[Any]				=   '''
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")
        >>> pipe_prior.to("cuda")

        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")
        >>> pipe.to("cuda")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save("cat.png")
        ```
'''



def   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Optional[Any]				,		lowerCamelCase   :						Any=8					):






	'''simple docstring'''




	__lowerCAmelCase												=  h // scale_factor**2
	if h % scale_factor**2 != 0:
		new_h += 1
	__lowerCAmelCase												=  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			,				)     ->			Any:
						super().__init__()

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			List[str]:
						if latents is None:
							__lowerCAmelCase												=  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}'''     )
							__lowerCAmelCase												=  latents.to(UpperCamelCase     )

						__lowerCAmelCase												=  latents * scheduler.init_noise_sigma
						return latents



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase=None			,				)     ->			Any:
						__lowerCAmelCase												=  len(UpperCamelCase     ) if isinstance(UpperCamelCase			,				UpperCamelCase     ) else 1
						# get prompt text embeddings
						__lowerCAmelCase												=  self.tokenizer(
						    UpperCamelCase			,				padding="max_length"			,				truncation=UpperCamelCase			,				max_length=77			,				return_attention_mask=UpperCamelCase			,				add_special_tokens=UpperCamelCase			,				return_tensors="pt"			,				)

						__lowerCAmelCase												=  text_inputs.input_ids
						__lowerCAmelCase												=  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     ):
							__lowerCAmelCase												=  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}'''     )

						__lowerCAmelCase												=  text_input_ids.to(UpperCamelCase     )
						__lowerCAmelCase												=  text_inputs.attention_mask.to(UpperCamelCase     )

						__lowerCAmelCase							,       __lowerCAmelCase												=  self.text_encoder(
						    input_ids=UpperCamelCase			,				attention_mask=UpperCamelCase     )

						__lowerCAmelCase												=  prompt_embeds.repeat_interleave(UpperCamelCase			,				dim=0     )
						__lowerCAmelCase												=  text_encoder_hidden_states.repeat_interleave(UpperCamelCase			,				dim=0     )
						__lowerCAmelCase												=  text_mask.repeat_interleave(UpperCamelCase			,				dim=0     )

						if do_classifier_free_guidance:
							__lowerCAmelCase												=  42
							if negative_prompt is None:
								__lowerCAmelCase												=  [""] * 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     ):
								__lowerCAmelCase												=  [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:
								__lowerCAmelCase												=  negative_prompt

							__lowerCAmelCase												=  self.tokenizer(
							    UpperCamelCase			,				padding="max_length"			,				max_length=77			,				truncation=UpperCamelCase			,				return_attention_mask=UpperCamelCase			,				add_special_tokens=UpperCamelCase			,				return_tensors="pt"			,				)
							__lowerCAmelCase												=  uncond_input.input_ids.to(UpperCamelCase     )
							__lowerCAmelCase												=  uncond_input.attention_mask.to(UpperCamelCase     )

							__lowerCAmelCase							,       __lowerCAmelCase												=  self.text_encoder(
							    input_ids=UpperCamelCase			,				attention_mask=UpperCamelCase     )

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

							__lowerCAmelCase												=  negative_prompt_embeds.shape[1]
							__lowerCAmelCase												=  negative_prompt_embeds.repeat(1			,				UpperCamelCase     )
							__lowerCAmelCase												=  negative_prompt_embeds.view(batch_size * num_images_per_prompt			,				UpperCamelCase     )

							__lowerCAmelCase												=  uncond_text_encoder_hidden_states.shape[1]
							__lowerCAmelCase												=  uncond_text_encoder_hidden_states.repeat(1			,				UpperCamelCase			,				1     )
							__lowerCAmelCase												=  uncond_text_encoder_hidden_states.view(
							    batch_size * num_images_per_prompt			,				UpperCamelCase			,				-1     )
							__lowerCAmelCase												=  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
							__lowerCAmelCase												=  torch.cat([negative_prompt_embeds, prompt_embeds]     )
							__lowerCAmelCase												=  torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]     )

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

						return prompt_embeds, text_encoder_hidden_states, text_mask



					def        UpperCAmelCase_ (    self			,				UpperCamelCase=0     )     ->			Optional[Any]:
						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.text_encoder,
						    self.movq,
						]
						for cpu_offloaded_model in models:
							if cpu_offloaded_model is not None:
								cpu_offload(UpperCamelCase			,				UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase=0     )     ->			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."     )

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

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

						if self.safety_checker is not None:
							__lowerCAmelCase							,       __lowerCAmelCase												=  cpu_offload_with_hook(self.safety_checker			,				UpperCamelCase			,				prev_module_hook=UpperCamelCase     )

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



					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def        UpperCAmelCase_ (    self     )     ->			Dict:
						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			,				)     ->			int:

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

						__lowerCAmelCase												=  self._execution_device

						__lowerCAmelCase												=  batch_size * num_images_per_prompt
						__lowerCAmelCase												=  guidance_scale > 1.0

						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  self._encode_prompt(
						    UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )

						if isinstance(UpperCamelCase			,				UpperCamelCase     ):
							__lowerCAmelCase												=  torch.cat(UpperCamelCase			,				dim=0     )
						if isinstance(UpperCamelCase			,				UpperCamelCase     ):
							__lowerCAmelCase												=  torch.cat(UpperCamelCase			,				dim=0     )

						if do_classifier_free_guidance:
							__lowerCAmelCase												=  image_embeds.repeat_interleave(UpperCamelCase			,				dim=0     )
							__lowerCAmelCase												=  negative_image_embeds.repeat_interleave(UpperCamelCase			,				dim=0     )

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

						self.scheduler.set_timesteps(UpperCamelCase			,				device=UpperCamelCase     )
						__lowerCAmelCase												=  self.scheduler.timesteps

						__lowerCAmelCase												=  self.unet.config.in_channels

						__lowerCAmelCase							,       __lowerCAmelCase												=  get_new_h_w(UpperCamelCase			,				UpperCamelCase			,				self.movq_scale_factor     )

						# create initial latent
						__lowerCAmelCase												=  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
							__lowerCAmelCase												=  torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents

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

							if do_classifier_free_guidance:
								__lowerCAmelCase							,       __lowerCAmelCase												=  noise_pred.split(latents.shape[1]			,				dim=1     )
								__lowerCAmelCase							,       __lowerCAmelCase												=  noise_pred.chunk(2     )
								__lowerCAmelCase							,       __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							,       __lowerCAmelCase												=  noise_pred.split(latents.shape[1]			,				dim=1     )

							# compute the previous noisy sample x_t -> x_t-1
							__lowerCAmelCase												=  self.scheduler.step(
							    UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				generator=UpperCamelCase			,				).prev_sample
						# post-processing
						__lowerCAmelCase												=  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"]:
							__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(UpperCamelCase     )

						if not return_dict:
							return (image,)

						return ImagePipelineOutput(images=UpperCamelCase     )




'''simple docstring'''



import importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Union

import torch

from ..utils import BaseOutput


lowerCAmelCase	:		Optional[Any]				=   '''scheduler_config.json'''






class 					UpperCAmelCase__	(		UpperCamelCase__ ):
					a      :     str        =							1
					a      :     Optional[int]        =							2
					a      :     int        =							3
					a      :     Union[str, Any]        =							4
					a      :     int        =							5
					a      :     Optional[int]        =							6
					a      :     str        =							7
					a      :     List[Any]        =							8
					a      :     List[str]        =							9
					a      :     List[str]        =							1_0
					a      :     int        =							1_1
					a      :     Any        =							1_2
					a      :     Any        =							1_3
					a      :     Tuple        =							1_4






@dataclass
class 					UpperCAmelCase__	(		UpperCamelCase__ ):
					a      :     torch.FloatTensor






class 					UpperCAmelCase__	:
					a      :     Tuple        =							SCHEDULER_CONFIG_NAME
					a      :     Union[str, Any]        =							[]
					a      :     str        =							True



					@classmethod
					def        UpperCAmelCase_ (    cls			,				UpperCamelCase = None			,				UpperCamelCase = None			,				UpperCamelCase=False			,				**UpperCamelCase			,				)     ->			int:
						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  cls.load_config(
						    pretrained_model_name_or_path=UpperCamelCase			,				subfolder=UpperCamelCase			,				return_unused_kwargs=UpperCamelCase			,				return_commit_hash=UpperCamelCase			,				**UpperCamelCase			,				)
						return cls.from_config(UpperCamelCase			,				return_unused_kwargs=UpperCamelCase			,				**UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase = False			,				**UpperCamelCase     )     ->			Dict:
						self.save_config(save_directory=UpperCamelCase			,				push_to_hub=UpperCamelCase			,				**UpperCamelCase     )



					@property
					def        UpperCAmelCase_ (    self     )     ->			str:
						return self._get_compatibles()



					@classmethod
					def        UpperCAmelCase_ (    cls     )     ->			Tuple:
						__lowerCAmelCase												=  list(set([cls.__name__] + cls._compatibles     )     )
						__lowerCAmelCase												=  importlib.import_module(__name__.split("."     )[0]     )
						__lowerCAmelCase												=  [
						    getattr(UpperCamelCase			,				UpperCamelCase     ) for c in compatible_classes_str if hasattr(UpperCamelCase			,				UpperCamelCase     )
						]
						return compatible_classes