





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


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

_a						=   """
    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__(__snake_case ,__snake_case ,__snake_case=8							)       ->   Dict:







	'''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       __A					(					__snake_case				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						):



				'''simple docstring'''







				super().__init__()

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				if latents is None:
					lowerCamelCase__              =	randn_tensor(A_ ,						generator=A_ ,						device=A_ ,						dtype=A_							)
				else:
					if latents.shape != shape:
						raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'							)
					lowerCamelCase__              =	latents.to(A_							)

				lowerCamelCase__              =	latents * scheduler.init_noise_sigma
				return latents





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase=None ,						):



				'''simple docstring'''







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

				lowerCamelCase__              =	text_inputs.input_ids
				lowerCamelCase__              =	self.tokenizer(A_ ,						padding='''longest''' ,						return_tensors='''pt'''							).input_ids

				if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(A_ ,						A_							):
					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(A_							)
				lowerCamelCase__              =	text_inputs.attention_mask.to(A_							)

				lowerCamelCase__              =	self.text_encoder(
				    input_ids=A_ ,						attention_mask=A_							)

				lowerCamelCase__              =	prompt_embeds.repeat_interleave(A_ ,						dim=0							)
				lowerCamelCase__              =	text_encoder_hidden_states.repeat_interleave(A_ ,						dim=0							)
				lowerCamelCase__              =	text_mask.repeat_interleave(A_ ,						dim=0							)

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

					lowerCamelCase__              =	self.text_encoder(
					    input_ids=A_ ,						attention_mask=A_							)

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

					lowerCamelCase__              =	negative_prompt_embeds.shape[1]
					lowerCamelCase__              =	negative_prompt_embeds.repeat(1 ,						A_							)
					lowerCamelCase__              =	negative_prompt_embeds.view(batch_size * num_images_per_prompt ,						A_							)

					lowerCamelCase__              =	uncond_text_encoder_hidden_states.shape[1]
					lowerCamelCase__              =	uncond_text_encoder_hidden_states.repeat(1 ,						A_ ,						1							)
					lowerCamelCase__              =	uncond_text_encoder_hidden_states.view(
					    batch_size * num_images_per_prompt ,						A_ ,						-1							)
					lowerCamelCase__              =	uncond_text_mask.repeat_interleave(A_ ,						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 	__lowerCamelCase       (					self ,						__lowerCAmelCase=0							):



				'''simple docstring'''







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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase=0							):



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

				lowerCamelCase__              =	torch.device(F'cuda:{gpu_id}'							)

				if self.device.type != "cpu":
					self.to('''cpu''' ,						silence_dtype_warnings=A_							)
					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__              =	cpu_offload_with_hook(A_ ,						A_ ,						prev_module_hook=A_							)

				if self.safety_checker is not None:
					lowerCamelCase__              =	cpu_offload_with_hook(self.safety_checker ,						A_ ,						prev_module_hook=A_							)

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





			@property
			# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				if not hasattr(self.unet ,						'''_hf_hook'''							):
					return self.device
				for module in self.unet.modules():
					if (
					    hasattr(A_ ,						'''_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(A_							)
			def __call__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = 5_1_2 ,						__lowerCAmelCase = 5_1_2 ,						__lowerCAmelCase = 1_0_0 ,						__lowerCAmelCase = 4.0 ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = "pil" ,						__lowerCAmelCase = True ,						):



				'''simple docstring'''







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

				lowerCamelCase__              =	self._execution_device

				lowerCamelCase__              =	batch_size * num_images_per_prompt
				lowerCamelCase__              =	guidance_scale > 1.0

				lowerCamelCase__              =	self._encode_prompt(
				    A_ ,						A_ ,						A_ ,						A_ ,						A_							)

				if isinstance(A_ ,						A_							):
					lowerCamelCase__              =	torch.cat(A_ ,						dim=0							)
				if isinstance(A_ ,						A_							):
					lowerCamelCase__              =	torch.cat(A_ ,						dim=0							)

				if do_classifier_free_guidance:
					lowerCamelCase__              =	image_embeds.repeat_interleave(A_ ,						dim=0							)
					lowerCamelCase__              =	negative_image_embeds.repeat_interleave(A_ ,						dim=0							)

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

				self.scheduler.set_timesteps(A_ ,						device=A_							)
				lowerCamelCase__              =	self.scheduler.timesteps

				lowerCamelCase__              =	self.unet.config.in_channels

				lowerCamelCase__              =	get_new_h_w(A_ ,						A_ ,						self.movq_scale_factor							)

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

				for i, t in enumerate(self.progress_bar(A_							)							):
					# 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=A_ ,						timestep=A_ ,						encoder_hidden_states=A_ ,						added_cond_kwargs=A_ ,						return_dict=A_ ,						)[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(
					    A_ ,						A_ ,						A_ ,						generator=A_ ,						).prev_sample
				# post-processing
				lowerCamelCase__              =	self.movq.decode(A_ ,						force_not_quantize=A_							)["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(A_							)

				if not return_dict:
					return (image,)

				return ImagePipelineOutput(images=A_							)









from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		42






class       __A					(					lowerCAmelCase    ,       lowerCAmelCase				):

			'''simple docstring'''







			@register_to_config
			def __init__(					self ,						__lowerCAmelCase = 1_6 ,						__lowerCAmelCase = 8_8 ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = 0.0 ,						__lowerCAmelCase = 3_2 ,						__lowerCAmelCase = None ,						__lowerCAmelCase = False ,						__lowerCAmelCase = None ,						__lowerCAmelCase = "geglu" ,						__lowerCAmelCase = True ,						__lowerCAmelCase = True ,						):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	num_attention_heads
				lowerCamelCase__              =	attention_head_dim
				lowerCamelCase__              =	num_attention_heads * attention_head_dim

				lowerCamelCase__              =	in_channels

				lowerCamelCase__              =	torch.nn.GroupNorm(num_groups=__lowerCAmelCase ,						num_channels=__lowerCAmelCase ,						eps=1E-6 ,						affine=__lowerCAmelCase							)
				lowerCamelCase__              =	nn.Linear(__lowerCAmelCase ,						__lowerCAmelCase							)

				# 3. Define transformers blocks
				lowerCamelCase__              =	nn.ModuleList(
				    [
				        BasicTransformerBlock(
				            __lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						dropout=__lowerCAmelCase ,						cross_attention_dim=__lowerCAmelCase ,						activation_fn=__lowerCAmelCase ,						attention_bias=__lowerCAmelCase ,						double_self_attention=__lowerCAmelCase ,						norm_elementwise_affine=__lowerCAmelCase ,						)
				        for d in range(__lowerCAmelCase							)
				    ]							)

				lowerCamelCase__              =	nn.Linear(__lowerCAmelCase ,						__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase=None ,						__lowerCAmelCase=None ,						__lowerCAmelCase=None ,						__lowerCAmelCase=1 ,						__lowerCAmelCase=None ,						__lowerCAmelCase = True ,						):



				'''simple docstring'''







				lowerCamelCase__       ,			lowerCamelCase__       ,			lowerCamelCase__       ,			lowerCamelCase__              =	hidden_states.shape
				lowerCamelCase__              =	batch_frames // num_frames

				lowerCamelCase__              =	hidden_states

				lowerCamelCase__              =	hidden_states[None, :].reshape(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							)
				lowerCamelCase__              =	hidden_states.permute(0 ,						2 ,						1 ,						3 ,						4							)

				lowerCamelCase__              =	self.norm(__lowerCAmelCase							)
				lowerCamelCase__              =	hidden_states.permute(0 ,						3 ,						4 ,						2 ,						1							).reshape(batch_size * height * width ,						__lowerCAmelCase ,						__lowerCAmelCase							)

				lowerCamelCase__              =	self.proj_in(__lowerCAmelCase							)

				# 2. Blocks
				for block in self.transformer_blocks:
					lowerCamelCase__              =	block(
					    __lowerCAmelCase ,						encoder_hidden_states=__lowerCAmelCase ,						timestep=__lowerCAmelCase ,						cross_attention_kwargs=__lowerCAmelCase ,						class_labels=__lowerCAmelCase ,						)

				# 3. Output
				lowerCamelCase__              =	self.proj_out(__lowerCAmelCase							)
				lowerCamelCase__              =	(
				    hidden_states[None, None, :]
				    .reshape(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							)
				    .permute(0 ,						3 ,						4 ,						1 ,						2							)
				    .contiguous()
				)
				lowerCamelCase__              =	hidden_states.reshape(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							)

				lowerCamelCase__              =	hidden_states + residual

				if not return_dict:
					return (output,)

				return TransformerTemporalModelOutput(sample=__lowerCAmelCase							)



