




'''simple docstring'''


import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
				import onnxruntime as ort

class     A					(	_a  ,unittest.TestCase					):
							lowercase_           =       'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'







							def      __lowerCAmelCase    (						self :						Tuple							,    lowerCAmelCase_ :						Tuple=0      )    ->							int:


														"""simple docstring"""


														_a			  =       floats_tensor((1, 3, 1_28, 1_28)							,    rng=random.Random(lowerCAmelCase_      )      )
														_a			  =       np.random.RandomState(lowerCAmelCase_      )
														_a			  =       {
														    '''prompt''': '''A painting of a squirrel eating a burger''',
														    '''image''': image,
														    '''generator''': generator,
														    '''num_inference_steps''': 3,
														    '''strength''': 0.7_5,
														    '''guidance_scale''': 7.5,
														    '''output_type''': '''numpy''',
														}
														return inputs







							def      __lowerCAmelCase    (						self :						Any      )    ->							Optional[int]:


														"""simple docstring"""


														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint							,    provider='''CPUExecutionProvider'''      )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														_a			  =       self.get_dummy_inputs()
														_a			  =       pipe(**lowerCAmelCase_      ).images
														_a			  =       image[0, -3:, -3:, -1].flatten()

														assert image.shape == (1, 1_28, 1_28, 3)
														_a			  =       np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]      )
														assert np.abs(image_slice - expected_slice      ).max() < 1e-1







							def      __lowerCAmelCase    (						self :						Union[str, Any]      )    ->							List[Any]:


														"""simple docstring"""


														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint							,    provider='''CPUExecutionProvider'''      )
														_a			  =       PNDMScheduler.from_config(pipe.scheduler.config							,    skip_prk_steps=lowerCAmelCase_      )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														_a			  =       self.get_dummy_inputs()
														_a			  =       pipe(**lowerCAmelCase_      ).images
														_a			  =       image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_28, 1_28, 3)
														_a			  =       np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]      )

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







							def      __lowerCAmelCase    (						self :						List[Any]      )    ->							str:


														"""simple docstring"""


														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint							,    provider='''CPUExecutionProvider'''      )
														_a			  =       LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														# warmup pass to apply optimizations
														_a			  =       pipe(**self.get_dummy_inputs()      )

														_a			  =       self.get_dummy_inputs()
														_a			  =       pipe(**lowerCAmelCase_      ).images
														_a			  =       image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_28, 1_28, 3)
														_a			  =       np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]      )

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







							def      __lowerCAmelCase    (						self :						Union[str, Any]      )    ->							Tuple:


														"""simple docstring"""


														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint							,    provider='''CPUExecutionProvider'''      )
														_a			  =       EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														_a			  =       self.get_dummy_inputs()
														_a			  =       pipe(**lowerCAmelCase_      ).images
														_a			  =       image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_28, 1_28, 3)
														_a			  =       np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

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







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


														"""simple docstring"""


														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint							,    provider='''CPUExecutionProvider'''      )
														_a			  =       EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														_a			  =       self.get_dummy_inputs()
														_a			  =       pipe(**lowerCAmelCase_      ).images
														_a			  =       image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_28, 1_28, 3)
														_a			  =       np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

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







							def      __lowerCAmelCase    (						self :						Optional[int]      )    ->							Dict:


														"""simple docstring"""


														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint							,    provider='''CPUExecutionProvider'''      )
														_a			  =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														_a			  =       self.get_dummy_inputs()
														_a			  =       pipe(**lowerCAmelCase_      ).images
														_a			  =       image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_28, 1_28, 3)
														_a			  =       np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]      )

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



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







							@property
							def      __lowerCAmelCase    (						self :						Optional[Any]      )    ->							Any:


														"""simple docstring"""


														return (
														    "CUDAExecutionProvider",
														    {
														        "gpu_mem_limit": "15000000000",  # 15GB
														        "arena_extend_strategy": "kSameAsRequested",
														    },
														)







							@property
							def      __lowerCAmelCase    (						self :						Optional[int]      )    ->							Any:


														"""simple docstring"""


														_a			  =       ort.SessionOptions()
														_a			  =       False
														return options







							def      __lowerCAmelCase    (						self :						Optional[int]      )    ->							Any:


														"""simple docstring"""


														_a			  =       load_image(
														    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
														    '''/img2img/sketch-mountains-input.jpg'''      )
														_a			  =       init_image.resize((7_68, 5_12)      )
														# using the PNDM scheduler by default
														_a			  =       OnnxStableDiffusionImgaImgPipeline.from_pretrained(
														    '''CompVis/stable-diffusion-v1-4'''							,    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			  =       '''A fantasy landscape, trending on artstation'''

														_a			  =       np.random.RandomState(0      )
														_a			  =       pipe(
														    prompt=lowerCAmelCase_							,    image=lowerCAmelCase_							,    strength=0.7_5							,    guidance_scale=7.5							,    num_inference_steps=10							,    generator=lowerCAmelCase_							,    output_type='''np'''							,    )
														_a			  =       output.images
														_a			  =       images[0, 2_55:2_58, 3_83:3_86, -1]

														assert images.shape == (1, 5_12, 7_68, 3)
														_a			  =       np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]      )
														# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







							def      __lowerCAmelCase    (						self :						List[str]      )    ->							Tuple:


														"""simple docstring"""


														_a			  =       load_image(
														    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
														    '''/img2img/sketch-mountains-input.jpg'''      )
														_a			  =       init_image.resize((7_68, 5_12)      )
														_a			  =       LMSDiscreteScheduler.from_pretrained(
														    '''runwayml/stable-diffusion-v1-5'''							,    subfolder='''scheduler'''							,    revision='''onnx'''      )
														_a			  =       OnnxStableDiffusionImgaImgPipeline.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							,    )
														pipe.set_progress_bar_config(disable=lowerCAmelCase_      )

														_a			  =       '''A fantasy landscape, trending on artstation'''

														_a			  =       np.random.RandomState(0      )
														_a			  =       pipe(
														    prompt=lowerCAmelCase_							,    image=lowerCAmelCase_							,    strength=0.7_5							,    guidance_scale=7.5							,    num_inference_steps=20							,    generator=lowerCAmelCase_							,    output_type='''np'''							,    )
														_a			  =       output.images
														_a			  =       images[0, 2_55:2_58, 3_83:3_86, -1]

														assert images.shape == (1, 5_12, 7_68, 3)
														_a			  =       np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]      )
														# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







