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


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

SCREAMING_SNAKE_CASE_:	List[Any]    ='\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       lowerCAmelCase_							(   snake_case_			:					Tuple   ,       snake_case_			:					List[str]   ,       snake_case_			:					List[str]=8     )       -> Optional[Any]:


     '''simple docstring'''




     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 (					UpperCamelCase__      ):



    def __init__(self       :       Dict  ,     __a       :       MultilingualCLIP  ,     __a       :       XLMRobertaTokenizer  ,     __a       :       UNetaDConditionModel  ,     __a       :       Union[DDIMScheduler, DDPMScheduler]  ,     __a       :       VQModel  ,     ):
         super().__init__()

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



    def    _lowercase		(self       :       Union[str, Any]  ,     __a       :       Optional[Any]  ,     __a       :       Tuple  ,     __a       :       Optional[Any]  ,     __a       :       Dict  ,     __a       :       Tuple  ,     __a       :       Dict					):
         if latents is None:
              UpperCAmelCase_			       =  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_			       =  latents.to(__a					)

         UpperCAmelCase_			       =  latents * scheduler.init_noise_sigma
         return latents



    def    _lowercase		(self       :       str  ,     __a       :       int  ,     __a       :       Union[str, Any]  ,     __a       :       List[Any]  ,     __a       :       Optional[Any]  ,     __a       :       Any=None  ,     ):
         UpperCAmelCase_			       =  len(__a					) if isinstance(__a  ,     __a					) else 1
         # get prompt text embeddings
         UpperCAmelCase_			       =  self.tokenizer(
             __a  ,     padding="max_length"  ,     truncation=__a  ,     max_length=77  ,     return_attention_mask=__a  ,     add_special_tokens=__a  ,     return_tensors="pt"  ,     )

         UpperCAmelCase_			       =  text_inputs.input_ids
         UpperCAmelCase_			       =  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_			       =  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(__a					)
         UpperCAmelCase_			       =  text_inputs.attention_mask.to(__a					)

         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.text_encoder(
             input_ids=__a  ,     attention_mask=__a					)

         UpperCAmelCase_			       =  prompt_embeds.repeat_interleave(__a  ,     dim=0					)
         UpperCAmelCase_			       =  text_encoder_hidden_states.repeat_interleave(__a  ,     dim=0					)
         UpperCAmelCase_			       =  text_mask.repeat_interleave(__a  ,     dim=0					)

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

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

              UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.text_encoder(
                  input_ids=__a  ,     attention_mask=__a					)

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

              UpperCAmelCase_			       =  negative_prompt_embeds.shape[1]
              UpperCAmelCase_			       =  negative_prompt_embeds.repeat(1  ,     __a					)
              UpperCAmelCase_			       =  negative_prompt_embeds.view(batch_size * num_images_per_prompt  ,     __a					)

              UpperCAmelCase_			       =  uncond_text_encoder_hidden_states.shape[1]
              UpperCAmelCase_			       =  uncond_text_encoder_hidden_states.repeat(1  ,     __a  ,     1					)
              UpperCAmelCase_			       =  uncond_text_encoder_hidden_states.view(
                  batch_size * num_images_per_prompt  ,     __a  ,     -1					)
              UpperCAmelCase_			       =  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_			       =  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       :       List[Any]  ,     __a       :       Tuple=0					):
         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(__a  ,     __a					)



    def    _lowercase		(self       :       Tuple  ,     __a       :       List[Any]=0					):
         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=__a					)
              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(__a  ,     __a  ,     prev_module_hook=__a					)

         if self.safety_checker is not None:
              UpperCAmelCase_       ,					UpperCAmelCase_			       =  cpu_offload_with_hook(self.safety_checker  ,     __a  ,     prev_module_hook=__a					)

         # 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       :       int					):
         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       :       List[str]  ,     __a       :       Union[str, List[str]]  ,     __a       :       Union[torch.FloatTensor, List[torch.FloatTensor]]  ,     __a       :       Union[torch.FloatTensor, List[torch.FloatTensor]]  ,     __a       :       Optional[Union[str, List[str]]] = None  ,     __a       :       int = 512  ,     __a       :       int = 512  ,     __a       :       int = 100  ,     __a       :       float = 4.0  ,     __a       :       int = 1  ,     __a       :       Optional[Union[torch.Generator, List[torch.Generator]]] = None  ,     __a       :       Optional[torch.FloatTensor] = None  ,     __a       :       Optional[str] = "pil"  ,     __a       :       bool = True  ,     ):

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

         UpperCAmelCase_			       =  self._execution_device

         UpperCAmelCase_			       =  batch_size * num_images_per_prompt
         UpperCAmelCase_			       =  guidance_scale > 1.0

         UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_			       =  self._encode_prompt(
             __a  ,     __a  ,     __a  ,     __a  ,     __a					)

         if isinstance(__a  ,     __a					):
              UpperCAmelCase_			       =  torch.cat(__a  ,     dim=0					)
         if isinstance(__a  ,     __a					):
              UpperCAmelCase_			       =  torch.cat(__a  ,     dim=0					)

         if do_classifier_free_guidance:
              UpperCAmelCase_			       =  image_embeds.repeat_interleave(__a  ,     dim=0					)
              UpperCAmelCase_			       =  negative_image_embeds.repeat_interleave(__a  ,     dim=0					)

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

         self.scheduler.set_timesteps(__a  ,     device=__a					)
         UpperCAmelCase_			       =  self.scheduler.timesteps

         UpperCAmelCase_			       =  self.unet.config.in_channels

         UpperCAmelCase_       ,					UpperCAmelCase_			       =  get_new_h_w(__a  ,     __a  ,     self.movq_scale_factor					)

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

              UpperCAmelCase_			       =  {"text_embeds": prompt_embeds, "image_embeds": image_embeds}
              UpperCAmelCase_			       =  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_			       =  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(
                  __a  ,     __a  ,     __a  ,     generator=__a  ,     ).prev_sample
         # post-processing
         UpperCAmelCase_			       =  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_			       =  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(__a					)

         if not return_dict:
              return (image,)

         return ImagePipelineOutput(images=__a					)








