


import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    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 						_UpperCAmelCase      (     lowercase		,		unittest.TestCase							):
		# TODO: is there an appropriate internal test set?
		lowerCamelCase_	:			Tuple       			=  """ssube/stable-diffusion-x4-upscaler-onnx"""






		def _snake_case		(						self				:     Tuple   ,  UpperCAmelCase				:     List[str]=0):
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						floats_tensor((1, 3, 1_28, 1_28)   ,  rng=random.Random(UpperCAmelCase))
				SCREAMING_SNAKE_CASE_				:str				     =						torch.manual_seed(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Any				     =						{
				    "prompt": "A painting of a squirrel eating a burger",
				    "image": image,
				    "generator": generator,
				    "num_inference_steps": 3,
				    "guidance_scale": 7.5,
				    "output_type": "numpy",
				}
				return inputs






		def _snake_case		(						self				:     Optional[Any]):
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,  provider="CPUExecutionProvider")
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:Any				     =						self.get_dummy_inputs()
				SCREAMING_SNAKE_CASE_				:List[str]				     =						pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						image[0, -3:, -3:, -1].flatten()

				# started as 128, should now be 512
				assert image.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:int				     =						np.array(
				    [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223])
				assert np.abs(image_slice - expected_slice).max() < 1E-1






		def _snake_case		(						self				:     str):
				SCREAMING_SNAKE_CASE_				:int				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,  provider="CPUExecutionProvider")
				SCREAMING_SNAKE_CASE_				:str				     =						PNDMScheduler.from_config(pipe.scheduler.config   ,  skip_prk_steps=UpperCAmelCase)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

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

				assert image.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:str				     =						np.array(
				    [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964])
				assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-1






		def _snake_case		(						self				:     Tuple):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,  provider="CPUExecutionProvider")
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						self.get_dummy_inputs()
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:Dict				     =						image[0, -3:, -3:, -1]

				assert image.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						np.array(
				    [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515])

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






		def _snake_case		(						self				:     Tuple):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,  provider="CPUExecutionProvider")
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						EulerDiscreteScheduler.from_config(pipe.scheduler.config)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

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

				assert image.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						np.array(
				    [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223])
				assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-1







		def _snake_case		(						self				:     Dict):
				SCREAMING_SNAKE_CASE_				:str				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,  provider="CPUExecutionProvider")
				SCREAMING_SNAKE_CASE_				:Dict				     =						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						self.get_dummy_inputs()
				SCREAMING_SNAKE_CASE_				:Any				     =						pipe(**UpperCAmelCase).images
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						image[0, -3:, -3:, -1]

				assert image.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:str				     =						np.array(
				    [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043])

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




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






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






		@property
		def _snake_case		(						self				:     str):
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						ort.SessionOptions()
				SCREAMING_SNAKE_CASE_				:str				     =						False
				return options






		def _snake_case		(						self				:     Optional[int]):
				SCREAMING_SNAKE_CASE_				:int				     =						load_image(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
				    "/img2img/sketch-mountains-input.jpg")
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						init_image.resize((1_28, 1_28))
				# using the PNDM scheduler by default
				SCREAMING_SNAKE_CASE_				:int				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(
				    "ssube/stable-diffusion-x4-upscaler-onnx"   ,  provider=self.gpu_provider   ,  sess_options=self.gpu_options   ,  )
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						"A fantasy landscape, trending on artstation"

				SCREAMING_SNAKE_CASE_				:int				     =						torch.manual_seed(0)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						pipe(
				    prompt=UpperCAmelCase   ,  image=UpperCAmelCase   ,  guidance_scale=7.5   ,  num_inference_steps=10   ,  generator=UpperCAmelCase   ,  output_type="np"   ,  )
				SCREAMING_SNAKE_CASE_				:int				     =						output.images
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						images[0, 2_55:2_58, 3_83:3_86, -1]

				assert images.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:Any				     =						np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972])
				# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







		def _snake_case		(						self				:     Tuple):
				SCREAMING_SNAKE_CASE_				:Tuple				     =						load_image(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
				    "/img2img/sketch-mountains-input.jpg")
				SCREAMING_SNAKE_CASE_				:Dict				     =						init_image.resize((1_28, 1_28))
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						LMSDiscreteScheduler.from_pretrained(
				    "ssube/stable-diffusion-x4-upscaler-onnx"   ,  subfolder="scheduler")
				SCREAMING_SNAKE_CASE_				:Dict				     =						OnnxStableDiffusionUpscalePipeline.from_pretrained(
				    "ssube/stable-diffusion-x4-upscaler-onnx"   ,  scheduler=UpperCAmelCase   ,  provider=self.gpu_provider   ,  sess_options=self.gpu_options   ,  )
				pipe.set_progress_bar_config(disable=UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:List[str]				     =						"A fantasy landscape, trending on artstation"

				SCREAMING_SNAKE_CASE_				:str				     =						torch.manual_seed(0)
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						pipe(
				    prompt=UpperCAmelCase   ,  image=UpperCAmelCase   ,  guidance_scale=7.5   ,  num_inference_steps=20   ,  generator=UpperCAmelCase   ,  output_type="np"   ,  )
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						output.images
				SCREAMING_SNAKE_CASE_				:Dict				     =						images[0, 2_55:2_58, 3_83:3_86, -1]

				assert images.shape == (1, 5_12, 5_12, 3)
				SCREAMING_SNAKE_CASE_				:Tuple				     =						np.array(
				    [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566])
				# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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










import os

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

import datasets


SCREAMING_SNAKE_CASE__     					=			datasets.logging.get_logger(__name__)


SCREAMING_SNAKE_CASE__     					=			"\\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"

SCREAMING_SNAKE_CASE__     					=			"\\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"

SCREAMING_SNAKE_CASE__     					=			"\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"

SCREAMING_SNAKE_CASE__     					=			{
    "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 						_UpperCAmelCase      (     datasets.Metric							):






		def _snake_case		(						self				:     List[str]):
				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 _snake_case		(						self				:     Any   ,  UpperCAmelCase				:     Any):
				# check that config name specifies a valid BLEURT model
				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').")
						SCREAMING_SNAKE_CASE_				:str				     =						"bleurt-base-128"

				if self.config_name.lower() in CHECKPOINT_URLS:
						SCREAMING_SNAKE_CASE_				:Dict				     =						self.config_name.lower()

				elif self.config_name.upper() in CHECKPOINT_URLS:
						SCREAMING_SNAKE_CASE_				:Dict				     =						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
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name])
				SCREAMING_SNAKE_CASE_				:Dict				     =						score.BleurtScorer(os.path.join(UpperCAmelCase   ,  UpperCAmelCase))







		def _snake_case		(						self				:     Union[str, Any]   ,  UpperCAmelCase				:     List[str]   ,  UpperCAmelCase				:     Union[str, Any]):
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						self.scorer.score(references=UpperCAmelCase   ,  candidates=UpperCAmelCase)
				return {"scores": scores}