'''simple docstring'''


import os

from bleurt import score  # From: git+https://github.com/google-research/bleurt.git

import datasets


_snake_case   :   Any     							=	datasets.logging.get_logger(__name__)


_snake_case   :   Tuple     							=	'\\n@inproceedings{bleurt,\n  title={BLEURT: Learning Robust Metrics for Text Generation},\n  author={Thibault Sellam and Dipanjan Das and Ankur P. Parikh},\n  booktitle={ACL},\n  year={2020},\n  url={https://arxiv.org/abs/2004.04696}\n}\n'

_snake_case   :   Optional[Any]     							=	'\\nBLEURT a learnt evaluation metric for Natural Language Generation. It is built using multiple phases of transfer learning starting from a pretrained BERT model (Devlin et al. 2018)\nand then employing another pre-training phrase using synthetic data. Finally it is trained on WMT human annotations. You may run BLEURT out-of-the-box or fine-tune\nit for your specific application (the latter is expected to perform better).\n\nSee the project\'s README at https://github.com/google-research/bleurt#readme for more information.\n'

_snake_case   :   Optional[int]     							=	'\nBLEURT score.\n\nArgs:\n    `predictions` (list of str): prediction/candidate sentences\n    `references` (list of str): reference sentences\n    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.\n\nReturns:\n    \'scores\': List of scores.\nExamples:\n\n    >>> predictions = ["hello there", "general kenobi"]\n    >>> references = ["hello there", "general kenobi"]\n    >>> bleurt = datasets.load_metric("bleurt")\n    >>> results = bleurt.compute(predictions=predictions, references=references)\n    >>> print([round(v, 2) for v in results["scores"]])\n    [1.03, 1.04]\n'

_snake_case   :   Dict     							=	{
    'bleurt-tiny-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-128.zip',
    'bleurt-tiny-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-512.zip',
    'bleurt-base-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-128.zip',
    'bleurt-base-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-512.zip',
    'bleurt-large-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-128.zip',
    'bleurt-large-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-512.zip',
    'BLEURT-20-D3': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D3.zip',
    'BLEURT-20-D6': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D6.zip',
    'BLEURT-20-D12': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D12.zip',
    'BLEURT-20': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip',
}

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







							def      __lowerCAmelCase    (						self :						int      )    ->							Dict:


														"""simple docstring"""


														return datasets.MetricInfo(
														    description=_DESCRIPTION							,    citation=_CITATION							,    homepage='''https://github.com/google-research/bleurt'''							,    inputs_description=_KWARGS_DESCRIPTION							,    features=datasets.Features(
														        {
														            '''predictions''': datasets.Value('''string'''							,    id='''sequence'''      ),
														            '''references''': datasets.Value('''string'''							,    id='''sequence'''      ),
														        }      )							,    codebase_urls=['''https://github.com/google-research/bleurt''']							,    reference_urls=['''https://github.com/google-research/bleurt''', '''https://arxiv.org/abs/2004.04696''']							,    )







							def      __lowerCAmelCase    (						self :						Dict							,    lowerCAmelCase_ :						str      )    ->							Tuple:


														"""simple docstring"""


														if self.config_name == "default":
																					logger.warning(
																					    '''Using default BLEURT-Base checkpoint for sequence maximum length 128. '''
																					    '''You can use a bigger model for better results with e.g.: datasets.load_metric(\'bleurt\', \'bleurt-large-512\').'''      )
																					_a			  =       '''bleurt-base-128'''

														if self.config_name.lower() in CHECKPOINT_URLS:
																					_a			  =       self.config_name.lower()

														elif self.config_name.upper() in CHECKPOINT_URLS:
																					_a			  =       self.config_name.upper()

														else:
																					raise KeyError(
																					    F'{self.config_name} model not found. You should supply the name of a model checkpoint for bleurt in {CHECKPOINT_URLS.keys()}'      )

														# download the model checkpoint specified by self.config_name and set up the scorer
														_a			  =       dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]      )
														_a			  =       score.BleurtScorer(os.path.join(lowerCAmelCase_							,    lowerCAmelCase_      )      )







							def      __lowerCAmelCase    (						self :						str							,    lowerCAmelCase_ :						int							,    lowerCAmelCase_ :						Union[str, Any]      )    ->							Optional[int]:


														"""simple docstring"""


														_a			  =       self.scorer.score(references=lowerCAmelCase_							,    candidates=lowerCAmelCase_      )
														return {"scores": scores}