from ...processing_utils import ProcessorMixin







class 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							_UpperCAmelCase      =						"""WhisperFeatureExtractor"""
							_UpperCAmelCase      =						"""WhisperTokenizer"""






							def __init__(      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										super().__init__(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.feature_extractor
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					False






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__=None			,						lowerCAmelCase__=None			,						lowerCAmelCase__=True       ):





										"""simple docstring"""




										return self.tokenizer.get_decoder_prompt_ids(task=lowerCAmelCase__			,						language=lowerCAmelCase__			,						no_timestamps=lowerCAmelCase__       )






							def __call__(      self			,						*lowerCAmelCase__			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										if self._in_target_context_manager:
													return self.current_processor(*lowerCAmelCase__			,						**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					kwargs.pop('audio'			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					kwargs.pop('sampling_rate'			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					kwargs.pop('text'			,						lowerCAmelCase__       )
										if len(lowerCAmelCase__       ) > 0:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					args[0]
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					args[1:]

										if audio is None and text is None:
													raise ValueError('You need to specify either an `audio` or `text` input to process.'       )

										if audio is not None:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					self.feature_extractor(lowerCAmelCase__			,						*lowerCAmelCase__			,						sampling_rate=lowerCAmelCase__			,						**lowerCAmelCase__       )
										if text is not None:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.tokenizer(lowerCAmelCase__			,						**lowerCAmelCase__       )

										if text is None:
													return inputs

										elif audio is None:
													return encodings
										else:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					encodings['input_ids']
													return inputs






							def 					UpperCamelCase__     (      self			,						*lowerCAmelCase__			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										return self.tokenizer.batch_decode(*lowerCAmelCase__			,						**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self			,						*lowerCAmelCase__			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										return self.tokenizer.decode(*lowerCAmelCase__			,						**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__="np"       ):





										"""simple docstring"""




										return self.tokenizer.get_prompt_ids(lowerCAmelCase__			,						return_tensors=lowerCAmelCase__       )





