





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 :  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 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase=8  )		->				Dict:
 UpperCamelCase			:      Union[str, Any] 			=     h // scale_factor**2
 if h % scale_factor**2 != 0:
  new_h += 1
 UpperCamelCase			:      Any 			=     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    ):





      def __init__(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							):
       '''simple docstring'''




       super().__init__()

       self.register_modules(
           text_encoder=A_ ,							tokenizer=A_ ,							unet=A_ ,							scheduler=A_ ,							movq=A_ ,							)
       UpperCamelCase			:      List[Any] 			=     2 ** (len(self.movq.config.block_out_channels     ) - 1)





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       if latents is None:
        UpperCamelCase			:      int 			=     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}"""     )
        UpperCamelCase			:      Optional[Any] 			=     latents.to(A_     )

       UpperCamelCase			:      Dict 			=     latents * scheduler.init_noise_sigma
       return latents





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_=None ,							):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     len(A_     ) if isinstance(A_ ,							A_     ) else 1
       # get prompt text embeddings
       UpperCamelCase			:      Any 			=     self.tokenizer(
           A_ ,							padding="max_length" ,							truncation=A_ ,							max_length=77 ,							return_attention_mask=A_ ,							add_special_tokens=A_ ,							return_tensors="pt" ,							)

       UpperCamelCase			:      List[Any] 			=     text_inputs.input_ids
       UpperCamelCase			:      List[Any] 			=     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_     ):
        UpperCamelCase			:      Union[str, Any] 			=     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			:      str 			=     text_input_ids.to(A_     )
       UpperCamelCase			:      int 			=     text_inputs.attention_mask.to(A_     )

       UpperCamelCase			,					UpperCamelCase			:      Optional[int] 			=     self.text_encoder(
           input_ids=A_ ,							attention_mask=A_     )

       UpperCamelCase			:      Optional[Any] 			=     prompt_embeds.repeat_interleave(A_ ,							dim=0     )
       UpperCamelCase			:      Any 			=     text_encoder_hidden_states.repeat_interleave(A_ ,							dim=0     )
       UpperCamelCase			:      int 			=     text_mask.repeat_interleave(A_ ,							dim=0     )

       if do_classifier_free_guidance:
        UpperCamelCase			:      List[str]
        if negative_prompt is None:
         UpperCamelCase			:      Optional[Any] 			=     [""] * 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_     ):
         UpperCamelCase			:      Union[str, Any] 			=     [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:
         UpperCamelCase			:      Optional[int] 			=     negative_prompt

        UpperCamelCase			:      Tuple 			=     self.tokenizer(
            A_ ,							padding="max_length" ,							max_length=77 ,							truncation=A_ ,							return_attention_mask=A_ ,							add_special_tokens=A_ ,							return_tensors="pt" ,							)
        UpperCamelCase			:      Optional[Any] 			=     uncond_input.input_ids.to(A_     )
        UpperCamelCase			:      Tuple 			=     uncond_input.attention_mask.to(A_     )

        UpperCamelCase			,					UpperCamelCase			:      Optional[int] 			=     self.text_encoder(
            input_ids=A_ ,							attention_mask=A_     )

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

        UpperCamelCase			:      int 			=     negative_prompt_embeds.shape[1]
        UpperCamelCase			:      Any 			=     negative_prompt_embeds.repeat(1 ,							A_     )
        UpperCamelCase			:      List[Any] 			=     negative_prompt_embeds.view(batch_size * num_images_per_prompt ,							A_     )

        UpperCamelCase			:      List[Any] 			=     uncond_text_encoder_hidden_states.shape[1]
        UpperCamelCase			:      int 			=     uncond_text_encoder_hidden_states.repeat(1 ,							A_ ,							1     )
        UpperCamelCase			:      Union[str, Any] 			=     uncond_text_encoder_hidden_states.view(
            batch_size * num_images_per_prompt ,							A_ ,							-1     )
        UpperCamelCase			:      Optional[int] 			=     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
        UpperCamelCase			:      List[Any] 			=     torch.cat([negative_prompt_embeds, prompt_embeds]     )
        UpperCamelCase			:      Optional[int] 			=     torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]     )

        UpperCamelCase			:      Tuple 			=     torch.cat([uncond_text_mask, text_mask]     )

       return prompt_embeds, text_encoder_hidden_states, text_mask





      def  __UpperCamelCase(				self ,							A_=0     ):
       '''simple docstring'''




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

       UpperCamelCase			:      Dict 			=     torch.device(F"""cuda:{gpu_id}"""     )

       UpperCamelCase			:      Optional[int] 			=     [
           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  __UpperCamelCase(				self ,							A_=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			:      Union[str, Any] 			=     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)

       UpperCamelCase			:      Tuple 			=     None
       for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
        UpperCamelCase			,					UpperCamelCase			:      Optional[int] 			=     cpu_offload_with_hook(A_ ,							A_ ,							prev_module_hook=A_     )

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

       # We'll offload the last model manually.
       UpperCamelCase			:      str 			=     hook





      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def  __UpperCamelCase(				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 ,							A_ ,							A_ ,							A_ ,							A_ = None ,							A_ = 512 ,							A_ = 512 ,							A_ = 100 ,							A_ = 4.0 ,							A_ = 1 ,							A_ = None ,							A_ = None ,							A_ = "pil" ,							A_ = True ,							):
       '''simple docstring'''




       if isinstance(A_ ,							A_     ):
        UpperCamelCase			:      Optional[Any] 			=     1
       elif isinstance(A_ ,							A_     ):
        UpperCamelCase			:      List[Any] 			=     len(A_     )
       else:
        raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(A_     )}"""     )

       UpperCamelCase			:      str 			=     self._execution_device

       UpperCamelCase			:      Dict 			=     batch_size * num_images_per_prompt
       UpperCamelCase			:      Optional[int] 			=     guidance_scale > 1.0

       UpperCamelCase			,					UpperCamelCase			,					UpperCamelCase			:      Union[str, Any] 			=     self._encode_prompt(
           A_ ,							A_ ,							A_ ,							A_ ,							A_     )

       if isinstance(A_ ,							A_     ):
        UpperCamelCase			:      Any 			=     torch.cat(A_ ,							dim=0     )
       if isinstance(A_ ,							A_     ):
        UpperCamelCase			:      int 			=     torch.cat(A_ ,							dim=0     )

       if do_classifier_free_guidance:
        UpperCamelCase			:      int 			=     image_embeds.repeat_interleave(A_ ,							dim=0     )
        UpperCamelCase			:      Optional[Any] 			=     negative_image_embeds.repeat_interleave(A_ ,							dim=0     )

       UpperCamelCase			:      List[Any] 			=     torch.cat([negative_image_embeds, image_embeds] ,							dim=0     ).to(
           dtype=prompt_embeds.dtype ,							device=A_     )

       self.scheduler.set_timesteps(A_ ,							device=A_     )
       UpperCamelCase			:      Any 			=     self.scheduler.timesteps

       UpperCamelCase			:      Tuple 			=     self.unet.config.in_channels

       UpperCamelCase			,					UpperCamelCase			:      Union[str, Any] 			=     get_new_h_w(A_ ,							A_ ,							self.movq_scale_factor     )

       # create initial latent
       UpperCamelCase			:      Tuple 			=     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
        UpperCamelCase			:      Union[str, Any] 			=     torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents

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

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

        # compute the previous noisy sample x_t -> x_t-1
        UpperCamelCase			:      Optional[int] 			=     self.scheduler.step(
            A_ ,							A_ ,							A_ ,							generator=A_ ,							).prev_sample
       # post-processing
       UpperCamelCase			:      Tuple 			=     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"]:
        UpperCamelCase			:      Tuple 			=     image * 0.5 + 0.5
        UpperCamelCase			:      Union[str, Any] 			=     image.clamp(0 ,							1     )
        UpperCamelCase			:      str 			=     image.cpu().permute(0 ,							2 ,							3 ,							1     ).float().numpy()

       if output_type == "pil":
        UpperCamelCase			:      Tuple 			=     self.numpy_to_pil(A_     )

       if not return_dict:
        return (image,)

       return ImagePipelineOutput(images=A_     )







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase :  Union[str, Any]										=  {
    """configuration_vision_encoder_decoder""": ["""VisionEncoderDecoderConfig""", """VisionEncoderDecoderOnnxConfig"""]
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase :  Dict										=  ["""VisionEncoderDecoderModel"""]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase :  List[str]										=  ["""TFVisionEncoderDecoderModel"""]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase :  int										=  ["""FlaxVisionEncoderDecoderModel"""]

if TYPE_CHECKING:
    from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig, VisionEncoderDecoderOnnxConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_vision_encoder_decoder import VisionEncoderDecoderModel

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_vision_encoder_decoder import TFVisionEncoderDecoderModel

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_vision_encoder_decoder import FlaxVisionEncoderDecoderModel

else:
    import sys

    __lowerCamelCase :  str										=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
