


"""simple docstring"""






from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


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

__snake_case		=       '''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''
def       A_   (   _lowerCAmelCase     :					Dict,     _lowerCAmelCase     :					Dict,     _lowerCAmelCase     :					Dict=8       ):






     """simple docstring"""

     _a    	=				height // scale_factor**2
     if height % scale_factor**2 != 0:
          new_height += 1
     _a    	=				width // scale_factor**2
     if width % scale_factor**2 != 0:
          new_width += 1
     return new_height * scale_factor, new_width * scale_factor


class     __lowerCamelCase  (					a__				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				)       ->							Optional[Any]:
          super().__init__()

          self.register_modules(
              unet=__UpperCAmelCase    ,				scheduler=__UpperCAmelCase    ,				movq=__UpperCAmelCase    ,				)
          _a    	=				2 ** (len(self.movq.config.block_out_channels							) - 1)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Tuple:
          if latents is None:
               _a    	=				randn_tensor(__UpperCAmelCase    ,				generator=__UpperCAmelCase    ,				device=__UpperCAmelCase    ,				dtype=__UpperCAmelCase							)
          else:
               if latents.shape != shape:
                    raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'							)
               _a    	=				latents.to(__UpperCAmelCase							)

          _a    	=				latents * scheduler.init_noise_sigma
          return latents


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase=0							)       ->							str:
          if is_accelerate_available():
               from accelerate import cpu_offload
          else:
               raise ImportError('''Please install accelerate via `pip install accelerate`'''							)

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

          _a    	=				[
              self.unet,
              self.movq,
          ]
          for cpu_offloaded_model in models:
               if cpu_offloaded_model is not None:
                    cpu_offload(__UpperCAmelCase    ,				__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase=0							)       ->							Dict:
          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.'''							)

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

          if self.device.type != "cpu":
               self.to('''cpu'''    ,				silence_dtype_warnings=__UpperCAmelCase							)
               torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

          _a    	=				None
          for cpu_offloaded_model in [self.unet, self.movq]:
               _a	,      _a    	=				cpu_offload_with_hook(__UpperCAmelCase    ,				__UpperCAmelCase    ,				prev_module_hook=__UpperCAmelCase							)

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


     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def 							_UpperCAmelCase      (      self							)       ->							int:
          if not hasattr(self.unet    ,				'''_hf_hook'''							):
               return self.device
          for module in self.unet.modules():
               if (
                   hasattr(__UpperCAmelCase    ,				'''_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(__UpperCAmelCase							)
     def __call__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase = 512    ,				__UpperCAmelCase = 512    ,				__UpperCAmelCase = 100    ,				__UpperCAmelCase = 4.0    ,				__UpperCAmelCase = 1    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = "pil"    ,				__UpperCAmelCase = True    ,				)       ->							str:
          _a    	=				self._execution_device

          _a    	=				guidance_scale > 1.0

          if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							):
               _a    	=				torch.cat(__UpperCAmelCase    ,				dim=0							)
          _a    	=				image_embeds.shape[0] * num_images_per_prompt
          if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							):
               _a    	=				torch.cat(__UpperCAmelCase    ,				dim=0							)

          if do_classifier_free_guidance:
               _a    	=				image_embeds.repeat_interleave(__UpperCAmelCase    ,				dim=0							)
               _a    	=				negative_image_embeds.repeat_interleave(__UpperCAmelCase    ,				dim=0							)

          _a    	=				torch.cat([negative_image_embeds, image_embeds]    ,				dim=0							).to(dtype=self.unet.dtype    ,				device=__UpperCAmelCase							)

          self.scheduler.set_timesteps(__UpperCAmelCase    ,				device=__UpperCAmelCase							)
          _a    	=				self.scheduler.timesteps

          _a    	=				self.unet.config.in_channels

          _a	,      _a    	=				downscale_height_and_width(__UpperCAmelCase    ,				__UpperCAmelCase    ,				self.movq_scale_factor							)

          # create initial latent
          _a    	=				self.prepare_latents(
              (batch_size, num_channels_latents, height, width)    ,				image_embeds.dtype    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				self.scheduler    ,				)

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

               _a    	=				{'''image_embeds''': image_embeds}
               _a    	=				self.unet(
                   sample=__UpperCAmelCase    ,				timestep=__UpperCAmelCase    ,				encoder_hidden_states=__UpperCAmelCase    ,				added_cond_kwargs=__UpperCAmelCase    ,				return_dict=__UpperCAmelCase    ,				)[0]

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

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

          if output_type == "pil":
               _a    	=				self.numpy_to_pil(__UpperCAmelCase							)

          if not return_dict:
               return (image,)

          return ImagePipelineOutput(images=__UpperCAmelCase							)



"""simple docstring"""






import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


__snake_case		=       logging.get_logger(__name__)
__snake_case		=       {'''vocab_file''': '''spiece.model'''}

__snake_case		=       {
    '''vocab_file''': {
        '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/spiece.model''',
        '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/spiece.model''',
        '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model''',
        '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model''',
        '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/spiece.model''',
        '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/spiece.model''',
        '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model''',
        '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model''',
    }
}

