





"""simple docstring"""






import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
  import onnxruntime as ort




class 			__lowercase     (     _UpperCamelCase				, unittest.TestCase							):







     '''simple docstring'''




     __lowerCAmelCase					=					'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''







     def 	_lowerCamelCase						(  self		,     _UpperCAmelCase=0       ):
         __a	:  List[Any]         =					np.random.RandomState(_lowerCAmelCase       )
         __a	:  List[str]         =					{
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs







     def 	_lowerCamelCase						(  self       ):
         __a	:  Dict         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  str         =					self.get_dummy_inputs()
         __a	:  List[Any]         =					pipe(**_lowerCAmelCase       ).images
         __a	:  List[Any]         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __a	:  Dict         =					np.array([0.6_5_0_7_2, 0.5_8_4_9_2, 0.4_8_2_1_9, 0.5_5_5_2_1, 0.5_3_1_8_0, 0.5_5_9_3_9, 0.5_0_6_9_7, 0.3_9_8_0_0, 0.4_6_4_5_5]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







     def 	_lowerCamelCase						(  self       ):
         __a	:  int         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         __a	:  Dict         =					PNDMScheduler.from_config(pipe.scheduler.config		,     skip_prk_steps=_lowerCAmelCase       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  int         =					self.get_dummy_inputs()
         __a	:  int         =					pipe(**_lowerCAmelCase       ).images
         __a	:  Dict         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __a	:  Optional[Any]         =					np.array([0.6_5_8_6_3, 0.5_9_4_2_5, 0.4_9_3_2_6, 0.5_6_3_1_3, 0.5_3_8_7_5, 0.5_6_6_2_7, 0.5_1_0_6_5, 0.3_9_7_7_7, 0.4_6_3_3_0]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







     def 	_lowerCamelCase						(  self       ):
         __a	:  str         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         __a	:  Any         =					LMSDiscreteScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Optional[Any]         =					self.get_dummy_inputs()
         __a	:  Union[str, Any]         =					pipe(**_lowerCAmelCase       ).images
         __a	:  List[str]         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __a	:  int         =					np.array([0.5_3_7_5_5, 0.6_0_7_8_6, 0.4_7_4_0_2, 0.4_9_4_8_8, 0.5_1_8_6_9, 0.4_9_8_1_9, 0.4_7_9_8_5, 0.3_8_9_5_7, 0.4_4_2_7_9]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







     def 	_lowerCamelCase						(  self       ):
         __a	:  List[str]         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         __a	:  Optional[int]         =					EulerDiscreteScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Optional[int]         =					self.get_dummy_inputs()
         __a	:  Any         =					pipe(**_lowerCAmelCase       ).images
         __a	:  str         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __a	:  List[Any]         =					np.array([0.5_3_7_5_5, 0.6_0_7_8_6, 0.4_7_4_0_2, 0.4_9_4_8_8, 0.5_1_8_6_9, 0.4_9_8_1_9, 0.4_7_9_8_5, 0.3_8_9_5_7, 0.4_4_2_7_9]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







     def 	_lowerCamelCase						(  self       ):
         __a	:  List[str]         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         __a	:  List[Any]         =					EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Tuple         =					self.get_dummy_inputs()
         __a	:  Optional[int]         =					pipe(**_lowerCAmelCase       ).images
         __a	:  int         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __a	:  int         =					np.array([0.5_3_8_1_7, 0.6_0_8_1_2, 0.4_7_3_8_4, 0.4_9_5_3_0, 0.5_1_8_9_4, 0.4_9_8_1_4, 0.4_7_9_8_4, 0.3_8_9_5_8, 0.4_4_2_7_1]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







     def 	_lowerCamelCase						(  self       ):
         __a	:  Any         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         __a	:  Optional[int]         =					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  int         =					self.get_dummy_inputs()
         __a	:  Optional[Any]         =					pipe(**_lowerCAmelCase       ).images
         __a	:  Union[str, Any]         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __a	:  Dict         =					np.array([0.5_3_8_9_5, 0.6_0_8_0_8, 0.4_7_9_3_3, 0.4_9_6_0_8, 0.5_1_8_8_6, 0.4_9_9_5_0, 0.4_8_0_5_3, 0.3_8_9_5_7, 0.4_4_2_0_0]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







     def 	_lowerCamelCase						(  self       ):
         __a	:  List[str]         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Any         =					self.get_dummy_inputs()
         __a	:  Union[str, Any]         =					3 * [inputs['''prompt''']]

         # forward
         __a	:  str         =					pipe(**_lowerCAmelCase       )
         __a	:  Dict         =					output.images[0, -3:, -3:, -1]

         __a	:  Optional[Any]         =					self.get_dummy_inputs()
         __a	:  List[Any]         =					3 * [inputs.pop('''prompt'''       )]

         __a	:  Any         =					pipe.tokenizer(
             _lowerCAmelCase		,     padding='''max_length'''		,     max_length=pipe.tokenizer.model_max_length		,     truncation=_lowerCAmelCase		,     return_tensors='''np'''		,     )
         __a	:  Any         =					text_inputs['''input_ids''']

         __a	:  List[Any]         =					pipe.text_encoder(input_ids=text_inputs.astype(np.intaa       )       )[0]

         __a	:  Tuple         =					prompt_embeds

         # forward
         __a	:  str         =					pipe(**_lowerCAmelCase       )
         __a	:  List[str]         =					output.images[0, -3:, -3:, -1]

         assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1e-4







     def 	_lowerCamelCase						(  self       ):
         __a	:  Tuple         =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''       )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Any         =					self.get_dummy_inputs()
         __a	:  int         =					3 * ['''this is a negative prompt''']
         __a	:  List[str]         =					negative_prompt
         __a	:  Union[str, Any]         =					3 * [inputs['''prompt''']]

         # forward
         __a	:  Tuple         =					pipe(**_lowerCAmelCase       )
         __a	:  Optional[Any]         =					output.images[0, -3:, -3:, -1]

         __a	:  Optional[Any]         =					self.get_dummy_inputs()
         __a	:  Union[str, Any]         =					3 * [inputs.pop('''prompt'''       )]

         __a	:  Optional[int]         =					[]
         for p in [prompt, negative_prompt]:
             __a	:  List[Any]         =					pipe.tokenizer(
                 _lowerCAmelCase		,     padding='''max_length'''		,     max_length=pipe.tokenizer.model_max_length		,     truncation=_lowerCAmelCase		,     return_tensors='''np'''		,     )
             __a	:  Any         =					text_inputs['''input_ids''']

             embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa       )       )[0]       )

         __a    ,  __a	:  Union[str, Any]         =					embeds

         # forward
         __a	:  Optional[Any]         =					pipe(**_lowerCAmelCase       )
         __a	:  Any         =					output.images[0, -3:, -3:, -1]

         assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1e-4






@nightly
@require_onnxruntime
@require_torch_gpu
class 			__lowercase     (     unittest.TestCase							):







     '''simple docstring'''




     @property
     def 	_lowerCamelCase						(  self       ):
         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )







     @property
     def 	_lowerCamelCase						(  self       ):
         __a	:  Dict         =					ort.SessionOptions()
         __a	:  Tuple         =					False
         return options







     def 	_lowerCamelCase						(  self       ):
         __a	:  Any         =					OnnxStableDiffusionPipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4'''		,     revision='''onnx'''		,     safety_checker=_lowerCAmelCase		,     feature_extractor=_lowerCAmelCase		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
         sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Optional[int]         =					'''A painting of a squirrel eating a burger'''
         np.random.seed(0       )
         __a	:  Optional[Any]         =					sd_pipe([prompt]		,     guidance_scale=6.0		,     num_inference_steps=10		,     output_type='''np'''       )
         __a	:  Any         =					output.images

         __a	:  Tuple         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         __a	:  Any         =					np.array([0.0_4_5_2, 0.0_3_9_0, 0.0_0_8_7, 0.0_3_5_0, 0.0_6_1_7, 0.0_3_6_4, 0.0_5_4_4, 0.0_5_2_3, 0.0_7_2_0]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-3







     def 	_lowerCamelCase						(  self       ):
         __a	:  Dict         =					DDIMScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''		,     subfolder='''scheduler'''		,     revision='''onnx'''       )
         __a	:  Union[str, Any]         =					OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''		,     revision='''onnx'''		,     scheduler=_lowerCAmelCase		,     safety_checker=_lowerCAmelCase		,     feature_extractor=_lowerCAmelCase		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
         sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Tuple         =					'''open neural network exchange'''
         __a	:  List[str]         =					np.random.RandomState(0       )
         __a	:  Any         =					sd_pipe([prompt]		,     guidance_scale=7.5		,     num_inference_steps=10		,     generator=_lowerCAmelCase		,     output_type='''np'''       )
         __a	:  Optional[Any]         =					output.images
         __a	:  str         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         __a	:  Union[str, Any]         =					np.array([0.2_8_6_7, 0.1_9_7_4, 0.1_4_8_1, 0.7_2_9_4, 0.7_2_5_1, 0.6_6_6_7, 0.4_1_9_4, 0.5_6_4_2, 0.6_4_8_6]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-3







     def 	_lowerCamelCase						(  self       ):
         __a	:  Any         =					LMSDiscreteScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''		,     subfolder='''scheduler'''		,     revision='''onnx'''       )
         __a	:  List[str]         =					OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''		,     revision='''onnx'''		,     scheduler=_lowerCAmelCase		,     safety_checker=_lowerCAmelCase		,     feature_extractor=_lowerCAmelCase		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
         sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  Union[str, Any]         =					'''open neural network exchange'''
         __a	:  List[Any]         =					np.random.RandomState(0       )
         __a	:  List[Any]         =					sd_pipe([prompt]		,     guidance_scale=7.5		,     num_inference_steps=10		,     generator=_lowerCAmelCase		,     output_type='''np'''       )
         __a	:  Optional[int]         =					output.images
         __a	:  int         =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         __a	:  Union[str, Any]         =					np.array([0.2_3_0_6, 0.1_9_5_9, 0.1_5_9_3, 0.6_5_4_9, 0.6_3_9_4, 0.5_4_0_8, 0.5_0_6_5, 0.6_0_1_0, 0.6_1_6_1]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-3







     def 	_lowerCamelCase						(  self       ):
         __a	:  Optional[int]         =					0

         def test_callback_fn(_UpperCAmelCase		,     _UpperCAmelCase		,     _UpperCAmelCase       ) -> None:
             __a	:  Optional[Any]         =					True
             nonlocal number_of_steps
             number_of_steps += 1
             if step == 0:
                 assert latents.shape == (1, 4, 64, 64)
                 __a	:  Any         =					latents[0, -3:, -3:, -1]
                 __a	:  Dict         =					np.array(
                     [-0.6_7_7_2, -0.3_8_3_5, -1.2_4_5_6, 0.1_9_0_5, -1.0_9_7_4, 0.6_9_6_7, -1.9_3_5_3, 0.0_1_7_8, 1.0_1_6_7]       )

                 assert np.abs(latents_slice.flatten() - expected_slice       ).max() < 1e-3
             elif step == 5:
                 assert latents.shape == (1, 4, 64, 64)
                 __a	:  int         =					latents[0, -3:, -3:, -1]
                 __a	:  List[str]         =					np.array(
                     [-0.3_3_5_1, 0.2_2_4_1, -0.1_8_3_7, -0.2_3_2_5, -0.6_5_7_7, 0.3_3_9_3, -0.0_2_4_1, 0.5_8_9_9, 1.3_8_7_5]       )

                 assert np.abs(latents_slice.flatten() - expected_slice       ).max() < 1e-3

         __a	:  Optional[int]         =					False

         __a	:  int         =					OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''		,     revision='''onnx'''		,     safety_checker=_lowerCAmelCase		,     feature_extractor=_lowerCAmelCase		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         __a	:  List[Any]         =					'''Andromeda galaxy in a bottle'''

         __a	:  str         =					np.random.RandomState(0       )
         pipe(
             prompt=_lowerCAmelCase		,     num_inference_steps=5		,     guidance_scale=7.5		,     generator=_lowerCAmelCase		,     callback=_lowerCAmelCase		,     callback_steps=1		,     )
         assert test_callback_fn.has_been_called
         assert number_of_steps == 6







     def 	_lowerCamelCase						(  self       ):
         __a	:  Dict         =					OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''		,     revision='''onnx'''		,     safety_checker=_lowerCAmelCase		,     feature_extractor=_lowerCAmelCase		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
         assert isinstance(_lowerCAmelCase		,     _lowerCAmelCase       )
         assert pipe.safety_checker is None

         __a	:  List[Any]         =					pipe('''example prompt'''		,     num_inference_steps=2       ).images[0]
         assert image is not None

         # check that there's no error when saving a pipeline with one of the models being None
         with tempfile.TemporaryDirectory() as tmpdirname:
             pipe.save_pretrained(_lowerCAmelCase       )
             __a	:  List[Any]         =					OnnxStableDiffusionPipeline.from_pretrained(_lowerCAmelCase       )

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         __a	:  Dict         =					pipe('''example prompt'''		,     num_inference_steps=2       ).images[0]
         assert image is not None






"""simple docstring"""






import comet  # From: unbabel-comet
import torch

import datasets


A	     =					datasets.logging.get_logger(__name__)

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

A	     =					'''\
Crosslingual 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).
With 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.

See the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.
'''

A	     =					'''
COMET score.

Args:

`sources` (list of str): Source sentences
`predictions` (list of str): candidate translations
`references` (list of str): reference translations
`cuda` (bool): If set to True, runs COMET using GPU
`show_progress` (bool): Shows progress
`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.

Returns:
    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.
    `scores`: List of scores.

Examples:

    >>> comet_metric = datasets.load_metric(\'comet\')
    >>> # comet_metric = load_metric(\'comet\', \'wmt20-comet-da\')  # you can also choose which model to use
    >>> source = ["Dem Feuer konnte Einhalt geboten werden", "Schulen und Kindergärten wurden eröffnet."]
    >>> hypothesis = ["The fire could be stopped", "Schools and kindergartens were open"]
    >>> reference = ["They were able to control the fire.", "Schools and kindergartens opened"]
    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)
    >>> print([round(v, 2) for v in results["scores"]])
    [0.19, 0.92]
'''




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







					'''simple docstring'''




					def 	_lowerCamelCase						(  self       ):
									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 	_lowerCamelCase						(  self		,     _UpperCAmelCase       ):
									if self.config_name == "default":
													__a	:  List[str]         =					comet.load_from_checkpoint(comet.download_model('''wmt20-comet-da'''       )       )
									else:
													__a	:  List[str]         =					comet.load_from_checkpoint(comet.download_model(self.config_name       )       )







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