





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

A    :					List[str] =      '''
    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					(							__a   :str    ,     __a   :Any    ,     __a   :List[str]=8    )		->				Dict:






		"""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   A			(SCREAMING_SNAKE_CASE    ):
			'''simple docstring'''




			def __init__(			self   :    Dict							,	__lowerCAmelCase   :    MultilingualCLIP							,	__lowerCAmelCase   :    XLMRobertaTokenizer							,	__lowerCAmelCase   :    UNetaDConditionModel							,	__lowerCAmelCase   :    Union[DDIMScheduler, DDPMScheduler]							,	__lowerCAmelCase   :    VQModel							,	)	->	Optional[int]:







					"""simple docstring"""






					super().__init__()

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






			def    a_       (			self   :    str							,	__lowerCAmelCase   :    Tuple							,	__lowerCAmelCase   :    Optional[int]							,	__lowerCAmelCase   :    Optional[int]							,	__lowerCAmelCase   :    Tuple							,	__lowerCAmelCase   :    Union[str, Any]							,	__lowerCAmelCase   :    List[str]   )	->	int:







					"""simple docstring"""






					if latents is None:
							A__        =   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}'   )
							A__        =   latents.to(__lowerCAmelCase   )

					A__        =   latents * scheduler.init_noise_sigma
					return latents






			def    a_       (			self   :    Any							,	__lowerCAmelCase   :    List[Any]							,	__lowerCAmelCase   :    List[Any]							,	__lowerCAmelCase   :    List[Any]							,	__lowerCAmelCase   :    Union[str, Any]							,	__lowerCAmelCase   :    int=None							,	)	->	Tuple:







					"""simple docstring"""






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

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

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

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

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

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

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

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

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

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







					"""simple docstring"""






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






			def    a_       (			self   :    str							,	__lowerCAmelCase   :    Tuple=0   )	->	Optional[int]:







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

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

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

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

					# 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   :    Optional[Any]   )	->	List[str]:







					"""simple docstring"""






					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   :    List[Any]							,	__lowerCAmelCase   :    Union[str, List[str]]							,	__lowerCAmelCase   :    Union[torch.FloatTensor, List[torch.FloatTensor]]							,	__lowerCAmelCase   :    Union[torch.FloatTensor, List[torch.FloatTensor]]							,	__lowerCAmelCase   :    Optional[Union[str, List[str]]] = None							,	__lowerCAmelCase   :    int = 5_12							,	__lowerCAmelCase   :    int = 5_12							,	__lowerCAmelCase   :    int = 1_00							,	__lowerCAmelCase   :    float = 4.0							,	__lowerCAmelCase   :    int = 1							,	__lowerCAmelCase   :    Optional[Union[torch.Generator, List[torch.Generator]]] = None							,	__lowerCAmelCase   :    Optional[torch.FloatTensor] = None							,	__lowerCAmelCase   :    Optional[str] = "pil"							,	__lowerCAmelCase   :    bool = True							,	)	->	Any:







					"""simple docstring"""






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

					A__        =   self._execution_device

					A__        =   batch_size * num_images_per_prompt
					A__        =   guidance_scale > 1.0

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

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

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

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

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

					A__        =   self.unet.config.in_channels

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

					# create initial latent
					A__        =   self.prepare_latents(
					    (batch_size, num_channels_latents, height, width)							,	text_encoder_hidden_states.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
							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=__lowerCAmelCase							,	timestep=__lowerCAmelCase							,	encoder_hidden_states=__lowerCAmelCase							,	added_cond_kwargs=__lowerCAmelCase							,	return_dict=__lowerCAmelCase							,	)[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(
							    __lowerCAmelCase							,	__lowerCAmelCase							,	__lowerCAmelCase							,	generator=__lowerCAmelCase							,	).prev_sample
					# post-processing
					A__        =   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"]:
							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(__lowerCAmelCase   )

					if not return_dict:
							return (image,)

					return ImagePipelineOutput(images=__lowerCAmelCase   )










import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)


class   A			(unittest.TestCase    ):
			'''simple docstring'''




			def    a_       (			self   :    Any   )	->	Union[str, Any]:







					"""simple docstring"""






					A__        =   ["""a""", """b""", """c"""]

					# Defaults to last layer if both are None
					A__      ,			A__        =   get_aligned_output_features_output_indices(__lowerCAmelCase							,	__lowerCAmelCase							,	__lowerCAmelCase   )
					self.assertEqual(__lowerCAmelCase							,	["""c"""]   )
					self.assertEqual(__lowerCAmelCase							,	[2]   )

					# Out indices set to match out features
					A__      ,			A__        =   get_aligned_output_features_output_indices(["""a""", """c"""]							,	__lowerCAmelCase							,	__lowerCAmelCase   )
					self.assertEqual(__lowerCAmelCase							,	["""a""", """c"""]   )
					self.assertEqual(__lowerCAmelCase							,	[0, 2]   )

					# Out features set to match out indices
					A__      ,			A__        =   get_aligned_output_features_output_indices(__lowerCAmelCase							,	[0, 2]							,	__lowerCAmelCase   )
					self.assertEqual(__lowerCAmelCase							,	["""a""", """c"""]   )
					self.assertEqual(__lowerCAmelCase							,	[0, 2]   )

					# Out features selected from negative indices
					A__      ,			A__        =   get_aligned_output_features_output_indices(__lowerCAmelCase							,	[-3, -1]							,	__lowerCAmelCase   )
					self.assertEqual(__lowerCAmelCase							,	["""a""", """c"""]   )
					self.assertEqual(__lowerCAmelCase							,	[-3, -1]   )






			def    a_       (			self   :    Optional[Any]   )	->	Optional[Any]:







					"""simple docstring"""






					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(["""a""", """b"""]							,	(0, 1)							,	__lowerCAmelCase   )

					# Out features must be a list
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(("""a""", """b""")							,	(0, 1)							,	["""a""", """b"""]   )

					# Out features must be a subset of stage names
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(["""a""", """b"""]							,	(0, 1)							,	["""a"""]   )

					# Out indices must be a list or tuple
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(__lowerCAmelCase							,	0							,	["""a""", """b"""]   )

					# Out indices must be a subset of stage names
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(__lowerCAmelCase							,	(0, 1)							,	["""a"""]   )

					# Out features and out indices must be the same length
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(["""a""", """b"""]							,	(0,)							,	["""a""", """b""", """c"""]   )

					# Out features should match out indices
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(["""a""", """b"""]							,	(0, 2)							,	["""a""", """b""", """c"""]   )

					# Out features and out indices should be in order
					with self.assertRaises(__lowerCAmelCase   ):
							verify_out_features_out_indices(["""b""", """a"""]							,	(0, 1)							,	["""a""", """b"""]   )

					# Check passes with valid inputs
					verify_out_features_out_indices(["""a""", """b""", """d"""]							,	(0, 1, -1)							,	["""a""", """b""", """c""", """d"""]   )







			def    a_       (			self   :    Union[str, Any]   )	->	List[str]:







					"""simple docstring"""






					A__        =   BackboneMixin()

					A__        =   ["""a""", """b""", """c"""]
					A__        =   ["""a""", """c"""]
					A__        =   [0, 2]

					# Check that the output features and indices are set correctly
					self.assertEqual(backbone.out_features							,	["""a""", """c"""]   )
					self.assertEqual(backbone.out_indices							,	[0, 2]   )

					# Check out features and indices are updated correctly
					A__        =   ["""a""", """b"""]
					self.assertEqual(backbone.out_features							,	["""a""", """b"""]   )
					self.assertEqual(backbone.out_indices							,	[0, 1]   )

					A__        =   [-3, -1]
					self.assertEqual(backbone.out_features							,	["""a""", """c"""]   )
					self.assertEqual(backbone.out_indices							,	[-3, -1]   )