__snake_case		=       {
    '''albert-base-v1''': 512,
    '''albert-large-v1''': 512,
    '''albert-xlarge-v1''': 512,
    '''albert-xxlarge-v1''': 512,
    '''albert-base-v2''': 512,
    '''albert-large-v2''': 512,
    '''albert-xlarge-v2''': 512,
    '''albert-xxlarge-v2''': 512,
}

__snake_case		=       '''▁'''


class     __lowerCamelCase  (					a__				):

     '''simple docstring'''
     A_       :	Optional[int]    =		VOCAB_FILES_NAMES
     A_       :	int    =		PRETRAINED_VOCAB_FILES_MAP
     A_       :	Optional[int]    =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase=True    ,				__UpperCAmelCase=True    ,				__UpperCAmelCase=False    ,				__UpperCAmelCase="[CLS]"    ,				__UpperCAmelCase="[SEP]"    ,				__UpperCAmelCase="<unk>"    ,				__UpperCAmelCase="[SEP]"    ,				__UpperCAmelCase="<pad>"    ,				__UpperCAmelCase="[CLS]"    ,				__UpperCAmelCase="[MASK]"    ,				__UpperCAmelCase = None    ,				**__UpperCAmelCase    ,				)       ->							None:
          # Mask token behave like a normal word, i.e. include the space before it and
          # is included in the raw text, there should be a match in a non-normalized sentence.
          _a    	=				(
              AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase    ,				normalized=__UpperCAmelCase							)
              if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							)
              else mask_token
          )

          _a    	=				{} if sp_model_kwargs is None else sp_model_kwargs

          super().__init__(
              do_lower_case=__UpperCAmelCase    ,				remove_space=__UpperCAmelCase    ,				keep_accents=__UpperCAmelCase    ,				bos_token=__UpperCAmelCase    ,				eos_token=__UpperCAmelCase    ,				unk_token=__UpperCAmelCase    ,				sep_token=__UpperCAmelCase    ,				pad_token=__UpperCAmelCase    ,				cls_token=__UpperCAmelCase    ,				mask_token=__UpperCAmelCase    ,				sp_model_kwargs=self.sp_model_kwargs    ,				**__UpperCAmelCase    ,				)

          _a    	=				do_lower_case
          _a    	=				remove_space
          _a    	=				keep_accents
          _a    	=				vocab_file

          _a    	=				spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(__UpperCAmelCase							)


     @property
     def 							_UpperCAmelCase      (      self							)       ->							str:
          return len(self.sp_model							)


     def 							_UpperCAmelCase      (      self							)       ->							Optional[Any]:
          _a    	=				{self.convert_ids_to_tokens(__UpperCAmelCase							): i for i in range(self.vocab_size							)}
          vocab.update(self.added_tokens_encoder							)
          return vocab


     def __getstate__(      self							)       ->							int:
          _a    	=				self.__dict__.copy()
          _a    	=				None
          return state


     def __setstate__(      self    ,				__UpperCAmelCase							)       ->							Tuple:
          _a    	=				d

          # for backward compatibility
          if not hasattr(self    ,				'''sp_model_kwargs'''							):
               _a    	=				{}

          _a    	=				spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(self.vocab_file							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          if self.remove_space:
               _a    	=				''' '''.join(inputs.strip().split()							)
          else:
               _a    	=				inputs
          _a    	=				outputs.replace('''``'''    ,				'''"'''							).replace('''\'\''''    ,				'''"'''							)

          if not self.keep_accents:
               _a    	=				unicodedata.normalize('''NFKD'''    ,				__UpperCAmelCase							)
               _a    	=				''''''.join([c for c in outputs if not unicodedata.combining(__UpperCAmelCase							)]							)
          if self.do_lower_case:
               _a    	=				outputs.lower()

          return outputs


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          _a    	=				self.preprocess_text(__UpperCAmelCase							)
          _a    	=				self.sp_model.encode(__UpperCAmelCase    ,				out_type=__UpperCAmelCase							)
          _a    	=				[]
          for piece in pieces:
               if len(__UpperCAmelCase							) > 1 and piece[-1] == str(''','''							) and piece[-2].isdigit():
                    _a    	=				self.sp_model.EncodeAsPieces(piece[:-1].replace(__UpperCAmelCase    ,				''''''							)							)
                    if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
                         if len(cur_pieces[0]							) == 1:
                              _a    	=				cur_pieces[1:]
                         else:
                              _a    	=				cur_pieces[0][1:]
                    cur_pieces.append(piece[-1]							)
                    new_pieces.extend(__UpperCAmelCase							)
               else:
                    new_pieces.append(__UpperCAmelCase							)

          return new_pieces


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							str:
          return self.sp_model.PieceToId(__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							int:
          return self.sp_model.IdToPiece(__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Optional[Any]:
          _a    	=				[]
          _a    	=				''''''
          _a    	=				False
          for token in tokens:
               # make sure that special tokens are not decoded using sentencepiece model
               if token in self.all_special_tokens:
                    if not prev_is_special:
                         out_string += " "
                    out_string += self.sp_model.decode(__UpperCAmelCase							) + token
                    _a    	=				True
                    _a    	=				[]
               else:
                    current_sub_tokens.append(__UpperCAmelCase							)
                    _a    	=				False
          out_string += self.sp_model.decode(__UpperCAmelCase							)
          return out_string.strip()


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				[self.sep_token_id]
          _a    	=				[self.cls_token_id]
          if token_ids_a is None:
               return cls + token_ids_a + sep
          return cls + token_ids_a + sep + token_ids_a + sep


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = False							)       ->							List[int]:

          if already_has_special_tokens:
               return super().get_special_tokens_mask(
                   token_ids_a=__UpperCAmelCase    ,				token_ids_a=__UpperCAmelCase    ,				already_has_special_tokens=__UpperCAmelCase							)

          if token_ids_a is not None:
               return [1] + ([0] * len(__UpperCAmelCase							)) + [1] + ([0] * len(__UpperCAmelCase							)) + [1]
          return [1] + ([0] * len(__UpperCAmelCase							)) + [1]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				[self.sep_token_id]
          _a    	=				[self.cls_token_id]

          if token_ids_a is None:
               return len(cls + token_ids_a + sep							) * [0]
          return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1]





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							Tuple[str]:
          if not os.path.isdir(__UpperCAmelCase							):
               logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
               return
          _a    	=				os.path.join(
              __UpperCAmelCase    ,				(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']							)

          if os.path.abspath(self.vocab_file							) != os.path.abspath(__UpperCAmelCase							) and os.path.isfile(self.vocab_file							):
               copyfile(self.vocab_file    ,				__UpperCAmelCase							)
          elif not os.path.isfile(self.vocab_file							):
               with open(__UpperCAmelCase    ,				'''wb'''							) as fi:
                    _a    	=				self.sp_model.serialized_model_proto()
                    fi.write(__UpperCAmelCase							)

          return (out_vocab_file,)