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     _snake_case  (  __lowerCamelCase    ,					unittest.TestCase ):
					lowerCamelCase__:						Dict     =   """ssube/stable-diffusion-x4-upscaler-onnx"""
					def       _lowerCamelCase       (     self: Any		,     __lowerCamelCase: Optional[Any]=0   )			-> Optional[int]:
												__UpperCAmelCase  :	Optional[int]   			=      floats_tensor((1, 3, 1_28, 1_28)		,     rng=random.Random(__lowerCamelCase   )   )
												__UpperCAmelCase  :	Union[str, Any]   			=      torch.manual_seed(__lowerCamelCase   )
												__UpperCAmelCase  :	int   			=      {
												    "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       _lowerCamelCase       (     self: List[Any]   )			-> Tuple:
												__UpperCAmelCase  :	List[Any]   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider="CPUExecutionProvider"   )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

												__UpperCAmelCase  :	Optional[int]   			=      self.get_dummy_inputs()
												__UpperCAmelCase  :	int   			=      pipe(**__lowerCamelCase   ).images
												__UpperCAmelCase  :	List[str]   			=      image[0, -3:, -3:, -1].flatten()

												# started as 128, should now be 512
												assert image.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	Dict   			=      np.array(
												    [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]   )
												assert np.abs(image_slice - expected_slice   ).max() < 1e-1
					def       _lowerCamelCase       (     self: int   )			-> Optional[Any]:
												__UpperCAmelCase  :	str   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider="CPUExecutionProvider"   )
												__UpperCAmelCase  :	List[str]   			=      PNDMScheduler.from_config(pipe.scheduler.config		,     skip_prk_steps=__lowerCamelCase   )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

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

												assert image.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	List[str]   			=      np.array(
												    [0.6_89_88_92, 0.59_24_05_56, 0.52_49_95_27, 0.58_86_62_15, 0.52_25_82_35, 0.52_57_27_15, 0.62_41_44_73, 0.6_17_43_87, 0.6_21_49_64]   )
												assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-1
					def       _lowerCamelCase       (     self: Tuple   )			-> Union[str, Any]:
												__UpperCAmelCase  :	List[str]   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider="CPUExecutionProvider"   )
												__UpperCAmelCase  :	Optional[int]   			=      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config   )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

												__UpperCAmelCase  :	Tuple   			=      self.get_dummy_inputs()
												__UpperCAmelCase  :	Union[str, Any]   			=      pipe(**__lowerCamelCase   ).images
												__UpperCAmelCase  :	Any   			=      image[0, -3:, -3:, -1]

												assert image.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	List[str]   			=      np.array(
												    [0.7_65_92_78, 0.76_43_76_64, 0.75_57_91_07, 0.7_69_11_16, 0.77_66_69_86, 0.7_72_76_72, 0.7_75_86_64, 0.7_81_22_26, 0.76_94_25_15]   )

												assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-1
					def       _lowerCamelCase       (     self: List[str]   )			-> List[str]:
												__UpperCAmelCase  :	List[str]   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider="CPUExecutionProvider"   )
												__UpperCAmelCase  :	List[Any]   			=      EulerDiscreteScheduler.from_config(pipe.scheduler.config   )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

												__UpperCAmelCase  :	Any   			=      self.get_dummy_inputs()
												__UpperCAmelCase  :	Any   			=      pipe(**__lowerCamelCase   ).images
												__UpperCAmelCase  :	int   			=      image[0, -3:, -3:, -1]

												assert image.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	Optional[Any]   			=      np.array(
												    [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]   )
												assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-1





					def       _lowerCamelCase       (     self: List[Any]   )			-> List[Any]:
												__UpperCAmelCase  :	Any   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider="CPUExecutionProvider"   )
												__UpperCAmelCase  :	Optional[Any]   			=      EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config   )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

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

												assert image.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	List[str]   			=      np.array(
												    [0.77_42_44_96, 0.77_36_01, 0.7_64_52_88, 0.7_76_95_98, 0.7_77_27_39, 0.7_73_86_88, 0.78_18_72_33, 0.77_87_95_84, 0.76_70_43]   )

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




