
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


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

UpperCamelCase     :    Tuple       =			"\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 							A						(    snake_case				:List[str]    , snake_case				:Union[str, Any]    , snake_case				:Optional[Any]=8		)				->     Dict:
 __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 __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):





     def __init__(     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					):

      '''simple docstring'''




      super().__init__()

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





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      if latents is None:
       __UpperCamelCase      =	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}'				)
       __UpperCamelCase      =	latents.to(__UpperCAmelCase				)

      __UpperCamelCase      =	latents * scheduler.init_noise_sigma
      return latents





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase=None     ,					):

      '''simple docstring'''




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

      __UpperCamelCase      =	text_inputs.input_ids
      __UpperCamelCase      =	self.tokenizer(__UpperCAmelCase     ,					padding='longest'     ,					return_tensors='pt'				).input_ids

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

      __UpperCamelCase  ,     __UpperCamelCase      =	self.text_encoder(
          input_ids=__UpperCAmelCase     ,					attention_mask=__UpperCAmelCase				)

      __UpperCamelCase      =	prompt_embeds.repeat_interleave(__UpperCAmelCase     ,					dim=0				)
      __UpperCamelCase      =	text_encoder_hidden_states.repeat_interleave(__UpperCAmelCase     ,					dim=0				)
      __UpperCamelCase      =	text_mask.repeat_interleave(__UpperCAmelCase     ,					dim=0				)

      if do_classifier_free_guidance:
       __UpperCamelCase      =	42
       if negative_prompt is None:
        __UpperCamelCase      =	[''] * batch_size
       elif type(__UpperCAmelCase				) is not type(__UpperCAmelCase				):
        raise TypeError(
            F'`negative_prompt` should be the same type to `prompt`, but got {type(__UpperCAmelCase				)} !='
            F' {type(__UpperCAmelCase				)}.'				)
       elif isinstance(__UpperCAmelCase     ,					__UpperCAmelCase				):
        __UpperCamelCase      =	[negative_prompt]
       elif batch_size != len(__UpperCAmelCase				):
        raise ValueError(
            F'`negative_prompt`: {negative_prompt} has batch size {len(__UpperCAmelCase				)}, 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(
           __UpperCAmelCase     ,					padding='max_length'     ,					max_length=77     ,					truncation=__UpperCAmelCase     ,					return_attention_mask=__UpperCAmelCase     ,					add_special_tokens=__UpperCAmelCase     ,					return_tensors='pt'     ,					)
       __UpperCamelCase      =	uncond_input.input_ids.to(__UpperCAmelCase				)
       __UpperCamelCase      =	uncond_input.attention_mask.to(__UpperCAmelCase				)

       __UpperCamelCase  ,     __UpperCamelCase      =	self.text_encoder(
           input_ids=__UpperCAmelCase     ,					attention_mask=__UpperCAmelCase				)

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

       __UpperCamelCase      =	negative_prompt_embeds.shape[1]
       __UpperCamelCase      =	negative_prompt_embeds.repeat(1     ,					__UpperCAmelCase				)
       __UpperCamelCase      =	negative_prompt_embeds.view(batch_size * num_images_per_prompt     ,					__UpperCAmelCase				)

       __UpperCamelCase      =	uncond_text_encoder_hidden_states.shape[1]
       __UpperCamelCase      =	uncond_text_encoder_hidden_states.repeat(1     ,					__UpperCAmelCase     ,					1				)
       __UpperCamelCase      =	uncond_text_encoder_hidden_states.view(
           batch_size * num_images_per_prompt     ,					__UpperCAmelCase     ,					-1				)
       __UpperCamelCase      =	uncond_text_mask.repeat_interleave(__UpperCAmelCase     ,					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 						UpperCAmelCase (     self     ,					__UpperCAmelCase=0				):

      '''simple docstring'''




      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(__UpperCAmelCase     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase=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      =	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)

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

      if self.safety_checker is not None:
       __UpperCamelCase  ,     __UpperCamelCase      =	cpu_offload_with_hook(self.safety_checker     ,					__UpperCAmelCase     ,					prev_module_hook=__UpperCAmelCase				)

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





     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      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     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = 512     ,					__UpperCAmelCase = 512     ,					__UpperCAmelCase = 100     ,					__UpperCAmelCase = 4.0     ,					__UpperCAmelCase = 1     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = "pil"     ,					__UpperCAmelCase = True     ,					):

      '''simple docstring'''




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

      __UpperCamelCase      =	self._execution_device

      __UpperCamelCase      =	batch_size * num_images_per_prompt
      __UpperCamelCase      =	guidance_scale > 1.0

      __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	self._encode_prompt(
          __UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)

      if isinstance(__UpperCAmelCase     ,					__UpperCAmelCase				):
       __UpperCamelCase      =	torch.cat(__UpperCAmelCase     ,					dim=0				)
      if isinstance(__UpperCAmelCase     ,					__UpperCAmelCase				):
       __UpperCamelCase      =	torch.cat(__UpperCAmelCase     ,					dim=0				)

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

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

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

      __UpperCamelCase      =	self.unet.config.in_channels

      __UpperCamelCase  ,     __UpperCamelCase      =	get_new_h_w(__UpperCAmelCase     ,					__UpperCAmelCase     ,					self.movq_scale_factor				)

      # create initial latent
      __UpperCamelCase      =	self.prepare_latents(
          (batch_size, num_channels_latents, height, width)     ,					text_encoder_hidden_states.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
       __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=__UpperCAmelCase     ,					timestep=__UpperCAmelCase     ,					encoder_hidden_states=__UpperCAmelCase     ,					added_cond_kwargs=__UpperCAmelCase     ,					return_dict=__UpperCAmelCase     ,					)[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(
           __UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					generator=__UpperCAmelCase     ,					).prev_sample
      # post-processing
      __UpperCamelCase      =	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"]:
       __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(__UpperCAmelCase				)

      if not return_dict:
       return (image,)

      return ImagePipelineOutput(images=__UpperCAmelCase				)





"""simple docstring"""

import json
import re
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
      if is_torch_available():
            import torch
      if is_tf_available():
            import tensorflow as tf

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_codegen import CodeGenTokenizer


UpperCamelCase     :    List[Any]       =			logging.get_logger(__name__)

UpperCamelCase     :    List[str]       =			{"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

UpperCamelCase     :    List[Any]       =			{
    "vocab_file": {
        "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json",
    },
    "merges_file": {
        "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "Salesforce/codegen-350M-mono": (
            "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/tokenizer.json"
        ),
    },
}

UpperCamelCase     :    List[Any]       =			{
    "Salesforce/codegen-350M-mono": 2_0_4_8,
}

class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):
     lowercase						       =						VOCAB_FILES_NAMES
     lowercase						       =						PRETRAINED_VOCAB_FILES_MAP
     lowercase						       =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     lowercase						       =						["input_ids", "attention_mask"]
     lowercase						       =						CodeGenTokenizer





     def __init__(     self     ,					__UpperCAmelCase=None     ,					__UpperCAmelCase=None     ,					__UpperCAmelCase=None     ,					__UpperCAmelCase="<|endoftext|>"     ,					__UpperCAmelCase="<|endoftext|>"     ,					__UpperCAmelCase="<|endoftext|>"     ,					__UpperCAmelCase=False     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      super().__init__(
          __UpperCAmelCase     ,					__UpperCAmelCase     ,					tokenizer_file=__UpperCAmelCase     ,					unk_token=__UpperCAmelCase     ,					bos_token=__UpperCAmelCase     ,					eos_token=__UpperCAmelCase     ,					add_prefix_space=__UpperCAmelCase     ,					**__UpperCAmelCase     ,					)

      if kwargs.pop('add_bos_token'     ,					__UpperCAmelCase				):
       __UpperCamelCase      =	kwargs.pop('name_or_path'     ,					''				)
       raise ValueError(
           'Currenty GPT2\'s fast tokenizer does NOT support adding a BOS token.'
           'Instead you should use GPT2\'s slow tokenizer class `CodeGenTokenizer` as follows: \n'
           F'`CodeGenTokenizer.from_pretrained(\'{model_id}\')`\nor\n'
           F'`AutoTokenizer.from_pretrained(\'{model_id}\', use_fast=False)`\n'
           'This issue will be fixed soon, see: https://github.com/huggingface/tokenizers/pull/1005.'
           ' so that the fast tokenizer works correctly.'				)

      __UpperCamelCase      =	json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()				)
      if pre_tok_state.get('add_prefix_space'     ,					__UpperCAmelCase				) != add_prefix_space:
       __UpperCamelCase      =	getattr(__UpperCAmelCase     ,					pre_tok_state.pop('type'				)				)
       __UpperCamelCase      =	add_prefix_space
       __UpperCamelCase      =	pre_tok_class(**__UpperCAmelCase				)

      __UpperCamelCase      =	add_prefix_space





     def 						UpperCAmelCase (     self     ,					*__UpperCAmelCase     ,					**__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	kwargs.get('is_split_into_words'     ,					__UpperCAmelCase				)
      assert self.add_prefix_space or not is_split_into_words, (
          F'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
          "to use it with pretokenized inputs."
      )

      return super()._batch_encode_plus(*__UpperCAmelCase     ,					**__UpperCAmelCase				)





     def 						UpperCAmelCase (     self     ,					*__UpperCAmelCase     ,					**__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	kwargs.get('is_split_into_words'     ,					__UpperCAmelCase				)

      assert self.add_prefix_space or not is_split_into_words, (
          F'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
          "to use it with pretokenized inputs."
      )

      return super()._encode_plus(*__UpperCAmelCase     ,					**__UpperCAmelCase				)





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase = None				):

      '''simple docstring'''




      __UpperCamelCase      =	self._tokenizer.model.save(__UpperCAmelCase     ,					name=__UpperCAmelCase				)
      return tuple(__UpperCAmelCase				)





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase = False     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = None     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      __UpperCamelCase      =	super().decode(
          token_ids=__UpperCAmelCase     ,					skip_special_tokens=__UpperCAmelCase     ,					clean_up_tokenization_spaces=__UpperCAmelCase     ,					**__UpperCAmelCase     ,					)

      if truncate_before_pattern is not None and len(__UpperCAmelCase				) > 0:
       __UpperCamelCase      =	self.truncate(__UpperCAmelCase     ,					__UpperCAmelCase				)

      return decoded_text







     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      def find_re(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):
       __UpperCamelCase      =	pattern.search(__UpperCAmelCase     ,					__UpperCAmelCase				)
       return m.start() if m else -1

      __UpperCamelCase      =	[re.compile(__UpperCAmelCase     ,					re.MULTILINE				) for pattern in truncate_before_pattern]

      __UpperCamelCase      =	list(re.finditer('^print'     ,					__UpperCAmelCase     ,					re.MULTILINE				)				)

      if len(__UpperCAmelCase				) > 1:
       __UpperCamelCase      =	completion[: prints[1].start()]

      __UpperCamelCase      =	list(re.finditer('^def'     ,					__UpperCAmelCase     ,					re.MULTILINE				)				)

      if len(__UpperCAmelCase				) > 1:
       __UpperCamelCase      =	completion[: defs[1].start()]

      __UpperCamelCase      =	0

      __UpperCamelCase      =	[
          pos for pos in [find_re(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				) for terminal in terminals] if pos != -1
      ]

      if len(__UpperCAmelCase				) > 0:
       return completion[: min(__UpperCAmelCase				)]
      else:
       return completion



