

'''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 							A__   (	A__					,       unittest.TestCase  ):
		# TODO: is there an appropriate internal test set?
		A__					=    'ssube/stable-diffusion-x4-upscaler-onnx'





		def 						A  (		self				:     str							, _a				:     Any=0			)     ->   Tuple:





							'''simple docstring'''






							_SCREAMING_SNAKE_CASE							   =floats_tensor((1, 3, 128, 128)							, rng=random.Random(_a			)			)
							_SCREAMING_SNAKE_CASE							   =torch.manual_seed(_a			)
							_SCREAMING_SNAKE_CASE							   ={
							    '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				:     int			)     ->   Tuple:





							'''simple docstring'''






							_SCREAMING_SNAKE_CASE							   =OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							, provider='CPUExecutionProvider'			)
							pipe.set_progress_bar_config(disable=_a			)

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

							# started as 128, should now be 512
							assert image.shape == (1, 512, 512, 3)
							_SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     str			)     ->   Any:





							'''simple docstring'''






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

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

							assert image.shape == (1, 512, 512, 3)
							_SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     Union[str, Any]			)     ->   Union[str, Any]:





							'''simple docstring'''






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

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

							assert image.shape == (1, 512, 512, 3)
							_SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     Optional[int]			)     ->   List[Any]:





							'''simple docstring'''






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

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

							assert image.shape == (1, 512, 512, 3)
							_SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     List[str]			)     ->   List[str]:





							'''simple docstring'''






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

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

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





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





							'''simple docstring'''






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





		@property
		def 						A  (		self				:     Tuple			)     ->   Any:





							'''simple docstring'''






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





		def 						A  (		self				:     List[Any]			)     ->   Tuple:





							'''simple docstring'''






							_SCREAMING_SNAKE_CASE							   =load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'			)
							_SCREAMING_SNAKE_CASE							   =init_image.resize((128, 128)			)
							# using the PNDM scheduler by default
							_SCREAMING_SNAKE_CASE							   =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			)

							_SCREAMING_SNAKE_CASE							   ='A fantasy landscape, trending on artstation'

							_SCREAMING_SNAKE_CASE							   =torch.manual_seed(0			)
							_SCREAMING_SNAKE_CASE							   =pipe(
							    prompt=_a							, image=_a							, guidance_scale=7.5							, num_inference_steps=10							, generator=_a							, output_type='np'							, )
							_SCREAMING_SNAKE_CASE							   =output.images
							_SCREAMING_SNAKE_CASE							   =images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 512, 3)
							_SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     Any			)     ->   Tuple:





							'''simple docstring'''






							_SCREAMING_SNAKE_CASE							   =load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'			)
							_SCREAMING_SNAKE_CASE							   =init_image.resize((128, 128)			)
							_SCREAMING_SNAKE_CASE							   =LMSDiscreteScheduler.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'							, subfolder='scheduler'			)
							_SCREAMING_SNAKE_CASE							   =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			)

							_SCREAMING_SNAKE_CASE							   ='A fantasy landscape, trending on artstation'

							_SCREAMING_SNAKE_CASE							   =torch.manual_seed(0			)
							_SCREAMING_SNAKE_CASE							   =pipe(
							    prompt=_a							, image=_a							, guidance_scale=7.5							, num_inference_steps=20							, generator=_a							, output_type='np'							, )
							_SCREAMING_SNAKE_CASE							   =output.images
							_SCREAMING_SNAKE_CASE							   =images[0, 255:258, 383:386, -1]

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



'''simple docstring'''
from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching


def 	_lowerCAmelCase (						)				->      int:
     """simple docstring"""




     import os as original_os
     from os import path as original_path
     from os import rename as original_rename
     from os.path import dirname as original_dirname
     from os.path import join as original_join

     assert _test_patching.os is original_os
     assert _test_patching.path is original_path
     assert _test_patching.join is original_join

     assert _test_patching.renamed_os is original_os
     assert _test_patching.renamed_path is original_path
     assert _test_patching.renamed_join is original_join

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_mock__'
     with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
          # Every way to access os.path.join must be patched, and the rest must stay untouched

          # check os.path.join
          assert isinstance(_test_patching.os							,			_PatchedModuleObj    )
          assert isinstance(_test_patching.os.path							,			_PatchedModuleObj    )
          assert _test_patching.os.path.join is mock

          # check path.join
          assert isinstance(_test_patching.path							,			_PatchedModuleObj    )
          assert _test_patching.path.join is mock

          # check join
          assert _test_patching.join is mock

          # check that the other attributes are untouched
          assert _test_patching.os.rename is original_rename
          assert _test_patching.path.dirname is original_dirname
          assert _test_patching.os.path.dirname is original_dirname

          # Even renamed modules or objects must be patched

          # check renamed_os.path.join
          assert isinstance(_test_patching.renamed_os							,			_PatchedModuleObj    )
          assert isinstance(_test_patching.renamed_os.path							,			_PatchedModuleObj    )
          assert _test_patching.renamed_os.path.join is mock

          # check renamed_path.join
          assert isinstance(_test_patching.renamed_path							,			_PatchedModuleObj    )
          assert _test_patching.renamed_path.join is mock

          # check renamed_join
          assert _test_patching.renamed_join is mock

          # check that the other attributes are untouched
          assert _test_patching.renamed_os.rename is original_rename
          assert _test_patching.renamed_path.dirname is original_dirname
          assert _test_patching.renamed_os.path.dirname is original_dirname

     # check that everthing is back to normal when the patch is over

     assert _test_patching.os is original_os
     assert _test_patching.path is original_path
     assert _test_patching.join is original_join

     assert _test_patching.renamed_os is original_os
     assert _test_patching.renamed_path is original_path
     assert _test_patching.renamed_join is original_join


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     assert _test_patching.open is open

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_builtin_mock__'
     # _test_patching has "open" in its globals
     assert _test_patching.open is open
     with patch_submodule(_test_patching							,			'open'							,			_UpperCamelCase    ):
          assert _test_patching.open is mock

     # check that everthing is back to normal when the patch is over

     assert _test_patching.open is open


def 	_lowerCAmelCase (						)				->      str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_missing_mock__'
     with patch_submodule(_test_patching							,			'pandas.read_csv'							,			_UpperCamelCase    ):
          pass


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_missing_builtin_mock__'
     # _test_patching doesn't have "len" in its globals
     assert getattr(_test_patching							,			'len'							,			_UpperCamelCase    ) is None
     with patch_submodule(_test_patching							,			'len'							,			_UpperCamelCase    ):
          assert _test_patching.len is mock
     assert _test_patching.len is len


def 	_lowerCAmelCase (						)				->      str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_start_and_stop_mock__'
     _SCREAMING_SNAKE_CASE							   =patch_submodule(_test_patching							,			'open'							,			_UpperCamelCase    )
     assert _test_patching.open is open
     patch.start()
     assert _test_patching.open is mock
     patch.stop()
     assert _test_patching.open is open


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     from os import rename as original_rename
     from os.path import dirname as original_dirname
     from os.path import join as original_join

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_join__'
     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_dirname__'
     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_rename__'
     assert _test_patching.os.path.join is original_join
     assert _test_patching.os.path.dirname is original_dirname
     assert _test_patching.os.rename is original_rename

     with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
          with patch_submodule(_test_patching							,			'os.rename'							,			_UpperCamelCase    ):
               with patch_submodule(_test_patching							,			'os.path.dirname'							,			_UpperCamelCase    ):
                    assert _test_patching.os.path.join is mock_join
                    assert _test_patching.os.path.dirname is mock_dirname
                    assert _test_patching.os.rename is mock_rename

    # try another order
     with patch_submodule(_test_patching							,			'os.rename'							,			_UpperCamelCase    ):
          with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
               with patch_submodule(_test_patching							,			'os.path.dirname'							,			_UpperCamelCase    ):
                    assert _test_patching.os.path.join is mock_join
                    assert _test_patching.os.path.dirname is mock_dirname
                    assert _test_patching.os.rename is mock_rename

     assert _test_patching.os.path.join is original_join
     assert _test_patching.os.path.dirname is original_dirname
     assert _test_patching.os.rename is original_rename







def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_doesnt_exist_mock__'
     with patch_submodule(_test_patching							,			'__module_that_doesn_exist__.__attribute_that_doesn_exist__'							,			_UpperCamelCase    ):
          pass
     with patch_submodule(_test_patching							,			'os.__attribute_that_doesn_exist__'							,			_UpperCamelCase    ):
          pass
