





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


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

SCREAMING_SNAKE_CASE_		     = """
    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      __lowercase     (   _SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE=8			)				->     str:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE	=		h // scale_factor**2
      if h % scale_factor**2 != 0:
            new_h += 1
      SCREAMING_SNAKE_CASE	=		w // scale_factor**2
      if w % scale_factor**2 != 0:
            new_w += 1
      return new_h * scale_factor, new_w * scale_factor



class 	UpperCamelCase__				(   lowerCAmelCase_     ):



     '''simple docstring'''

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




           '''simple docstring'''


           super().__init__()

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






     def        SCREAMING_SNAKE_CASE__						(							self		:					str		,lowerCamelCase__		:					Any		,lowerCamelCase__		:					List[str]		,lowerCamelCase__		:					List[Any]		,lowerCamelCase__		:					Optional[int]		,lowerCamelCase__		:					List[str]		,lowerCamelCase__		:					List[str] )  ->	Union[str, Any]:




           '''simple docstring'''


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

           SCREAMING_SNAKE_CASE	=		latents * scheduler.init_noise_sigma
           return latents






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str]		,lowerCamelCase__		:					Union[str, Any]		,lowerCamelCase__		:					Any		,lowerCamelCase__		:					Optional[int]		,lowerCamelCase__		:					Any		,lowerCamelCase__		:					List[str]=None		,)  ->	Any:




           '''simple docstring'''


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

           SCREAMING_SNAKE_CASE	=		text_inputs.input_ids
           SCREAMING_SNAKE_CASE	=		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__ ):
                 SCREAMING_SNAKE_CASE	=		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}""" )

           SCREAMING_SNAKE_CASE	=		text_input_ids.to(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		text_inputs.attention_mask.to(lowerCamelCase__ )

           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.text_encoder(
               input_ids=lowerCamelCase__		,attention_mask=lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		prompt_embeds.repeat_interleave(lowerCamelCase__		,dim=0 )
           SCREAMING_SNAKE_CASE	=		text_encoder_hidden_states.repeat_interleave(lowerCamelCase__		,dim=0 )
           SCREAMING_SNAKE_CASE	=		text_mask.repeat_interleave(lowerCamelCase__		,dim=0 )

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

                 SCREAMING_SNAKE_CASE	=		self.tokenizer(
                     lowerCamelCase__		,padding="""max_length"""		,max_length=77		,truncation=lowerCamelCase__		,return_attention_mask=lowerCamelCase__		,add_special_tokens=lowerCamelCase__		,return_tensors="""pt"""		,)
                 SCREAMING_SNAKE_CASE	=		uncond_input.input_ids.to(lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		uncond_input.attention_mask.to(lowerCamelCase__ )

                 SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.text_encoder(
                     input_ids=lowerCamelCase__		,attention_mask=lowerCamelCase__ )

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

                 SCREAMING_SNAKE_CASE	=		negative_prompt_embeds.shape[1]
                 SCREAMING_SNAKE_CASE	=		negative_prompt_embeds.repeat(1		,lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		negative_prompt_embeds.view(batch_size * num_images_per_prompt		,lowerCamelCase__ )

                 SCREAMING_SNAKE_CASE	=		uncond_text_encoder_hidden_states.shape[1]
                 SCREAMING_SNAKE_CASE	=		uncond_text_encoder_hidden_states.repeat(1		,lowerCamelCase__		,1 )
                 SCREAMING_SNAKE_CASE	=		uncond_text_encoder_hidden_states.view(
                     batch_size * num_images_per_prompt		,lowerCamelCase__		,-1 )
                 SCREAMING_SNAKE_CASE	=		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
                 SCREAMING_SNAKE_CASE	=		torch.cat([negative_prompt_embeds, prompt_embeds] )
                 SCREAMING_SNAKE_CASE	=		torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states] )

                 SCREAMING_SNAKE_CASE	=		torch.cat([uncond_text_mask, text_mask] )

           return prompt_embeds, text_encoder_hidden_states, text_mask






     def        SCREAMING_SNAKE_CASE__						(							self		:					Any		,lowerCamelCase__		:					Tuple=0 )  ->	Optional[int]:




           '''simple docstring'''


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

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

           SCREAMING_SNAKE_CASE	=		[
               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        SCREAMING_SNAKE_CASE__						(							self		:					Union[str, Any]		,lowerCamelCase__		:					List[str]=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.""" )

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

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

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

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






     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple )  ->	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		:					Union[str, 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 = 512		,lowerCamelCase__		:					int = 512		,lowerCamelCase__		:					int = 100		,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		,)  ->	Union[str, Any]:




           '''simple docstring'''


           if isinstance(lowerCamelCase__		,lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		1
           elif isinstance(lowerCamelCase__		,lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		len(lowerCamelCase__ )
           else:
                 raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(lowerCamelCase__ )}""" )

           SCREAMING_SNAKE_CASE	=		self._execution_device

           SCREAMING_SNAKE_CASE	=		batch_size * num_images_per_prompt
           SCREAMING_SNAKE_CASE	=		guidance_scale > 1.0

           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self._encode_prompt(
               lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__ )

           if isinstance(lowerCamelCase__		,lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		torch.cat(lowerCamelCase__		,dim=0 )
           if isinstance(lowerCamelCase__		,lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		torch.cat(lowerCamelCase__		,dim=0 )

           if do_classifier_free_guidance:
                 SCREAMING_SNAKE_CASE	=		image_embeds.repeat_interleave(lowerCamelCase__		,dim=0 )
                 SCREAMING_SNAKE_CASE	=		negative_image_embeds.repeat_interleave(lowerCamelCase__		,dim=0 )

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

           self.scheduler.set_timesteps(lowerCamelCase__		,device=lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		self.scheduler.timesteps

           SCREAMING_SNAKE_CASE	=		self.unet.config.in_channels

           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		get_new_h_w(lowerCamelCase__		,lowerCamelCase__		,self.movq_scale_factor )

           # create initial latent
           SCREAMING_SNAKE_CASE	=		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
                 SCREAMING_SNAKE_CASE	=		torch.cat([latents] * 2 ) if do_classifier_free_guidance else latents

                 SCREAMING_SNAKE_CASE	=		{"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
                 SCREAMING_SNAKE_CASE	=		self.unet(
                     sample=lowerCamelCase__		,timestep=lowerCamelCase__		,encoder_hidden_states=lowerCamelCase__		,added_cond_kwargs=lowerCamelCase__		,return_dict=lowerCamelCase__		,)[0]

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

                 # compute the previous noisy sample x_t -> x_t-1
                 SCREAMING_SNAKE_CASE	=		self.scheduler.step(
                     lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__		,generator=lowerCamelCase__		,).prev_sample
           # post-processing
           SCREAMING_SNAKE_CASE	=		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"]:
                 SCREAMING_SNAKE_CASE	=		image * 0.5 + 0.5
                 SCREAMING_SNAKE_CASE	=		image.clamp(0		,1 )
                 SCREAMING_SNAKE_CASE	=		image.cpu().permute(0		,2		,3		,1 ).float().numpy()

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

           if not return_dict:
                 return (image,)

           return ImagePipelineOutput(images=lowerCamelCase__ )









def      __lowercase     (   _SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE			)				->     float:


      '''simple docstring'''


      _validate_point(_SCREAMING_SNAKE_CASE			)
      _validate_point(_SCREAMING_SNAKE_CASE			)
      if len(_SCREAMING_SNAKE_CASE			) != len(_SCREAMING_SNAKE_CASE			):
            raise ValueError("""Both points must be in the same n-dimensional space"""			)

      return float(sum(abs(a - b			) for a, b in zip(_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE			)			)			)



def      __lowercase     (   _SCREAMING_SNAKE_CASE			)				->     None:


      '''simple docstring'''


      if point:
            if isinstance(_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE			):
                  for item in point:
                        if not isinstance(_SCREAMING_SNAKE_CASE				,		(int, float)			):
                              SCREAMING_SNAKE_CASE	=		(
                                  """Expected a list of numbers as input, found """
                                  F"""{type(_SCREAMING_SNAKE_CASE			).__name__}"""
                              )
                              raise TypeError(_SCREAMING_SNAKE_CASE			)
            else:
                  SCREAMING_SNAKE_CASE	=		F"""Expected a list of numbers as input, found {type(_SCREAMING_SNAKE_CASE			).__name__}"""
                  raise TypeError(_SCREAMING_SNAKE_CASE			)
      else:
            raise ValueError("""Missing an input"""			)



def      __lowercase     (   _SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE			)				->     float:


      '''simple docstring'''


      _validate_point(_SCREAMING_SNAKE_CASE			)
      _validate_point(_SCREAMING_SNAKE_CASE			)
      if len(_SCREAMING_SNAKE_CASE			) != len(_SCREAMING_SNAKE_CASE			):
            raise ValueError("""Both points must be in the same n-dimensional space"""			)

      return float(sum(abs(x - y			) for x, y in zip(_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE			)			)			)


if __name__ == "__main__":
      import doctest

      doctest.testmod()


