




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


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

A     ='\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 						snake_case_							(_a	:							Dict    ,   _a	:							Optional[int]    ,   _a	:							Any=8     ):
    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 			_a  (	__a ):
       def __init__(	self						:	List[Any]   ,				lowercase						:	MultilingualCLIP   ,				lowercase						:	XLMRobertaTokenizer   ,				lowercase						:	UNetaDConditionModel   ,				lowercase						:	Union[DDIMScheduler, DDPMScheduler]   ,				lowercase						:	VQModel   ,				):



           '''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						:	int   ,				lowercase						:	List[Any]   ,				lowercase						:	int   ,				lowercase						:	Dict   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[Any]   ,				lowercase						:	str				):



           '''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						:	List[str]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Optional[Any]   ,				lowercase						:	int   ,				lowercase						:	str   ,				lowercase						:	Dict=None   ,				):



           '''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=77   ,				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=77   ,				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						:	str=0				):



           '''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						:	str   ,				lowercase						:	str=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.'''				)

           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						:	str				):



           '''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[Any]   ,				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 = 512   ,				lowercase						:	int = 512   ,				lowercase						:	int = 100   ,				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   ,				):



           '''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 (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_snake_case     :		Union[str, Any]       						= {
    'configuration_encodec': [
        'ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EncodecConfig',
    ],
    'feature_extraction_encodec': ['EncodecFeatureExtractor'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_snake_case     :		int       						= [
		    'ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'EncodecModel',
		    'EncodecPreTrainedModel',
		]

if TYPE_CHECKING:
		from .configuration_encodec import (
		    ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    EncodecConfig,
		)
		from .feature_extraction_encodec import EncodecFeatureExtractor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_encodec import (
				    ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
				    EncodecModel,
				    EncodecPreTrainedModel,
				)

else:
		import sys

		_snake_case     :		str       						= _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


