


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   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=8   )						->List[str]:
   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   __lowercase				(  __snake_case  ):





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






         """simple docstring"""







         super().__init__()

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





      def        _lowercase				(		self						:					str					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					List[str]					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					str					,       __lowerCamelCase						:					Any       )    ->   Tuple:






         """simple docstring"""







         if latents is None:
            UpperCAmelCase							=     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}"""       )
            UpperCAmelCase							=     latents.to(__lowerCamelCase       )

         UpperCAmelCase							=     latents * scheduler.init_noise_sigma
         return latents





      def        _lowercase				(		self						:					List[Any]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					str					,       __lowerCamelCase						:					Union[str, Any]=None					,       )    ->   List[str]:






         """simple docstring"""







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

         UpperCAmelCase							=     text_inputs.input_ids
         UpperCAmelCase							=     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       ):
            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(__lowerCamelCase       )
         UpperCAmelCase							=     text_inputs.attention_mask.to(__lowerCamelCase       )

         UpperCAmelCase				,  UpperCAmelCase							=     self.text_encoder(
             input_ids=__lowerCamelCase					,       attention_mask=__lowerCamelCase       )

         UpperCAmelCase							=     prompt_embeds.repeat_interleave(__lowerCamelCase					,       dim=0       )
         UpperCAmelCase							=     text_encoder_hidden_states.repeat_interleave(__lowerCamelCase					,       dim=0       )
         UpperCAmelCase							=     text_mask.repeat_interleave(__lowerCamelCase					,       dim=0       )

         if do_classifier_free_guidance:
            UpperCAmelCase							=     42
            if negative_prompt is None:
               UpperCAmelCase							=     [""""""] * 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       ):
               UpperCAmelCase							=     [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:
               UpperCAmelCase							=     negative_prompt

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

            UpperCAmelCase				,  UpperCAmelCase							=     self.text_encoder(
                input_ids=__lowerCamelCase					,       attention_mask=__lowerCamelCase       )

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

            UpperCAmelCase							=     negative_prompt_embeds.shape[1]
            UpperCAmelCase							=     negative_prompt_embeds.repeat(1					,       __lowerCamelCase       )
            UpperCAmelCase							=     negative_prompt_embeds.view(batch_size * num_images_per_prompt					,       __lowerCamelCase       )

            UpperCAmelCase							=     uncond_text_encoder_hidden_states.shape[1]
            UpperCAmelCase							=     uncond_text_encoder_hidden_states.repeat(1					,       __lowerCamelCase					,       1       )
            UpperCAmelCase							=     uncond_text_encoder_hidden_states.view(
                batch_size * num_images_per_prompt					,       __lowerCamelCase					,       -1       )
            UpperCAmelCase							=     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
            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        _lowercase				(		self						:					Tuple					,       __lowerCamelCase						:					Dict=0       )    ->   Any:






         """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(__lowerCamelCase					,       __lowerCamelCase       )





      def        _lowercase				(		self						:					Optional[int]					,       __lowerCamelCase						:					List[Any]=0       )    ->   List[str]:






         """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=__lowerCamelCase       )
            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(__lowerCamelCase					,       __lowerCamelCase					,       prev_module_hook=__lowerCamelCase       )

         if self.safety_checker is not None:
            UpperCAmelCase				,  UpperCAmelCase							=     cpu_offload_with_hook(self.safety_checker					,       __lowerCamelCase					,       prev_module_hook=__lowerCamelCase       )

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





      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def        _lowercase				(		self						:					Tuple       )    ->   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[str]					,       __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_1_2					,       __lowerCamelCase						:					int = 5_1_2					,       __lowerCamelCase						:					int = 1_0_0					,       __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       ):
            UpperCAmelCase							=     1
         elif isinstance(__lowerCamelCase					,       __lowerCamelCase       ):
            UpperCAmelCase							=     len(__lowerCamelCase       )
         else:
            raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__lowerCamelCase       )}"""       )

         UpperCAmelCase							=     self._execution_device

         UpperCAmelCase							=     batch_size * num_images_per_prompt
         UpperCAmelCase							=     guidance_scale > 1.0

         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     self._encode_prompt(
             __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase       )

         if isinstance(__lowerCamelCase					,       __lowerCamelCase       ):
            UpperCAmelCase							=     torch.cat(__lowerCamelCase					,       dim=0       )
         if isinstance(__lowerCamelCase					,       __lowerCamelCase       ):
            UpperCAmelCase							=     torch.cat(__lowerCamelCase					,       dim=0       )

         if do_classifier_free_guidance:
            UpperCAmelCase							=     image_embeds.repeat_interleave(__lowerCamelCase					,       dim=0       )
            UpperCAmelCase							=     negative_image_embeds.repeat_interleave(__lowerCamelCase					,       dim=0       )

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

         self.scheduler.set_timesteps(__lowerCamelCase					,       device=__lowerCamelCase       )
         UpperCAmelCase							=     self.scheduler.timesteps

         UpperCAmelCase							=     self.unet.config.in_channels

         UpperCAmelCase				,  UpperCAmelCase							=     get_new_h_w(__lowerCamelCase					,       __lowerCamelCase					,       self.movq_scale_factor       )

         # create initial latent
         UpperCAmelCase							=     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
            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=__lowerCamelCase					,       timestep=__lowerCamelCase					,       encoder_hidden_states=__lowerCamelCase					,       added_cond_kwargs=__lowerCamelCase					,       return_dict=__lowerCamelCase					,       )[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(
                __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       generator=__lowerCamelCase					,       ).prev_sample
         # post-processing
         UpperCAmelCase							=     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"]:
            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(__lowerCamelCase       )

         if not return_dict:
            return (image,)

         return ImagePipelineOutput(images=__lowerCamelCase       )







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__a     =			{"""configuration_ibert""": ["""IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """IBertConfig""", """IBertOnnxConfig"""]}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __a     =			[
     """IBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """IBertForMaskedLM""",
     """IBertForMultipleChoice""",
     """IBertForQuestionAnswering""",
     """IBertForSequenceClassification""",
     """IBertForTokenClassification""",
     """IBertModel""",
     """IBertPreTrainedModel""",
 ]

if TYPE_CHECKING:
 from .configuration_ibert import IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, IBertConfig, IBertOnnxConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_ibert import (
      IBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      IBertForMaskedLM,
      IBertForMultipleChoice,
      IBertForQuestionAnswering,
      IBertForSequenceClassification,
      IBertForTokenClassification,
      IBertModel,
      IBertPreTrainedModel,
  )

else:
 import sys

 __a     =			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



