

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   ):
							"""simple docstring"""







							UpperCAmelCase__       :						List[str]	    =    "ssube/stable-diffusion-x4-upscaler-onnx"


							def 				__lowercase				(			self       ,							_a=0       )    ->    Any:
												_a	:    Optional[int]         =		floats_tensor((1, 3, 1_2_8, 1_2_8)       ,							rng=random.Random(_a       )       )
												_a	:    Tuple         =		torch.manual_seed(_a       )
												_a	:    str         =		{
												    '''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 				__lowercase				(			self       )    ->    Optional[int]:
												_a	:    Union[str, Any]         =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,							provider='''CPUExecutionProvider'''       )
												pipe.set_progress_bar_config(disable=_a       )

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

												# started as 128, should now be 512
												assert image.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    List[Any]         =		np.array(
												    [0.697_4782, 0.6890_2093, 0.7013_5885, 0.758_3618, 0.780_4545, 0.785_4912, 0.7866_7426, 0.7874_3863, 0.7807_0223]       )
												assert np.abs(image_slice - expected_slice       ).max() < 1e-1


							def 				__lowercase				(			self       )    ->    int:
												_a	:    Optional[Any]         =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,							provider='''CPUExecutionProvider'''       )
												_a	:    Tuple         =		PNDMScheduler.from_config(pipe.scheduler.config       ,							skip_prk_steps=_a       )
												pipe.set_progress_bar_config(disable=_a       )

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

												assert image.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    List[Any]         =		np.array(
												    [0.689_8892, 0.5924_0556, 0.5249_9527, 0.5886_6215, 0.5225_8235, 0.5257_2715, 0.6241_4473, 0.617_4387, 0.621_4964]       )
												assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-1


							def 				__lowercase				(			self       )    ->    Optional[int]:
												_a	:    Dict         =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,							provider='''CPUExecutionProvider'''       )
												_a	:    Dict         =		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
												pipe.set_progress_bar_config(disable=_a       )

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

												assert image.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    Optional[int]         =		np.array(
												    [0.765_9278, 0.7643_7664, 0.7557_9107, 0.769_1116, 0.7766_6986, 0.772_7672, 0.775_8664, 0.781_2226, 0.7694_2515]       )

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


							def 				__lowercase				(			self       )    ->    List[str]:
												_a	:    Tuple         =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,							provider='''CPUExecutionProvider'''       )
												_a	:    Tuple         =		EulerDiscreteScheduler.from_config(pipe.scheduler.config       )
												pipe.set_progress_bar_config(disable=_a       )

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

												assert image.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    Tuple         =		np.array(
												    [0.697_4782, 0.6890_2093, 0.7013_5885, 0.758_3618, 0.780_4545, 0.785_4912, 0.7866_7426, 0.7874_3863, 0.7807_0223]       )
												assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-1




							def 				__lowercase				(			self       )    ->    List[str]:
												_a	:    Optional[Any]         =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,							provider='''CPUExecutionProvider'''       )
												_a	:    Optional[int]         =		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config       )
												pipe.set_progress_bar_config(disable=_a       )

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

												assert image.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    Optional[int]         =		np.array(
												    [0.7742_4496, 0.77_3601, 0.764_5288, 0.776_9598, 0.777_2739, 0.773_8688, 0.7818_7233, 0.7787_9584, 0.76_7043]       )

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



@nightly
@require_onnxruntime
@require_torch_gpu
class 				UpperCAmelCase_		(			unittest.TestCase   ):
							"""simple docstring"""







							@property
							def 				__lowercase				(			self       )    ->    Optional[int]:
												return (
												    "CUDAExecutionProvider",
												    {
												        "gpu_mem_limit": "15000000000",  # 15GB
												        "arena_extend_strategy": "kSameAsRequested",
												    },
												)


							@property
							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												_a	:    Tuple         =		ort.SessionOptions()
												_a	:    List[str]         =		False
												return options


							def 				__lowercase				(			self       )    ->    Optional[int]:
												_a	:    Any         =		load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg'''       )
												_a	:    Any         =		init_image.resize((1_2_8, 1_2_8)       )
												# using the PNDM scheduler by default
												_a	:    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=_a       )

												_a	:    List[str]         =		'''A fantasy landscape, trending on artstation'''

												_a	:    Tuple         =		torch.manual_seed(0       )
												_a	:    Tuple         =		pipe(
												    prompt=_a       ,							image=_a       ,							guidance_scale=7.5       ,							num_inference_steps=1_0       ,							generator=_a       ,							output_type='''np'''       ,							)
												_a	:    Optional[Any]         =		output.images
												_a	:    Optional[Any]         =		images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

												assert images.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    List[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 				__lowercase				(			self       )    ->    Tuple:
												_a	:    Dict         =		load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg'''       )
												_a	:    Tuple         =		init_image.resize((1_2_8, 1_2_8)       )
												_a	:    Dict         =		LMSDiscreteScheduler.from_pretrained(
												    '''ssube/stable-diffusion-x4-upscaler-onnx'''       ,							subfolder='''scheduler'''       )
												_a	:    Optional[Any]         =		OnnxStableDiffusionUpscalePipeline.from_pretrained(
												    '''ssube/stable-diffusion-x4-upscaler-onnx'''       ,							scheduler=_a       ,							provider=self.gpu_provider       ,							sess_options=self.gpu_options       ,							)
												pipe.set_progress_bar_config(disable=_a       )

												_a	:    List[Any]         =		'''A fantasy landscape, trending on artstation'''

												_a	:    str         =		torch.manual_seed(0       )
												_a	:    Union[str, Any]         =		pipe(
												    prompt=_a       ,							image=_a       ,							guidance_scale=7.5       ,							num_inference_steps=2_0       ,							generator=_a       ,							output_type='''np'''       ,							)
												_a	:    Dict         =		output.images
												_a	:    Optional[int]         =		images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

												assert images.shape == (1, 5_1_2, 5_1_2, 3)
												_a	:    Optional[int]         =		np.array(
												    [0.5017_3753, 0.5022_3356, 0.50_2039, 0.5023_3036, 0.502_3725, 0.502_2601, 0.501_8758, 0.5023_4085, 0.5024_1566]       )
												# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
			import torch

if is_vision_available():
			import PIL


a__				=		logging.get_logger(__name__)



def    __UpperCAmelCase (	__a      :			np.ndarray		,__a      :			Union[int, Iterable[int]]		,__a      :			bool		,__a      :			int	)			->     Tuple[int, int]:

					"""simple docstring"""

					def constraint_to_multiple_of(__a      :			List[str]		,__a      :			Dict		,__a      :			Any=0		,__a      :			int=None	):
										_a	:    Dict         =		round(val / multiple	) * multiple

										if max_val is not None and x > max_val:
															_a	:    Any         =		math.floor(val / multiple	) * multiple

										if x < min_val:
															_a	:    Dict         =		math.ceil(val / multiple	) * multiple

										return x

					_a	:    Union[str, Any]         =		(output_size, output_size) if isinstance(__a		,__a	) else output_size

					_a      ,					_a	:    List[Any]         =		get_image_size(__a	)
					_a      ,					_a	:    Any         =		output_size

					# determine new height and width
					_a	:    Union[str, Any]         =		output_height / input_height
					_a	:    Tuple         =		output_width / input_width

					if keep_aspect_ratio:
										# scale as little as possible
										if abs(1 - scale_width	) < abs(1 - scale_height	):
															# fit width
															_a	:    Optional[Any]         =		scale_width
										else:
															# fit height
															_a	:    Tuple         =		scale_height

					_a	:    Optional[Any]         =		constraint_to_multiple_of(scale_height * input_height		,multiple=__a	)
					_a	:    int         =		constraint_to_multiple_of(scale_width * input_width		,multiple=__a	)

					return (new_height, new_width)







class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Optional[Any]	    =    ["pixel_values"]


							def __init__(			self       ,							_a = True       ,							_a = None       ,							_a = PILImageResampling.BILINEAR       ,							_a = False       ,							_a = 1       ,							_a = True       ,							_a = 1 / 2_5_5       ,							_a = True       ,							_a = None       ,							_a = None       ,							**_a       ,							)    ->    None:
												super().__init__(**_a       )
												_a	:    Optional[int]         =		size if size is not None else {'''height''': 3_8_4, '''width''': 3_8_4}
												_a	:    Optional[Any]         =		get_size_dict(_a       )
												_a	:    Any         =		do_resize
												_a	:    Dict         =		size
												_a	:    str         =		keep_aspect_ratio
												_a	:    Any         =		ensure_multiple_of
												_a	:    Optional[Any]         =		resample
												_a	:    List[Any]         =		do_rescale
												_a	:    int         =		rescale_factor
												_a	:    Any         =		do_normalize
												_a	:    Union[str, Any]         =		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
												_a	:    List[str]         =		image_std if image_std is not None else IMAGENET_STANDARD_STD


							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a = False       ,							_a = 1       ,							_a = PILImageResampling.BICUBIC       ,							_a = None       ,							**_a       ,							)    ->    np.ndarray:
												_a	:    str         =		get_size_dict(_a       )
												if "height" not in size or "width" not in size:
																	raise ValueError(F"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"""       )
												_a	:    Optional[Any]         =		get_resize_output_image_size(
												    _a       ,							output_size=(size['''height'''], size['''width'''])       ,							keep_aspect_ratio=_a       ,							multiple=_a       ,							)
												return resize(_a       ,							size=_a       ,							resample=_a       ,							data_format=_a       ,							**_a       )


							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a = None       ,							**_a       ,							)    ->    int:
												return rescale(_a       ,							scale=_a       ,							data_format=_a       ,							**_a       )


							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a       ,							_a = None       ,							**_a       ,							)    ->    np.ndarray:
												return normalize(_a       ,							mean=_a       ,							std=_a       ,							data_format=_a       ,							**_a       )


							def 				__lowercase				(			self       ,							_a       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = ChannelDimension.FIRST       ,							**_a       ,							)    ->    PIL.Image.Image:
												_a	:    Optional[int]         =		do_resize if do_resize is not None else self.do_resize
												_a	:    Union[str, Any]         =		size if size is not None else self.size
												_a	:    str         =		get_size_dict(_a       )
												_a	:    List[Any]         =		keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
												_a	:    Optional[Any]         =		ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
												_a	:    str         =		resample if resample is not None else self.resample
												_a	:    str         =		do_rescale if do_rescale is not None else self.do_rescale
												_a	:    Union[str, Any]         =		rescale_factor if rescale_factor is not None else self.rescale_factor
												_a	:    List[Any]         =		do_normalize if do_normalize is not None else self.do_normalize
												_a	:    str         =		image_mean if image_mean is not None else self.image_mean
												_a	:    Tuple         =		image_std if image_std is not None else self.image_std

												_a	:    Dict         =		make_list_of_images(_a       )

												if not valid_images(_a       ):
																	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 or resample is None:
																	raise ValueError('''Size and resample must be specified if do_resize 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.'''       )

												# All transformations expect numpy arrays.
												_a	:    Dict         =		[to_numpy_array(_a       ) for image in images]

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

												if do_rescale:
																	_a	:    Optional[Any]         =		[self.rescale(image=_a       ,							scale=_a       ) for image in images]

												if do_normalize:
																	_a	:    Optional[int]         =		[self.normalize(image=_a       ,							mean=_a       ,							std=_a       ) for image in images]

												_a	:    int         =		[to_channel_dimension_format(_a       ,							_a       ) for image in images]

												_a	:    Tuple         =		{'''pixel_values''': images}
												return BatchFeature(data=_a       ,							tensor_type=_a       )




							def 				__lowercase				(			self       ,							_a       ,							_a = None       )    ->    Any:
												_a	:    Optional[Any]         =		outputs.logits

												# Resize logits and compute semantic segmentation maps
												if target_sizes is not None:
																	if len(_a       ) != len(_a       ):
																						raise ValueError(
																						    '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''       )

																	if is_torch_tensor(_a       ):
																						_a	:    List[Any]         =		target_sizes.numpy()

																	_a	:    str         =		[]

																	for idx in range(len(_a       )       ):
																						_a	:    str         =		torch.nn.functional.interpolate(
																						    logits[idx].unsqueeze(dim=0       )       ,							size=target_sizes[idx]       ,							mode='''bilinear'''       ,							align_corners=_a       )
																						_a	:    Union[str, Any]         =		resized_logits[0].argmax(dim=0       )
																						semantic_segmentation.append(_a       )
												else:
																	_a	:    Tuple         =		logits.argmax(dim=1       )
																	_a	:    Union[str, Any]         =		[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]       )]

												return semantic_segmentation



