






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

lowerCamelCase	:							Dict				='''
    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 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase=8			) ->	str:
     UpperCamelCase__			:   Optional[int]       				=  h // scale_factor**2
     if h % scale_factor**2 != 0:
          new_h += 1
     UpperCamelCase__			:   Optional[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  (    A__       ):







 def __init__(		self					:   Tuple							,   SCREAMING_SNAKE_CASE					:   MultilingualCLIP							,   SCREAMING_SNAKE_CASE					:   XLMRobertaTokenizer							,   SCREAMING_SNAKE_CASE					:   UNetaDConditionModel							,   SCREAMING_SNAKE_CASE					:   Union[DDIMScheduler, DDPMScheduler]							,   SCREAMING_SNAKE_CASE					:   VQModel							,   ):


      '''simple docstring'''






      super().__init__()

      self.register_modules(
          text_encoder=SCREAMING_SNAKE_CASE							,   tokenizer=SCREAMING_SNAKE_CASE							,   unet=SCREAMING_SNAKE_CASE							,   scheduler=SCREAMING_SNAKE_CASE							,   movq=SCREAMING_SNAKE_CASE							,   )
      UpperCamelCase__			:   Dict       				=  2 ** (len(self.movq.config.block_out_channels						) - 1)







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   str						):


      '''simple docstring'''






      if latents is None:
           UpperCamelCase__			:   List[str]       				=  randn_tensor(SCREAMING_SNAKE_CASE							,   generator=SCREAMING_SNAKE_CASE							,   device=SCREAMING_SNAKE_CASE							,   dtype=SCREAMING_SNAKE_CASE						)
      else:
           if latents.shape != shape:
                raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'						)
           UpperCamelCase__			:   Optional[Any]       				=  latents.to(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Dict       				=  latents * scheduler.init_noise_sigma
      return latents







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=None							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   str       				=  len(SCREAMING_SNAKE_CASE						) if isinstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						) else 1
      # get prompt text embeddings
      UpperCamelCase__			:   Union[str, Any]       				=  self.tokenizer(
          SCREAMING_SNAKE_CASE							,   padding="max_length"							,   truncation=SCREAMING_SNAKE_CASE							,   max_length=77							,   return_attention_mask=SCREAMING_SNAKE_CASE							,   add_special_tokens=SCREAMING_SNAKE_CASE							,   return_tensors="pt"							,   )

      UpperCamelCase__			:   List[Any]       				=  text_inputs.input_ids
      UpperCamelCase__			:   Any       				=  self.tokenizer(SCREAMING_SNAKE_CASE							,   padding="longest"							,   return_tensors="pt"						).input_ids

      if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						):
           UpperCamelCase__			:   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__			:   List[str]       				=  text_input_ids.to(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  text_inputs.attention_mask.to(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__				,     UpperCamelCase__			:   Union[str, Any]       				=  self.text_encoder(
          input_ids=SCREAMING_SNAKE_CASE							,   attention_mask=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Union[str, Any]       				=  prompt_embeds.repeat_interleave(SCREAMING_SNAKE_CASE							,   dim=0						)
      UpperCamelCase__			:   List[str]       				=  text_encoder_hidden_states.repeat_interleave(SCREAMING_SNAKE_CASE							,   dim=0						)
      UpperCamelCase__			:   Tuple       				=  text_mask.repeat_interleave(SCREAMING_SNAKE_CASE							,   dim=0						)

      if do_classifier_free_guidance:
           UpperCamelCase__			:   List[str]
           if negative_prompt is None:
                UpperCamelCase__			:   Optional[int]       				=  [""] * batch_size
           elif type(SCREAMING_SNAKE_CASE						) is not type(SCREAMING_SNAKE_CASE						):
                raise TypeError(
                    F'`negative_prompt` should be the same type to `prompt`, but got {type(SCREAMING_SNAKE_CASE						)} !='
                    F' {type(SCREAMING_SNAKE_CASE						)}.'						)
           elif isinstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						):
                UpperCamelCase__			:   Union[str, Any]       				=  [negative_prompt]
           elif batch_size != len(SCREAMING_SNAKE_CASE						):
                raise ValueError(
                    F'`negative_prompt`: {negative_prompt} has batch size {len(SCREAMING_SNAKE_CASE						)}, but `prompt`:'
                    F' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
                    " the batch size of `prompt`."						)
           else:
                UpperCamelCase__			:   Tuple       				=  negative_prompt

           UpperCamelCase__			:   Dict       				=  self.tokenizer(
               SCREAMING_SNAKE_CASE							,   padding="max_length"							,   max_length=77							,   truncation=SCREAMING_SNAKE_CASE							,   return_attention_mask=SCREAMING_SNAKE_CASE							,   add_special_tokens=SCREAMING_SNAKE_CASE							,   return_tensors="pt"							,   )
           UpperCamelCase__			:   Optional[int]       				=  uncond_input.input_ids.to(SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   Union[str, Any]       				=  uncond_input.attention_mask.to(SCREAMING_SNAKE_CASE						)

           UpperCamelCase__				,     UpperCamelCase__			:   List[Any]       				=  self.text_encoder(
               input_ids=SCREAMING_SNAKE_CASE							,   attention_mask=SCREAMING_SNAKE_CASE						)

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

           UpperCamelCase__			:   Optional[Any]       				=  negative_prompt_embeds.shape[1]
           UpperCamelCase__			:   List[Any]       				=  negative_prompt_embeds.repeat(1							,   SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   List[str]       				=  negative_prompt_embeds.view(batch_size * num_images_per_prompt							,   SCREAMING_SNAKE_CASE						)

           UpperCamelCase__			:   List[str]       				=  uncond_text_encoder_hidden_states.shape[1]
           UpperCamelCase__			:   Tuple       				=  uncond_text_encoder_hidden_states.repeat(1							,   SCREAMING_SNAKE_CASE							,   1						)
           UpperCamelCase__			:   Optional[int]       				=  uncond_text_encoder_hidden_states.view(
               batch_size * num_images_per_prompt							,   SCREAMING_SNAKE_CASE							,   -1						)
           UpperCamelCase__			:   List[Any]       				=  uncond_text_mask.repeat_interleave(SCREAMING_SNAKE_CASE							,   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__			:   Tuple       				=  torch.cat([negative_prompt_embeds, prompt_embeds]						)
           UpperCamelCase__			:   List[str]       				=  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 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   List[str]=0						):


      '''simple docstring'''






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

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

      UpperCamelCase__			:   int       				=  [
          self.unet,
          self.text_encoder,
          self.movq,
      ]
      for cpu_offloaded_model in models:
           if cpu_offloaded_model is not None:
                cpu_offload(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Tuple=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__			:   Optional[int]       				=  torch.device(F'cuda:{gpu_id}'						)

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

      if self.safety_checker is not None:
           UpperCamelCase__				,     UpperCamelCase__			:   Any       				=  cpu_offload_with_hook(self.safety_checker							,   SCREAMING_SNAKE_CASE							,   prev_module_hook=SCREAMING_SNAKE_CASE						)

      # We'll offload the last model manually.
      UpperCamelCase__			:   Optional[int]       				=  hook







 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      if not hasattr(self.unet							,   "_hf_hook"						):
           return self.device
      for module in self.unet.modules():
           if (
               hasattr(SCREAMING_SNAKE_CASE							,   "_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(SCREAMING_SNAKE_CASE						)
 def __call__(		self					:   str							,   SCREAMING_SNAKE_CASE					:   Union[str, List[str]]							,   SCREAMING_SNAKE_CASE					:   Union[torch.FloatTensor, List[torch.FloatTensor]]							,   SCREAMING_SNAKE_CASE					:   Union[torch.FloatTensor, List[torch.FloatTensor]]							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, List[str]]] = None							,   SCREAMING_SNAKE_CASE					:   int = 5_12							,   SCREAMING_SNAKE_CASE					:   int = 5_12							,   SCREAMING_SNAKE_CASE					:   int = 1_00							,   SCREAMING_SNAKE_CASE					:   float = 4.0							,   SCREAMING_SNAKE_CASE					:   int = 1							,   SCREAMING_SNAKE_CASE					:   Optional[Union[torch.Generator, List[torch.Generator]]] = None							,   SCREAMING_SNAKE_CASE					:   Optional[torch.FloatTensor] = None							,   SCREAMING_SNAKE_CASE					:   Optional[str] = "pil"							,   SCREAMING_SNAKE_CASE					:   bool = True							,   ):


      '''simple docstring'''






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

      UpperCamelCase__			:   int       				=  self._execution_device

      UpperCamelCase__			:   int       				=  batch_size * num_images_per_prompt
      UpperCamelCase__			:   Tuple       				=  guidance_scale > 1.0

      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   Union[str, Any]       				=  self._encode_prompt(
          SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)

      if isinstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						):
           UpperCamelCase__			:   Tuple       				=  torch.cat(SCREAMING_SNAKE_CASE							,   dim=0						)
      if isinstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						):
           UpperCamelCase__			:   Any       				=  torch.cat(SCREAMING_SNAKE_CASE							,   dim=0						)

      if do_classifier_free_guidance:
           UpperCamelCase__			:   List[Any]       				=  image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE							,   dim=0						)
           UpperCamelCase__			:   List[Any]       				=  negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE							,   dim=0						)

      UpperCamelCase__			:   Union[str, Any]       				=  torch.cat([negative_image_embeds, image_embeds]							,   dim=0						).to(
          dtype=prompt_embeds.dtype							,   device=SCREAMING_SNAKE_CASE						)

      self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE							,   device=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[str]       				=  self.scheduler.timesteps

      UpperCamelCase__			:   List[Any]       				=  self.unet.config.in_channels

      UpperCamelCase__				,     UpperCamelCase__			:   str       				=  get_new_h_w(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   self.movq_scale_factor						)

      # create initial latent
      UpperCamelCase__			:   str       				=  self.prepare_latents(
          (batch_size, num_channels_latents, height, width)							,   text_encoder_hidden_states.dtype							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   self.scheduler							,   )

      for i, t in enumerate(self.progress_bar(SCREAMING_SNAKE_CASE						)						):
           # expand the latents if we are doing classifier free guidance
           UpperCamelCase__			:   str       				=  torch.cat([latents] * 2						) if do_classifier_free_guidance else latents

           UpperCamelCase__			:   List[Any]       				=  {"text_embeds": prompt_embeds, "image_embeds": image_embeds}
           UpperCamelCase__			:   Optional[int]       				=  self.unet(
               sample=SCREAMING_SNAKE_CASE							,   timestep=SCREAMING_SNAKE_CASE							,   encoder_hidden_states=SCREAMING_SNAKE_CASE							,   added_cond_kwargs=SCREAMING_SNAKE_CASE							,   return_dict=SCREAMING_SNAKE_CASE							,   )[0]

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

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

      if output_type == "pil":
           UpperCamelCase__			:   Tuple       				=  self.numpy_to_pil(SCREAMING_SNAKE_CASE						)

      if not return_dict:
           return (image,)

      return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE						)







import numpy as np







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase			) ->	np.array:
     return 1 / (1 + np.exp(-vector			))







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase			) ->	np.array:
     return vector * sigmoid(1.7_0_2 * vector			)


if __name__ == "__main__":
 import doctest

 doctest.testmod()