"""simple docstring"""
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						( _snake_case ,			unittest.TestCase						):






    """simple docstring"""
    # TODO: is there an appropriate internal test set?
    __snake_case    =       'ssube/stable-diffusion-x4-upscaler-onnx'







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

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

        # started as 128, should now be 512
        assert image.shape == (1, 512, 512, 3)
        _lowerCamelCase	: Tuple       =       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 				a__					(			self					)	->	str:
        _lowerCamelCase	: Dict       =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    , provider='''CPUExecutionProvider'''					)
        _lowerCamelCase	: Optional[Any]       =       PNDMScheduler.from_config(pipe.scheduler.config    , skip_prk_steps=_lowercase					)
        pipe.set_progress_bar_config(disable=_lowercase					)

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

        assert image.shape == (1, 512, 512, 3)
        _lowerCamelCase	: Optional[Any]       =       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 				a__					(			self					)	->	Union[str, Any]:
        _lowerCamelCase	: str       =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    , provider='''CPUExecutionProvider'''					)
        _lowerCamelCase	: Optional[int]       =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config					)
        pipe.set_progress_bar_config(disable=_lowercase					)

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

        assert image.shape == (1, 512, 512, 3)
        _lowerCamelCase	: Dict       =       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 				a__					(			self					)	->	List[str]:
        _lowerCamelCase	: List[str]       =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    , provider='''CPUExecutionProvider'''					)
        _lowerCamelCase	: Dict       =       EulerDiscreteScheduler.from_config(pipe.scheduler.config					)
        pipe.set_progress_bar_config(disable=_lowercase					)

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

        assert image.shape == (1, 512, 512, 3)
        _lowerCamelCase	: Dict       =       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 				a__					(			self					)	->	List[str]:
        _lowerCamelCase	: Optional[Any]       =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    , provider='''CPUExecutionProvider'''					)
        _lowerCamelCase	: List[Any]       =       EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config					)
        pipe.set_progress_bar_config(disable=_lowercase					)

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

        assert image.shape == (1, 512, 512, 3)
        _lowerCamelCase	: 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						):






    """simple docstring"""







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







    @property
    def 				a__					(			self					)	->	Optional[int]:
        _lowerCamelCase	: Optional[int]       =       ort.SessionOptions()
        _lowerCamelCase	: Optional[Any]       =       False
        return options







    def 				a__					(			self					)	->	Any:
        _lowerCamelCase	: str       =       load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/img2img/sketch-mountains-input.jpg'''					)
        _lowerCamelCase	: Dict       =       init_image.resize((128, 128)					)
        # using the PNDM scheduler by default
        _lowerCamelCase	: Optional[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=_lowercase					)

        _lowerCamelCase	: str       =       '''A fantasy landscape, trending on artstation'''

        _lowerCamelCase	: Optional[Any]       =       torch.manual_seed(0					)
        _lowerCamelCase	: List[Any]       =       pipe(
            prompt=_lowercase    , image=_lowercase    , guidance_scale=7.5    , num_inference_steps=10    , generator=_lowercase    , output_type='''np'''    , )
        _lowerCamelCase	: str       =       output.images
        _lowerCamelCase	: str       =       images[0, 255:258, 383:386, -1]

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

        _lowerCamelCase	: str       =       '''A fantasy landscape, trending on artstation'''

        _lowerCamelCase	: Dict       =       torch.manual_seed(0					)
        _lowerCamelCase	: Union[str, Any]       =       pipe(
            prompt=_lowercase    , image=_lowercase    , guidance_scale=7.5    , num_inference_steps=20    , generator=_lowercase    , output_type='''np'''    , )
        _lowerCamelCase	: Optional[int]       =       output.images
        _lowerCamelCase	: Optional[int]       =       images[0, 255:258, 383:386, -1]

        assert images.shape == (1, 512, 512, 3)
        _lowerCamelCase	: 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









'''simple docstring'''

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 rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
    import PIL


a   	=    logging.get_logger(__name__)

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       str:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							b.T
   __SCREAMING_SNAKE_CASE   				=							np.sum(np.square(__UpperCAmelCase   )				,						axis=1   )
   __SCREAMING_SNAKE_CASE   				=							np.sum(np.square(__UpperCAmelCase   )				,						axis=0   )
   __SCREAMING_SNAKE_CASE   				=							np.matmul(__UpperCAmelCase				,						__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							aa[:, None] - 2 * ab + ba[None, :]
   return d

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       List[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							x.reshape(-1				,						3   )
   __SCREAMING_SNAKE_CASE   				=							squared_euclidean_distance(__UpperCAmelCase				,						__UpperCAmelCase   )
   return np.argmin(__UpperCAmelCase				,						axis=1   )







class __a			( _snake_case      ):
   __UpperCamelCase :			Any						   =						['pixel_values']





   def __init__(       self					:      Any					,lowerCamelCase					:      Optional[Union[List[List[int]], np.ndarray]] = None					,lowerCamelCase					:      bool = True					,lowerCamelCase					:      Dict[str, int] = None					,lowerCamelCase					:      PILImageResampling = PILImageResampling.BILINEAR					,lowerCamelCase					:      bool = True					,lowerCamelCase					:      bool = True					,**lowerCamelCase					:      Optional[Any]					,):

      '''simple docstring'''

      super().__init__(**lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							size if size is not None else {"""height""": 256, """width""": 256}
      __SCREAMING_SNAKE_CASE   				=							get_size_dict(lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							np.array(lowerCamelCase     ) if clusters is not None else None
      __SCREAMING_SNAKE_CASE   				=							do_resize
      __SCREAMING_SNAKE_CASE   				=							size
      __SCREAMING_SNAKE_CASE   				=							resample
      __SCREAMING_SNAKE_CASE   				=							do_normalize
      __SCREAMING_SNAKE_CASE   				=							do_color_quantize





   def   UpperCAmelCase__  (       self					:      Optional[int]					,lowerCamelCase					:      np.ndarray					,lowerCamelCase					:      Dict[str, int]					,lowerCamelCase					:      PILImageResampling = PILImageResampling.BILINEAR					,lowerCamelCase					:      Optional[Union[str, ChannelDimension]] = None					,**lowerCamelCase					:      int					,):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							get_size_dict(lowerCamelCase     )
      if "height" not in size or "width" not in size:
         raise ValueError(f"""Size dictionary must contain both height and width keys. Got {size.keys()}"""     )
      return resize(
          lowerCamelCase					,size=(size["""height"""], size["""width"""])					,resample=lowerCamelCase					,data_format=lowerCamelCase					,**lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Optional[Any]					,lowerCamelCase					:      np.ndarray					,lowerCamelCase					:      Optional[Union[str, ChannelDimension]] = None					,):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							rescale(image=lowerCamelCase					,scale=1 / 127.5					,data_format=lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							image - 1
      return image





   def   UpperCAmelCase__  (       self					:      Dict					,lowerCamelCase					:      ImageInput					,lowerCamelCase					:      bool = None					,lowerCamelCase					:      Dict[str, int] = None					,lowerCamelCase					:      PILImageResampling = None					,lowerCamelCase					:      bool = None					,lowerCamelCase					:      Optional[bool] = None					,lowerCamelCase					:      Optional[Union[List[List[int]], np.ndarray]] = None					,lowerCamelCase					:      Optional[Union[str, TensorType]] = None					,lowerCamelCase					:      Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST					,**lowerCamelCase					:      Any					,):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							do_resize if do_resize is not None else self.do_resize
      __SCREAMING_SNAKE_CASE   				=							size if size is not None else self.size
      __SCREAMING_SNAKE_CASE   				=							get_size_dict(lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							resample if resample is not None else self.resample
      __SCREAMING_SNAKE_CASE   				=							do_normalize if do_normalize is not None else self.do_normalize
      __SCREAMING_SNAKE_CASE   				=							do_color_quantize if do_color_quantize is not None else self.do_color_quantize
      __SCREAMING_SNAKE_CASE   				=							clusters if clusters is not None else self.clusters
      __SCREAMING_SNAKE_CASE   				=							np.array(lowerCamelCase     )

      __SCREAMING_SNAKE_CASE   				=							make_list_of_images(lowerCamelCase     )

      if not valid_images(lowerCamelCase     ):
         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_color_quantize and clusters is None:
         raise ValueError("""Clusters must be specified if do_color_quantize is True."""     )

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

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

      if do_normalize:
         __SCREAMING_SNAKE_CASE   				=							[self.normalize(image=lowerCamelCase     ) for image in images]

      if do_color_quantize:
         __SCREAMING_SNAKE_CASE   				=							[to_channel_dimension_format(lowerCamelCase					,ChannelDimension.LAST     ) for image in images]
         # color quantize from (batch_size, height, width, 3) to (batch_size, height, width)
         __SCREAMING_SNAKE_CASE   				=							np.array(lowerCamelCase     )
         __SCREAMING_SNAKE_CASE   				=							color_quantize(lowerCamelCase					,lowerCamelCase     ).reshape(images.shape[:-1]     )

         # flatten to (batch_size, height*width)
         __SCREAMING_SNAKE_CASE   				=							images.shape[0]
         __SCREAMING_SNAKE_CASE   				=							images.reshape(lowerCamelCase					,-1     )

         # We need to convert back to a list of images to keep consistent behaviour across processors.
         __SCREAMING_SNAKE_CASE   				=							list(lowerCamelCase     )
      else:
         __SCREAMING_SNAKE_CASE   				=							[to_channel_dimension_format(lowerCamelCase					,lowerCamelCase     ) for image in images]

      __SCREAMING_SNAKE_CASE   				=							{"""input_ids""": images}
      return BatchFeature(data=lowerCamelCase					,tensor_type=lowerCamelCase     )





