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


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

A__   :					List[str]							=   '\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)["depth"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline("depth-estimation")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to("cuda")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to("cuda")\n\n\n        >>> img = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/cat.png"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")\n\n        >>> prompt = "A robot, 4k photo"\n        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"\n\n        >>> generator = torch.Generator(device="cuda").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save("robot_cat.png")\n        ```\n'




def  _lowerCAmelCase		(		_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase=8						):






  """simple docstring"""



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







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





         """simple docstring"""




         super().__init__()

         self.register_modules(
             unet=A_ ,				scheduler=A_ ,				movq=A_ ,				)
         _lowercase:			Dict       =     2 ** (len(self.movq.config.block_out_channels       ) - 1)
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	str:





         """simple docstring"""




         if latents is None:
           _lowercase:			Optional[int]       =     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:			List[str]       =     latents.to(A_       )

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





         """simple docstring"""




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

         _lowercase:			str       =     torch.device(f'''cuda:{gpu_id}'''       )

         _lowercase:			List[str]       =     [
             self.unet,
             self.movq,
         ]
         for cpu_offloaded_model in models:
           if cpu_offloaded_model is not None:
             cpu_offload(A_ ,				A_       )
       def 							lowercase_							(     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.'''       )

         _lowercase:			Tuple       =     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:			Dict       =     None
         for cpu_offloaded_model in [self.unet, self.movq]:
           _lowercase ,				_lowercase:			Dict       =     cpu_offload_with_hook(A_ ,				A_ ,				prev_module_hook=A_       )

         # We'll offload the last model manually.
         _lowercase:			str       =     hook
       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def 							lowercase_							(     self       )     ->	Union[str, 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_ = 512 ,				A_ = 512 ,				A_ = 100 ,				A_ = 4.0 ,				A_ = 1 ,				A_ = None ,				A_ = None ,				A_ = "pil" ,				A_ = True ,				)     ->	Dict:





         """simple docstring"""




         _lowercase:			str       =     self._execution_device

         _lowercase:			Optional[Any]       =     guidance_scale > 1.0

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

         _lowercase:			Dict       =     image_embeds.shape[0] * num_images_per_prompt

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

         _lowercase:			Dict       =     torch.cat([negative_image_embeds, image_embeds] ,				dim=0       ).to(dtype=self.unet.dtype ,				device=A_       )
         _lowercase:			Tuple       =     torch.cat([hint, hint] ,				dim=0       ).to(dtype=self.unet.dtype ,				device=A_       )
         self.scheduler.set_timesteps(A_ ,				device=A_       )
         _lowercase:			Optional[Any]       =     self.scheduler.timesteps

         _lowercase:			str       =     self.movq.config.latent_channels

         _lowercase ,				_lowercase:			Union[str, Any]       =     downscale_height_and_width(A_ ,				A_ ,				self.movq_scale_factor       )

         # create initial latent
         _lowercase:			Any       =     self.prepare_latents(
             (batch_size, num_channels_latents, height, width) ,				image_embeds.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:			Optional[int]       =     torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents

           _lowercase:			Tuple       =     {'''image_embeds''': image_embeds, '''hint''': hint}
           _lowercase:			Dict       =     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:			Any       =     noise_pred.split(latents.shape[1] ,				dim=1       )
             _lowercase ,				_lowercase:			int       =     noise_pred.chunk(2       )
             _lowercase ,				_lowercase:			Tuple       =     variance_pred.chunk(2       )
             _lowercase:			Optional[Any]       =     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
             _lowercase:			str       =     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:			Tuple       =     noise_pred.split(latents.shape[1] ,				dim=1       )

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

         if output_type == "pil":
           _lowercase:			Optional[Any]       =     self.numpy_to_pil(A_       )

         if not return_dict:
           return (image,)

         return ImagePipelineOutput(images=A_       )


"""simple docstring"""




import comet  # From: unbabel-comet
import torch

import datasets


A__   :					int							=   datasets.logging.get_logger(__name__)

A__   :					Optional[Any]							=   '\\n@inproceedings{rei-EtAl:2020:WMT,\n   author    = {Rei, Ricardo  and  Stewart, Craig  and  Farinha, Ana C  and  Lavie, Alon},\n   title     = {Unbabel\'s Participation in the WMT20 Metrics Shared Task},\n   booktitle      = {Proceedings of the Fifth Conference on Machine Translation},\n   month          = {November},\n   year           = {2020},\n   address        = {Online},\n   publisher      = {Association for Computational Linguistics},\n   pages     = {909--918},\n}\n@inproceedings{rei-etal-2020-comet,\n   title = "{COMET}: A Neural Framework for {MT} Evaluation",\n   author = "Rei, Ricardo  and\n      Stewart, Craig  and\n      Farinha, Ana C  and\n      Lavie, Alon",\n   booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)",\n   month = nov,\n   year = "2020",\n   address = "Online",\n   publisher = "Association for Computational Linguistics",\n   url = "https://www.aclweb.org/anthology/2020.emnlp-main.213",\n   pages = "2685--2702",\n}\n'

A__   :					Union[str, Any]							=   '\\nCrosslingual Optimized Metric for Evaluation of Translation (COMET) is an open-source framework used to train Machine Translation metrics that achieve high levels of correlation with different types of human judgments (HTER, DA\'s or MQM).\nWith the release of the framework the authors also released fully trained models that were used to compete in the WMT20 Metrics Shared Task achieving SOTA in that years competition.\n\nSee the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.\n'

A__   :					int							=   '\nCOMET score.\n\nArgs:\n\n`sources` (list of str): Source sentences\n`predictions` (list of str): candidate translations\n`references` (list of str): reference translations\n`cuda` (bool): If set to True, runs COMET using GPU\n`show_progress` (bool): Shows progress\n`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.\n\nReturns:\n    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.\n    `scores`: List of scores.\n\nExamples:\n\n    >>> comet_metric = datasets.load_metric(\'comet\')\n    >>> # comet_metric = load_metric(\'comet\', \'wmt20-comet-da\')  # you can also choose which model to use\n    >>> source = ["Dem Feuer konnte Einhalt geboten werden", "Schulen und Kindergärten wurden eröffnet."]\n    >>> hypothesis = ["The fire could be stopped", "Schools and kindergartens were open"]\n    >>> reference = ["They were able to control the fire.", "Schools and kindergartens opened"]\n    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)\n    >>> print([round(v, 2) for v in results["scores"]])\n    [0.19, 0.92]\n'


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION   ,					_KWARGS_DESCRIPTION )
class 			__magic_name__       (       datasets.Metric ):
       def 							lowercase_							(     self       )     ->	Tuple:





         """simple docstring"""




         return datasets.MetricInfo(
             description=_DESCRIPTION ,				citation=_CITATION ,				homepage='''https://unbabel.github.io/COMET/html/index.html''' ,				inputs_description=_KWARGS_DESCRIPTION ,				features=datasets.Features(
                 {
                     '''sources''': datasets.Value('''string''' ,				id='''sequence'''       ),
                     '''predictions''': datasets.Value('''string''' ,				id='''sequence'''       ),
                     '''references''': datasets.Value('''string''' ,				id='''sequence'''       ),
                 }       ) ,				codebase_urls=['''https://github.com/Unbabel/COMET'''] ,				reference_urls=[
                 '''https://github.com/Unbabel/COMET''',
                 '''https://www.aclweb.org/anthology/2020.emnlp-main.213/''',
                 '''http://www.statmt.org/wmt20/pdf/2020.wmt-1.101.pdf6''',
             ] ,				)
       def 							lowercase_							(     self ,				A_       )     ->	Dict:





         """simple docstring"""




         if self.config_name == "default":
           _lowercase:			str       =     comet.load_from_checkpoint(comet.download_model('''wmt20-comet-da'''       )       )
         else:
           _lowercase:			Optional[Any]       =     comet.load_from_checkpoint(comet.download_model(self.config_name       )       )



       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_=None ,				A_=False       )     ->	List[str]:





         """simple docstring"""




         if gpus is None:
           _lowercase:			Union[str, Any]       =     1 if torch.cuda.is_available() else 0
         _lowercase:			Any       =     {'''src''': sources, '''mt''': predictions, '''ref''': references}
         _lowercase:			int       =     [dict(zip(A_ ,				A_       )       ) for t in zip(*data.values()       )]
         _lowercase ,				_lowercase:			Tuple       =     self.scorer.predict(A_ ,				gpus=A_ ,				progress_bar=A_       )
         return {"mean_score": mean_score, "scores": scores}