@nightly
@require_onnxruntime
@require_torch_gpu
class     _snake_case  (  unittest.TestCase ):
					@property
					def       _lowerCamelCase       (     self: str   )			-> Optional[Any]:
												return (
												    "CUDAExecutionProvider",
												    {
												        "gpu_mem_limit": "15000000000",  # 15GB
												        "arena_extend_strategy": "kSameAsRequested",
												    },
												)
					@property
					def       _lowerCamelCase       (     self: List[str]   )			-> str:
												__UpperCAmelCase  :	Optional[int]   			=      ort.SessionOptions()
												__UpperCAmelCase  :	str   			=      False
												return options
					def       _lowerCamelCase       (     self: List[str]   )			-> Tuple:
												__UpperCAmelCase  :	Union[str, Any]   			=      load_image(
												    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
												    "/img2img/sketch-mountains-input.jpg"   )
												__UpperCAmelCase  :	Optional[int]   			=      init_image.resize((1_28, 1_28)   )
												# using the PNDM scheduler by default
												__UpperCAmelCase  :	Union[str, Any]   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(
												    "ssube/stable-diffusion-x4-upscaler-onnx"		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

												__UpperCAmelCase  :	Optional[int]   			=      "A fantasy landscape, trending on artstation"

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

												assert images.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	Optional[int]   			=      np.array([0.48_83, 0.49_47, 0.49_80, 0.49_75, 0.49_82, 0.49_80, 0.50_00, 0.50_06, 0.49_72]   )
												# 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[Any]   )			-> str:
												__UpperCAmelCase  :	List[Any]   			=      load_image(
												    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
												    "/img2img/sketch-mountains-input.jpg"   )
												__UpperCAmelCase  :	Dict   			=      init_image.resize((1_28, 1_28)   )
												__UpperCAmelCase  :	List[str]   			=      LMSDiscreteScheduler.from_pretrained(
												    "ssube/stable-diffusion-x4-upscaler-onnx"		,     subfolder="scheduler"   )
												__UpperCAmelCase  :	List[Any]   			=      OnnxStableDiffusionUpscalePipeline.from_pretrained(
												    "ssube/stable-diffusion-x4-upscaler-onnx"		,     scheduler=__lowerCamelCase		,     provider=self.gpu_provider		,     sess_options=self.gpu_options		,     )
												pipe.set_progress_bar_config(disable=__lowerCamelCase   )

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

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

												assert images.shape == (1, 5_12, 5_12, 3)
												__UpperCAmelCase  :	Union[str, Any]   			=      np.array(
												    [0.50_17_37_53, 0.50_22_33_56, 0.50_20_39, 0.50_23_30_36, 0.5_02_37_25, 0.5_02_26_01, 0.5_01_87_58, 0.50_23_40_85, 0.50_24_15_66]   )
												# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


UpperCAmelCase	      = logging.get_logger(__name__)


if is_vision_available():
	import PIL







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     str							 =       ["""pixel_values"""]





	def __init__(       self		,    snake_case = True		,    snake_case = None		,    snake_case = PILImageResampling.BICUBIC		,    snake_case = True		,    snake_case = None		,    snake_case = True		,    snake_case = 1 / 255		,    snake_case = True		,    snake_case = None		,    snake_case = None		,    snake_case = True		,    **snake_case		,    ):
		super().__init__(**snake_case       )
		lowercase       =    size if size is not None else {'shortest_edge': 224}
		lowercase       =    get_size_dict(snake_case		,    default_to_square=snake_case       )
		lowercase       =    crop_size if crop_size is not None else {'height': 224, 'width': 224}
		lowercase       =    get_size_dict(snake_case		,    default_to_square=snake_case		,    param_name='crop_size'       )

		lowercase       =    do_resize
		lowercase       =    size
		lowercase       =    resample
		lowercase       =    do_center_crop
		lowercase       =    crop_size
		lowercase       =    do_rescale
		lowercase       =    rescale_factor
		lowercase       =    do_normalize
		lowercase       =    image_mean if image_mean is not None else OPENAI_CLIP_MEAN
		lowercase       =    image_std if image_std is not None else OPENAI_CLIP_STD
		lowercase       =    do_convert_rgb





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case = PILImageResampling.BICUBIC		,    snake_case = None		,    **snake_case		,    ):
		lowercase       =    get_size_dict(snake_case		,    default_to_square=snake_case       )
		if "shortest_edge" not in size:
			raise ValueError(F'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''       )
		lowercase       =    get_resize_output_image_size(snake_case		,    size=size['shortest_edge']		,    default_to_square=snake_case       )
		return resize(snake_case		,    size=snake_case		,    resample=snake_case		,    data_format=snake_case		,    **snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case = None		,    **snake_case		,    ):
		lowercase       =    get_size_dict(snake_case       )
		if "height" not in size or "width" not in size:
			raise ValueError(F'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''       )
		return center_crop(snake_case		,    size=(size['height'], size['width'])		,    data_format=snake_case		,    **snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case = None		,    **snake_case		,    ):
		return rescale(snake_case		,    scale=snake_case		,    data_format=snake_case		,    **snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case = None		,    **snake_case		,    ):
		return normalize(snake_case		,    mean=snake_case		,    std=snake_case		,    data_format=snake_case		,    **snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = ChannelDimension.FIRST		,    **snake_case		,    ):
		lowercase       =    do_resize if do_resize is not None else self.do_resize
		lowercase       =    size if size is not None else self.size
		lowercase       =    get_size_dict(snake_case		,    param_name='size'		,    default_to_square=snake_case       )
		lowercase       =    resample if resample is not None else self.resample
		lowercase       =    do_center_crop if do_center_crop is not None else self.do_center_crop
		lowercase       =    crop_size if crop_size is not None else self.crop_size
		lowercase       =    get_size_dict(snake_case		,    param_name='crop_size'		,    default_to_square=snake_case       )
		lowercase       =    do_rescale if do_rescale is not None else self.do_rescale
		lowercase       =    rescale_factor if rescale_factor is not None else self.rescale_factor
		lowercase       =    do_normalize if do_normalize is not None else self.do_normalize
		lowercase       =    image_mean if image_mean is not None else self.image_mean
		lowercase       =    image_std if image_std is not None else self.image_std
		lowercase       =    do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

		lowercase       =    make_list_of_images(snake_case       )

		if not valid_images(snake_case       ):
			raise ValueError(
			    'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
			    'torch.Tensor, tf.Tensor or jax.ndarray.'       )

		if do_resize and size is None:
			raise ValueError('Size must be specified if do_resize is True.'       )

		if do_center_crop and crop_size is None:
			raise ValueError('Crop size must be specified if do_center_crop is True.'       )

		if do_rescale and rescale_factor is None:
			raise ValueError('Rescale factor must be specified if do_rescale is True.'       )

		if do_normalize and (image_mean is None or image_std is None):
			raise ValueError('Image mean and std must be specified if do_normalize is True.'       )

		# PIL RGBA images are converted to RGB
		if do_convert_rgb:
			lowercase       =    [convert_to_rgb(snake_case       ) for image in images]

		# All transformations expect numpy arrays.
		lowercase       =    [to_numpy_array(snake_case       ) for image in images]

		if do_resize:
			lowercase       =    [self.resize(image=snake_case		,    size=snake_case		,    resample=snake_case       ) for image in images]

		if do_center_crop:
			lowercase       =    [self.center_crop(image=snake_case		,    size=snake_case       ) for image in images]

		if do_rescale:
			lowercase       =    [self.rescale(image=snake_case		,    scale=snake_case       ) for image in images]

		if do_normalize:
			lowercase       =    [self.normalize(image=snake_case		,    mean=snake_case		,    std=snake_case       ) for image in images]

		lowercase       =    [to_channel_dimension_format(snake_case		,    snake_case       ) for image in images]

		lowercase       =    {'pixel_values': images}
		return BatchFeature(data=snake_case		,    tensor_type=snake_case       )

