
"""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       lowerCAmelCase__				(						A__ ,	unittest.TestCase	):
     __a						=					'''ssube/stable-diffusion-x4-upscaler-onnx'''







     def      lowercase	(							self			:			List[str]    ,      _lowerCamelCase			:			List[Any]=0			):
       _snake_case					     =       floats_tensor((1, 3, 128, 128)    ,      rng=random.Random(lowerCAmelCase__			)			)
       _snake_case					     =       torch.manual_seed(lowerCAmelCase__			)
       _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      lowercase	(							self			:			Any			):
       _snake_case					     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    ,      provider='''CPUExecutionProvider'''			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

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

       # started as 128, should now be 512
       assert image.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array(
           [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]			)
       assert np.abs(image_slice - expected_slice			).max() < 1e-1







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    ,      provider='''CPUExecutionProvider'''			)
       _snake_case					     =       PNDMScheduler.from_config(pipe.scheduler.config    ,      skip_prk_steps=lowerCAmelCase__			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

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

       assert image.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array(
           [0.6_8_9_8_8_9_2, 0.5_9_2_4_0_5_5_6, 0.5_2_4_9_9_5_2_7, 0.5_8_8_6_6_2_1_5, 0.5_2_2_5_8_2_3_5, 0.5_2_5_7_2_7_1_5, 0.6_2_4_1_4_4_7_3, 0.6_1_7_4_3_8_7, 0.6_2_1_4_9_6_4]			)
       assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-1







     def      lowercase	(							self			:			int			):
       _snake_case					     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    ,      provider='''CPUExecutionProvider'''			)
       _snake_case					     =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

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

       assert image.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array(
           [0.7_6_5_9_2_7_8, 0.7_6_4_3_7_6_6_4, 0.7_5_5_7_9_1_0_7, 0.7_6_9_1_1_1_6, 0.7_7_6_6_6_9_8_6, 0.7_7_2_7_6_7_2, 0.7_7_5_8_6_6_4, 0.7_8_1_2_2_2_6, 0.7_6_9_4_2_5_1_5]			)

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







     def      lowercase	(							self			:			Dict			):
       _snake_case					     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    ,      provider='''CPUExecutionProvider'''			)
       _snake_case					     =       EulerDiscreteScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

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

       assert image.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array(
           [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]			)
       assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-1







     def      lowercase	(							self			:			Dict			):
       _snake_case					     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint    ,      provider='''CPUExecutionProvider'''			)
       _snake_case					     =       EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

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

       assert image.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array(
           [0.7_7_4_2_4_4_9_6, 0.7_7_3_6_0_1, 0.7_6_4_5_2_8_8, 0.7_7_6_9_5_9_8, 0.7_7_7_2_7_3_9, 0.7_7_3_8_6_8_8, 0.7_8_1_8_7_2_3_3, 0.7_7_8_7_9_5_8_4, 0.7_6_7_0_4_3]			)

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






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







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







     @property
     def      lowercase	(							self			:			Tuple			):
       _snake_case					     =       ort.SessionOptions()
       _snake_case					     =       False
       return options







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

       _snake_case					     =       '''A fantasy landscape, trending on artstation'''

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

       assert images.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array([0.4_8_8_3, 0.4_9_4_7, 0.4_9_8_0, 0.4_9_7_5, 0.4_9_8_2, 0.4_9_8_0, 0.5_0_0_0, 0.5_0_0_6, 0.4_9_7_2]			)
       # 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			:			int			):
       _snake_case					     =       load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/img2img/sketch-mountains-input.jpg'''			)
       _snake_case					     =       init_image.resize((128, 128)			)
       _snake_case					     =       LMSDiscreteScheduler.from_pretrained(
           '''ssube/stable-diffusion-x4-upscaler-onnx'''    ,      subfolder='''scheduler'''			)
       _snake_case					     =       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__			)

       _snake_case					     =       '''A fantasy landscape, trending on artstation'''

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

       assert images.shape == (1, 512, 512, 3)
       _snake_case					     =       np.array(
           [0.5_0_1_7_3_7_5_3, 0.5_0_2_2_3_3_5_6, 0.5_0_2_0_3_9, 0.5_0_2_3_3_0_3_6, 0.5_0_2_3_7_2_5, 0.5_0_2_2_6_0_1, 0.5_0_1_8_7_5_8, 0.5_0_2_3_4_0_8_5, 0.5_0_2_4_1_5_6_6]			)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







from importlib import import_module

from .logging import get_logger


__lowerCAmelCase	: str  				=get_logger(__name__)
class        _lowercase      :






				'''simple docstring'''



				def __init__(       self			:List[Any]			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:str=None	)		->			int:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		attrs or []
										if module is not None:
																for key in module.__dict__:
																						if key in attrs or not key.startswith('''__'''	):
																												setattr(self			,			lowerCAmelCase__			,			getattr(lowerCAmelCase__			,			lowerCAmelCase__	)	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		module._original_module if isinstance(lowerCAmelCase__			,			_PatchedModuleObj	) else module




class        _lowercase      :






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			Optional[int]						=		[]


				def __init__(       self			:Tuple			,			lowerCAmelCase__			:Optional[Any]			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:Any			,			lowerCAmelCase__			:Dict=None	)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		obj
										__SCREAMING_SNAKE_CASE						:     str					=		target
										__SCREAMING_SNAKE_CASE						:     Dict					=		new
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		target.split('''.'''	)[0]
										__SCREAMING_SNAKE_CASE						:     List[str]					=		{}
										__SCREAMING_SNAKE_CASE						:     Tuple					=		attrs or []


				def __enter__(       self			:int	)		->			Dict:
										*__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.target.split('''.'''	)

										# Patch modules:
										# it's used to patch attributes of submodules like "os.path.join";
										# in this case we need to patch "os" and "os.path"

										for i in range(len(lowerCAmelCase__	)	):
																try:
																						__SCREAMING_SNAKE_CASE						:     Any					=		import_module('''.'''.join(submodules[: i + 1]	)	)
																except ModuleNotFoundError:
																						continue
																# We iterate over all the globals in self.obj in case we find "os" or "os.path"
																for attr in self.obj.__dir__():
																						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		getattr(self.obj			,			lowerCAmelCase__	)
																						# We don't check for the name of the global, but rather if its value *is* "os" or "os.path".
																						# This allows to patch renamed modules like "from os import path as ospath".
																						if obj_attr is submodule or (
																						    (isinstance(lowerCAmelCase__			,			_PatchedModuleObj	) and obj_attr._original_module is submodule)
																						):
																												__SCREAMING_SNAKE_CASE						:     int					=		obj_attr
																												# patch at top level
																												setattr(self.obj			,			lowerCAmelCase__			,			_PatchedModuleObj(lowerCAmelCase__			,			attrs=self.attrs	)	)
																												__SCREAMING_SNAKE_CASE						:     List[str]					=		getattr(self.obj			,			lowerCAmelCase__	)
																												# construct lower levels patches
																												for key in submodules[i + 1 :]:
																																		setattr(lowerCAmelCase__			,			lowerCAmelCase__			,			_PatchedModuleObj(getattr(lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__	)			,			attrs=self.attrs	)	)
																																		__SCREAMING_SNAKE_CASE						:     Tuple					=		getattr(lowerCAmelCase__			,			lowerCAmelCase__	)
																												# finally set the target attribute
																												setattr(lowerCAmelCase__			,			lowerCAmelCase__			,			self.new	)

        # Patch attribute itself:
        # it's used for builtins like "open",
        # and also to patch "os.path.join" we may also need to patch "join"
        # itself if it was imported as "from os.path import join".

										if submodules:  # if it's an attribute of a submodule like "os.path.join"
																try:
																						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		getattr(import_module('''.'''.join(lowerCAmelCase__	)	)			,			lowerCAmelCase__	)
																except (AttributeError, ModuleNotFoundError):
																						return
																# We iterate over all the globals in self.obj in case we find "os.path.join"
																for attr in self.obj.__dir__():
																						# We don't check for the name of the global, but rather if its value *is* "os.path.join".
																						# This allows to patch renamed attributes like "from os.path import join as pjoin".
																						if getattr(self.obj			,			lowerCAmelCase__	) is attr_value:
																												__SCREAMING_SNAKE_CASE						:     Any					=		getattr(self.obj			,			lowerCAmelCase__	)
																												setattr(self.obj			,			lowerCAmelCase__			,			self.new	)
										elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
																__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		globals()['''__builtins__'''][target_attr]
																setattr(self.obj			,			lowerCAmelCase__			,			self.new	)
										else:
																raise RuntimeError(f'''Tried to patch attribute {target_attr} instead of a submodule.'''	)


				def __exit__(       self			:str			,			*lowerCAmelCase__			:Union[str, Any]	)		->			Optional[int]:
										for attr in list(self.original	):
																setattr(self.obj			,			lowerCAmelCase__			,			self.original.pop(lowerCAmelCase__	)	)


				def 		__magic_name__(       self			:List[Any]	)		->			List[Any]:
										self.__enter__()
										self._active_patches.append(self	)



				def 		__magic_name__(       self			:Optional[int]	)		->			int:
										try:
																self._active_patches.remove(self	)
										except ValueError:
																# If the patch hasn't been started this will fail
																return None

										return self.__exit__()
