






'''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_				: Optional[Any]    	=							logging.get_logger(__name__)  # pylint: disable=invalid-name

lowerCamelCase_				: Union[str, 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 	__magic_name__( _A						,      _A						,      _A=8					):
						'''simple docstring'''





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







class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			def __init__(   self							:						str ,      lowercase							:						MultilingualCLIP ,      lowercase							:						XLMRobertaTokenizer ,      lowercase							:						UNetaDConditionModel ,      lowercase							:						Union[DDIMScheduler, DDPMScheduler] ,      lowercase							:						VQModel ,      )							->    Optional[Any]:

									'''simple docstring'''







									super().__init__()

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






			def 				A      (   self							:						Optional[Any] ,      lowercase							:						List[str] ,      lowercase							:						Any ,      lowercase							:						int ,      lowercase							:						str ,      lowercase							:						Any ,      lowercase							:						Any			)							->    Optional[Any]:

									'''simple docstring'''







									if latents is None:
															UpperCamelCase__         = randn_tensor(lowercase ,      generator=lowercase ,      device=lowercase ,      dtype=lowercase			)
									else:
															if latents.shape != shape:
																					raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"			)
															UpperCamelCase__         = latents.to(lowercase			)

									UpperCamelCase__         = latents * scheduler.init_noise_sigma
									return latents






			def 				A      (   self							:						Union[str, Any] ,      lowercase							:						str ,      lowercase							:						Optional[Any] ,      lowercase							:						int ,      lowercase							:						List[Any] ,      lowercase							:						Optional[Any]=None ,      )							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = len(lowercase			) if isinstance(lowercase ,      lowercase			) else 1
									# get prompt text embeddings
									UpperCamelCase__         = self.tokenizer(
									    lowercase ,      padding="""max_length""" ,      truncation=lowercase ,      max_length=7_7 ,      return_attention_mask=lowercase ,      add_special_tokens=lowercase ,      return_tensors="""pt""" ,      )

									UpperCamelCase__         = text_inputs.input_ids
									UpperCamelCase__         = self.tokenizer(lowercase ,      padding="""longest""" ,      return_tensors="""pt"""			).input_ids

									if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowercase ,      lowercase			):
															UpperCamelCase__         = 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}"			)

									UpperCamelCase__         = text_input_ids.to(lowercase			)
									UpperCamelCase__         = text_inputs.attention_mask.to(lowercase			)

									UpperCamelCase__    ,   UpperCamelCase__         = self.text_encoder(
									    input_ids=lowercase ,      attention_mask=lowercase			)

									UpperCamelCase__         = prompt_embeds.repeat_interleave(lowercase ,      dim=0			)
									UpperCamelCase__         = text_encoder_hidden_states.repeat_interleave(lowercase ,      dim=0			)
									UpperCamelCase__         = text_mask.repeat_interleave(lowercase ,      dim=0			)

									if do_classifier_free_guidance:
															UpperCamelCase__         = 42
															if negative_prompt is None:
																					UpperCamelCase__         = [""""""] * batch_size
															elif type(lowercase			) is not type(lowercase			):
																					raise TypeError(
																					    f"`negative_prompt` should be the same type to `prompt`, but got {type(lowercase			)} !="
																					    f" {type(lowercase			)}."			)
															elif isinstance(lowercase ,      lowercase			):
																					UpperCamelCase__         = [negative_prompt]
															elif batch_size != len(lowercase			):
																					raise ValueError(
																					    f"`negative_prompt`: {negative_prompt} has batch size {len(lowercase			)}, but `prompt`:"
																					    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
																					    """ the batch size of `prompt`."""			)
															else:
																					UpperCamelCase__         = negative_prompt

															UpperCamelCase__         = self.tokenizer(
															    lowercase ,      padding="""max_length""" ,      max_length=7_7 ,      truncation=lowercase ,      return_attention_mask=lowercase ,      add_special_tokens=lowercase ,      return_tensors="""pt""" ,      )
															UpperCamelCase__         = uncond_input.input_ids.to(lowercase			)
															UpperCamelCase__         = uncond_input.attention_mask.to(lowercase			)

															UpperCamelCase__    ,   UpperCamelCase__         = self.text_encoder(
															    input_ids=lowercase ,      attention_mask=lowercase			)

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

															UpperCamelCase__         = negative_prompt_embeds.shape[1]
															UpperCamelCase__         = negative_prompt_embeds.repeat(1 ,      lowercase			)
															UpperCamelCase__         = negative_prompt_embeds.view(batch_size * num_images_per_prompt ,      lowercase			)

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

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

									return prompt_embeds, text_encoder_hidden_states, text_mask






			def 				A      (   self							:						Any ,      lowercase							:						int=0			)							->    int:

									'''simple docstring'''







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

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

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






			def 				A      (   self							:						int ,      lowercase							:						List[Any]=0			)							->    Dict:

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

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

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

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

									if self.safety_checker is not None:
															UpperCamelCase__    ,   UpperCamelCase__         = cpu_offload_with_hook(self.safety_checker ,      lowercase ,      prev_module_hook=lowercase			)

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






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

									'''simple docstring'''







									if not hasattr(self.unet ,      """_hf_hook"""			):
															return self.device
									for module in self.unet.modules():
															if (
															    hasattr(lowercase ,      """_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(lowercase			)
			def __call__(   self							:						Optional[int] ,      lowercase							:						Union[str, List[str]] ,      lowercase							:						Union[torch.FloatTensor, List[torch.FloatTensor]] ,      lowercase							:						Union[torch.FloatTensor, List[torch.FloatTensor]] ,      lowercase							:						Optional[Union[str, List[str]]] = None ,      lowercase							:						int = 5_1_2 ,      lowercase							:						int = 5_1_2 ,      lowercase							:						int = 1_0_0 ,      lowercase							:						float = 4.0 ,      lowercase							:						int = 1 ,      lowercase							:						Optional[Union[torch.Generator, List[torch.Generator]]] = None ,      lowercase							:						Optional[torch.FloatTensor] = None ,      lowercase							:						Optional[str] = "pil" ,      lowercase							:						bool = True ,      )							->    Tuple:

									'''simple docstring'''







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

									UpperCamelCase__         = self._execution_device

									UpperCamelCase__         = batch_size * num_images_per_prompt
									UpperCamelCase__         = guidance_scale > 1.0

									UpperCamelCase__    ,   UpperCamelCase__    ,   UpperCamelCase__         = self._encode_prompt(
									    lowercase ,      lowercase ,      lowercase ,      lowercase ,      lowercase			)

									if isinstance(lowercase ,      lowercase			):
															UpperCamelCase__         = torch.cat(lowercase ,      dim=0			)
									if isinstance(lowercase ,      lowercase			):
															UpperCamelCase__         = torch.cat(lowercase ,      dim=0			)

									if do_classifier_free_guidance:
															UpperCamelCase__         = image_embeds.repeat_interleave(lowercase ,      dim=0			)
															UpperCamelCase__         = negative_image_embeds.repeat_interleave(lowercase ,      dim=0			)

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

									self.scheduler.set_timesteps(lowercase ,      device=lowercase			)
									UpperCamelCase__         = self.scheduler.timesteps

									UpperCamelCase__         = self.unet.config.in_channels

									UpperCamelCase__    ,   UpperCamelCase__         = get_new_h_w(lowercase ,      lowercase ,      self.movq_scale_factor			)

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

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

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

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

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

									if output_type == "pil":
															UpperCamelCase__         = self.numpy_to_pil(lowercase			)

									if not return_dict:
															return (image,)

									return ImagePipelineOutput(images=lowercase			)








'''simple docstring'''


from typing import TYPE_CHECKING

from ..utils import _LazyModule


lowerCamelCase_				: List[Any]    	=							{
    '''config''': [
        '''EXTERNAL_DATA_FORMAT_SIZE_LIMIT''',
        '''OnnxConfig''',
        '''OnnxConfigWithPast''',
        '''OnnxSeq2SeqConfigWithPast''',
        '''PatchingSpec''',
    ],
    '''convert''': ['''export''', '''validate_model_outputs'''],
    '''features''': ['''FeaturesManager'''],
    '''utils''': ['''ParameterFormat''', '''compute_serialized_parameters_size'''],
}


if TYPE_CHECKING:
					from .config import (
					    EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
					    OnnxConfig,
					    OnnxConfigWithPast,
					    OnnxSeqaSeqConfigWithPast,
					    PatchingSpec,
					)
					from .convert import export, validate_model_outputs
					from .features import FeaturesManager
					from .utils import ParameterFormat, compute_serialized_parameters_size

else:
					import sys

					lowerCamelCase_				: List[Any]    	=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
