




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


lowercase__				:      Optional[int]	   = logging.get_logger(__name__)  # pylint: disable=invalid-name

lowercase__				:      Dict	   = "\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  SCREAMING_SNAKE_CASE							(       __UpperCamelCase		, __UpperCamelCase		, __UpperCamelCase=8)      ->					Tuple:
     a     =				h // scale_factor**2
     if h % scale_factor**2 != 0:
          new_h += 1
     a     =				w // scale_factor**2
     if w % scale_factor**2 != 0:
          new_w += 1
     return new_h * scale_factor, new_w * scale_factor





class       a__	( SCREAMING_SNAKE_CASE__				):







     def __init__(     self  ,		A  ,		A  ,		A  ,		A  ,		A  ,		)  ->   Dict:



          '''simple docstring'''



          super().__init__()

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







     def       lowerCAmelCase_    (     self  ,		A  ,		A  ,		A  ,		A  ,		A  ,		A			)  ->   Tuple:



          '''simple docstring'''



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

          a     =				latents * scheduler.init_noise_sigma
          return latents







     def       lowerCAmelCase_    (     self  ,		A  ,		A  ,		A  ,		A  ,		A=None  ,		)  ->   Dict:



          '''simple docstring'''



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

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

          a     =				text_input_ids.to(A__			)
          a     =				text_inputs.attention_mask.to(A__			)

          a  ,				a     =				self.text_encoder(
              input_ids=A__  ,		attention_mask=A__			)

          a     =				prompt_embeds.repeat_interleave(A__  ,		dim=0			)
          a     =				text_encoder_hidden_states.repeat_interleave(A__  ,		dim=0			)
          a     =				text_mask.repeat_interleave(A__  ,		dim=0			)

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

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

               a  ,				a     =				self.text_encoder(
                   input_ids=A__  ,		attention_mask=A__			)

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

               a     =				negative_prompt_embeds.shape[1]
               a     =				negative_prompt_embeds.repeat(1  ,		A__			)
               a     =				negative_prompt_embeds.view(batch_size * num_images_per_prompt  ,		A__			)

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

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

          return prompt_embeds, text_encoder_hidden_states, text_mask







     def       lowerCAmelCase_    (     self  ,		A=0			)  ->   List[str]:



          '''simple docstring'''



          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.text_encoder,
              self.movq,
          ]
          for cpu_offloaded_model in models:
               if cpu_offloaded_model is not None:
                    cpu_offload(A__  ,		A__			)







     def       lowerCAmelCase_    (     self  ,		A=0			)  ->   Dict:



          '''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."			)

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

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

          if self.safety_checker is not None:
               a  ,				a     =				cpu_offload_with_hook(self.safety_checker  ,		A__  ,		prev_module_hook=A__			)

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







     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def       lowerCAmelCase_    (     self			)  ->   List[Any]:



          '''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  ,		)  ->   List[str]:



          '''simple docstring'''



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

          a     =				self._execution_device

          a     =				batch_size * num_images_per_prompt
          a     =				guidance_scale > 1.0

          a  ,				a  ,				a     =				self._encode_prompt(
              A__  ,		A__  ,		A__  ,		A__  ,		A__			)

          if isinstance(A__  ,		A__			):
               a     =				torch.cat(A__  ,		dim=0			)
          if isinstance(A__  ,		A__			):
               a     =				torch.cat(A__  ,		dim=0			)

          if do_classifier_free_guidance:
               a     =				image_embeds.repeat_interleave(A__  ,		dim=0			)
               a     =				negative_image_embeds.repeat_interleave(A__  ,		dim=0			)

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

          self.scheduler.set_timesteps(A__  ,		device=A__			)
          a     =				self.scheduler.timesteps

          a     =				self.unet.config.in_channels

          a  ,				a     =				get_new_h_w(A__  ,		A__  ,		self.movq_scale_factor			)

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

               a     =				{"text_embeds": prompt_embeds, "image_embeds": image_embeds}
               a     =				self.unet(
                   sample=A__  ,		timestep=A__  ,		encoder_hidden_states=A__  ,		added_cond_kwargs=A__  ,		return_dict=A__  ,		)[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(
                   A__  ,		A__  ,		A__  ,		generator=A__  ,		).prev_sample
          # post-processing
          a     =				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"]:
               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(A__			)

          if not return_dict:
               return (image,)

          return ImagePipelineOutput(images=A__			)







import numpy as np
from nltk.translate import meteor_score

import datasets
from datasets.config import importlib_metadata, version


lowercase__				:      Optional[int]	   = version.parse(importlib_metadata.version("nltk"))
if NLTK_VERSION >= version.Version("3.6.4"):
   from nltk import word_tokenize


lowercase__				:      int	   = "\\n@inproceedings{banarjee2005,\n  title     = {{METEOR}: An Automatic Metric for {MT} Evaluation with Improved Correlation with Human Judgments},\n  author    = {Banerjee, Satanjeev  and Lavie, Alon},\n  booktitle = {Proceedings of the {ACL} Workshop on Intrinsic and Extrinsic Evaluation Measures for Machine Translation and/or Summarization},\n  month     = jun,\n  year      = {2005},\n  address   = {Ann Arbor, Michigan},\n  publisher = {Association for Computational Linguistics},\n  url       = {https://www.aclweb.org/anthology/W05-0909},\n  pages     = {65--72},\n}\n"

lowercase__				:      str	   = "\\nMETEOR, an automatic metric for machine translation evaluation\nthat is based on a generalized concept of unigram matching between the\nmachine-produced translation and human-produced reference translations.\nUnigrams can be matched based on their surface forms, stemmed forms,\nand meanings; furthermore, METEOR can be easily extended to include more\nadvanced matching strategies. Once all generalized unigram matches\nbetween the two strings have been found, METEOR computes a score for\nthis matching using a combination of unigram-precision, unigram-recall, and\na measure of fragmentation that is designed to directly capture how\nwell-ordered the matched words in the machine translation are in relation\nto the reference.\n\nMETEOR gets an R correlation value of 0.347 with human evaluation on the Arabic\ndata and 0.331 on the Chinese data. This is shown to be an improvement on\nusing simply unigram-precision, unigram-recall and their harmonic F1\ncombination.\n"

lowercase__				:      Tuple	   = "\nComputes METEOR score of translated segments against one or more references.\nArgs:\n    predictions: list of predictions to score. Each prediction\n        should be a string with tokens separated by spaces.\n    references: list of reference for each prediction. Each\n        reference should be a string with tokens separated by spaces.\n    alpha: Parameter for controlling relative weights of precision and recall. default: 0.9\n    beta: Parameter for controlling shape of penalty as a function of fragmentation. default: 3\n    gamma: Relative weight assigned to fragmentation penalty. default: 0.5\nReturns:\n    'meteor': meteor score.\nExamples:\n\n    >>> meteor = datasets.load_metric('meteor')\n    >>> predictions = [\"It is a guide to action which ensures that the military always obeys the commands of the party\"]\n    >>> references = [\"It is a guide to action that ensures that the military will forever heed Party commands\"]\n    >>> results = meteor.compute(predictions=predictions, references=references)\n    >>> print(round(results[\"meteor\"], 4))\n    0.6944\n"




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION       , _KWARGS_DESCRIPTION				)
class       a__	( datasets.Metric				):







     def       lowerCAmelCase_    (     self			)  ->   List[Any]:



          '''simple docstring'''



          return datasets.MetricInfo(
              description=_DESCRIPTION  ,		citation=_CITATION  ,		inputs_description=_KWARGS_DESCRIPTION  ,		features=datasets.Features(
                  {
                      "predictions": datasets.Value("string"  ,		id="sequence"			),
                      "references": datasets.Value("string"  ,		id="sequence"			),
                  }			)  ,		codebase_urls=["https://github.com/nltk/nltk/blob/develop/nltk/translate/meteor_score.py"]  ,		reference_urls=[
                  "https://www.nltk.org/api/nltk.translate.html#module-nltk.translate.meteor_score",
                  "https://en.wikipedia.org/wiki/METEOR",
              ]  ,		)







     def       lowerCAmelCase_    (     self  ,		A			)  ->   List[str]:



          '''simple docstring'''



          import nltk

          nltk.download("wordnet"			)
          if NLTK_VERSION >= version.Version("3.6.5"			):
               nltk.download("punkt"			)
          if NLTK_VERSION >= version.Version("3.6.6"			):
               nltk.download("omw-1.4"			)







     def       lowerCAmelCase_    (     self  ,		A  ,		A  ,		A=0.9  ,		A=3  ,		A=0.5			)  ->   Tuple:



          '''simple docstring'''



          if NLTK_VERSION >= version.Version("3.6.5"			):
               a     =				[
                   meteor_score.single_meteor_score(
                       word_tokenize(A			)  ,		word_tokenize(A			)  ,		alpha=A  ,		beta=A  ,		gamma=A			)
                   for ref, pred in zip(A  ,		A			)
               ]
          else:
               a     =				[
                   meteor_score.single_meteor_score(A  ,		A  ,		alpha=A  ,		beta=A  ,		gamma=A			)
                   for ref, pred in zip(A  ,		A			)
               ]

          return {"meteor": np.mean(A			)}

