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


A__   :					List[str]							=   logging.get_logger(__name__)  # pylint: disable=invalid-name

A__   :					Union[str, 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		(		_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase=8						):






  """simple docstring"""



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







class 			__magic_name__       (       SCREAMING_SNAKE_CASE__ ):
       def __init__(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				)     ->	Union[str, Any]:





         """simple docstring"""




         super().__init__()

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





         """simple docstring"""




         if latents is None:
           _lowercase:			Tuple       =     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}'''       )
           _lowercase:			int       =     latents.to(A_       )

         _lowercase:			Optional[Any]       =     latents * scheduler.init_noise_sigma
         return latents
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_=None ,				)     ->	Any:





         """simple docstring"""




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

         _lowercase:			Dict       =     text_inputs.input_ids
         _lowercase:			int       =     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_       ):
           _lowercase:			Optional[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}'''       )

         _lowercase:			List[str]       =     text_input_ids.to(A_       )
         _lowercase:			Union[str, Any]       =     text_inputs.attention_mask.to(A_       )

         _lowercase ,				_lowercase:			Optional[Any]       =     self.text_encoder(
             input_ids=A_ ,				attention_mask=A_       )

         _lowercase:			Tuple       =     prompt_embeds.repeat_interleave(A_ ,				dim=0       )
         _lowercase:			Tuple       =     text_encoder_hidden_states.repeat_interleave(A_ ,				dim=0       )
         _lowercase:			Union[str, Any]       =     text_mask.repeat_interleave(A_ ,				dim=0       )

         if do_classifier_free_guidance:
           _lowercase:			List[str]
           if negative_prompt is None:
             _lowercase:			List[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_       ):
             _lowercase:			List[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:
             _lowercase:			str       =     negative_prompt

           _lowercase:			int       =     self.tokenizer(
               A_ ,				padding='''max_length''' ,				max_length=77 ,				truncation=A_ ,				return_attention_mask=A_ ,				add_special_tokens=A_ ,				return_tensors='''pt''' ,				)
           _lowercase:			Dict       =     uncond_input.input_ids.to(A_       )
           _lowercase:			List[Any]       =     uncond_input.attention_mask.to(A_       )

           _lowercase ,				_lowercase:			int       =     self.text_encoder(
               input_ids=A_ ,				attention_mask=A_       )

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

           _lowercase:			str       =     negative_prompt_embeds.shape[1]
           _lowercase:			Tuple       =     negative_prompt_embeds.repeat(1 ,				A_       )
           _lowercase:			Tuple       =     negative_prompt_embeds.view(batch_size * num_images_per_prompt ,				A_       )

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

           _lowercase:			str       =     torch.cat([uncond_text_mask, text_mask]       )

         return prompt_embeds, text_encoder_hidden_states, text_mask
       def 							lowercase_							(     self ,				A_=0       )     ->	Dict:





         """simple docstring"""




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

         _lowercase:			Optional[Any]       =     torch.device(f'''cuda:{gpu_id}'''       )

         _lowercase:			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 							lowercase_							(     self ,				A_=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.'''       )

         _lowercase:			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)

         _lowercase:			Union[str, Any]       =     None
         for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
           _lowercase ,				_lowercase:			Any       =     cpu_offload_with_hook(A_ ,				A_ ,				prev_module_hook=A_       )

         if self.safety_checker is not None:
           _lowercase ,				_lowercase:			Dict       =     cpu_offload_with_hook(self.safety_checker ,				A_ ,				prev_module_hook=A_       )

         # We'll offload the last model manually.
         _lowercase:			int       =     hook
       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def 							lowercase_							(     self       )     ->	Tuple:





         """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 ,				)     ->	Any:





         """simple docstring"""




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

         _lowercase:			List[Any]       =     self._execution_device

         _lowercase:			Any       =     batch_size * num_images_per_prompt
         _lowercase:			Dict       =     guidance_scale > 1.0

         _lowercase ,				_lowercase ,				_lowercase:			Tuple       =     self._encode_prompt(
             A_ ,				A_ ,				A_ ,				A_ ,				A_       )

         if isinstance(A_ ,				A_       ):
           _lowercase:			Optional[Any]       =     torch.cat(A_ ,				dim=0       )
         if isinstance(A_ ,				A_       ):
           _lowercase:			int       =     torch.cat(A_ ,				dim=0       )

         if do_classifier_free_guidance:
           _lowercase:			List[Any]       =     image_embeds.repeat_interleave(A_ ,				dim=0       )
           _lowercase:			Union[str, Any]       =     negative_image_embeds.repeat_interleave(A_ ,				dim=0       )

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

         self.scheduler.set_timesteps(A_ ,				device=A_       )
         _lowercase:			Union[str, Any]       =     self.scheduler.timesteps

         _lowercase:			Optional[int]       =     self.unet.config.in_channels

         _lowercase ,				_lowercase:			Any       =     get_new_h_w(A_ ,				A_ ,				self.movq_scale_factor       )

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

           _lowercase:			List[Any]       =     {'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
           _lowercase:			List[Any]       =     self.unet(
               sample=A_ ,				timestep=A_ ,				encoder_hidden_states=A_ ,				added_cond_kwargs=A_ ,				return_dict=A_ ,				)[0]

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

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

         if output_type == "pil":
           _lowercase:			str       =     self.numpy_to_pil(A_       )

         if not return_dict:
           return (image,)

         return ImagePipelineOutput(images=A_       )


"""simple docstring"""




from __future__ import annotations

import numpy as np




def  _lowerCAmelCase		(		_UpperCamelCase						):






  """simple docstring"""



  _lowercase ,				_lowercase:			List[str]       =     np.shape(_UpperCamelCase						)
  if rows != columns:
    _lowercase:			int       =     (
        '''\'table\' has to be of square shaped array but got a '''
        f'''{rows}x{columns} array:\n{table}'''
    )
    raise ValueError(_UpperCamelCase						)

  _lowercase:			Tuple       =     np.zeros((rows, columns)						)
  _lowercase:			Tuple       =     np.zeros((rows, columns)						)
  for i in range(_UpperCamelCase						):
    for j in range(_UpperCamelCase						):
      _lowercase:			List[Any]       =     sum(lower[i][k] * upper[k][j] for k in range(_UpperCamelCase						)						)
      if upper[j][j] == 0:
        raise ArithmeticError('''No LU decomposition exists'''						)
      _lowercase:			Any       =     (table[i][j] - total) / upper[j][j]
    _lowercase:			Any       =     1
    for j in range(_UpperCamelCase ,				_UpperCamelCase						):
      _lowercase:			Optional[int]       =     sum(lower[i][k] * upper[k][j] for k in range(_UpperCamelCase						)						)
      _lowercase:			List[str]       =     table[i][j] - total
  return lower, upper


if __name__ == "__main__":
  import doctest

  doctest.testmod()
