






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 						SCREAMING_SNAKE_CASE				( snake_case  ,	unittest.TestCase					):





							"""simple docstring"""

							A_						  =					"ssube/stable-diffusion-x4-upscaler-onnx"

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

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

												# started as 128, should now be 512
												assert image.shape == (1, 5_12, 5_12, 3)
												_A								= np.array(
												    [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]   )
												assert np.abs(image_slice - expected_slice   ).max() < 1e-1

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

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

												assert image.shape == (1, 5_12, 5_12, 3)
												_A								= np.array(
												    [0.6_898_892, 0.59_240_556, 0.52_499_527, 0.58_866_215, 0.52_258_235, 0.52_572_715, 0.62_414_473, 0.6_174_387, 0.6_214_964]   )
												assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-1

							def       __A (  self:   Tuple   )    ->				Any:
												_A								= OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''   )
												_A								= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config   )
												pipe.set_progress_bar_config(disable=__A   )

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

												assert image.shape == (1, 5_12, 5_12, 3)
												_A								= np.array(
												    [0.7_659_278, 0.76_437_664, 0.75_579_107, 0.7_691_116, 0.77_666_986, 0.7_727_672, 0.7_758_664, 0.7_812_226, 0.76_942_515]   )

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

							def       __A (  self:   List[Any]   )    ->				Optional[Any]:
												_A								= OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''   )
												_A								= EulerDiscreteScheduler.from_config(pipe.scheduler.config   )
												pipe.set_progress_bar_config(disable=__A   )

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

												assert image.shape == (1, 5_12, 5_12, 3)
												_A								= np.array(
												    [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]   )
												assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-1



							def       __A (  self:   List[Any]   )    ->				Dict:
												_A								= OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint		,     provider='''CPUExecutionProvider'''   )
												_A								= EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config   )
												pipe.set_progress_bar_config(disable=__A   )

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

												assert image.shape == (1, 5_12, 5_12, 3)
												_A								= np.array(
												    [0.77_424_496, 0.773_601, 0.7_645_288, 0.7_769_598, 0.7_772_739, 0.7_738_688, 0.78_187_233, 0.77_879_584, 0.767_043]   )

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





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





							"""simple docstring"""

							@property
							def       __A (  self:   List[Any]   )    ->				Union[str, Any]:
												return (
												    "CUDAExecutionProvider",
												    {
												        "gpu_mem_limit": "15000000000",  # 15GB
												        "arena_extend_strategy": "kSameAsRequested",
												    },
												)

							@property
							def       __A (  self:   str   )    ->				str:
												_A								= ort.SessionOptions()
												_A								= False
												return options

							def       __A (  self:   List[Any]   )    ->				Optional[int]:
												_A								= load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg'''   )
												_A								= init_image.resize((1_28, 1_28)   )
												# using the PNDM scheduler by default
												_A								= 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								= '''A fantasy landscape, trending on artstation'''

												_A								= torch.manual_seed(0   )
												_A								= pipe(
												    prompt=__A		,     image=__A		,     guidance_scale=7.5		,     num_inference_steps=10		,     generator=__A		,     output_type='''np'''		,     )
												_A								= output.images
												_A								= images[0, 2_55:2_58, 3_83:3_86, -1]

												assert images.shape == (1, 5_12, 5_12, 3)
												_A								= np.array([0.4_883, 0.4_947, 0.4_980, 0.4_975, 0.4_982, 0.4_980, 0.5_000, 0.5_006, 0.4_972]   )
												# 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:   Tuple   )    ->				Tuple:
												_A								= load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg'''   )
												_A								= init_image.resize((1_28, 1_28)   )
												_A								= LMSDiscreteScheduler.from_pretrained(
												    '''ssube/stable-diffusion-x4-upscaler-onnx'''		,     subfolder='''scheduler'''   )
												_A								= 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								= '''A fantasy landscape, trending on artstation'''

												_A								= torch.manual_seed(0   )
												_A								= pipe(
												    prompt=__A		,     image=__A		,     guidance_scale=7.5		,     num_inference_steps=20		,     generator=__A		,     output_type='''np'''		,     )
												_A								= output.images
												_A								= images[0, 2_55:2_58, 3_83:3_86, -1]

												assert images.shape == (1, 5_12, 5_12, 3)
												_A								= np.array(
												    [0.50_173_753, 0.50_223_356, 0.502_039, 0.50_233_036, 0.5_023_725, 0.5_022_601, 0.5_018_758, 0.50_234_085, 0.50_241_566]   )
												# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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











import flax.linen as nn
import jax
import jax.numpy as jnp





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					jnp.floataa

							def       __A (  self:   Tuple   )    ->				Tuple:
												_A								= nn.Conv(
												    self.out_channels		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )



							def __call__(  self:   Dict		,     __A:   Dict   )    ->				Tuple:
												_A							,_A							,_A							,_A								= hidden_states.shape
												_A								= jax.image.resize(
												    __A		,     shape=(batch, height * 2, width * 2, channels)		,     method='''nearest'''		,     )
												_A								= self.conv(__A   )
												return hidden_states





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					jnp.floataa

							def       __A (  self:   List[str]   )    ->				Tuple:
												_A								= nn.Conv(
												    self.out_channels		,     kernel_size=(3, 3)		,     strides=(2, 2)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )



							def __call__(  self:   Union[str, Any]		,     __A:   List[Any]   )    ->				Union[str, Any]:
												# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
												# hidden_states = jnp.pad(hidden_states, pad_width=pad)
												_A								= self.conv(__A   )
												return hidden_states





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					None
							A_						  =					0.0
							A_						  =					None
							A_						  =					jnp.floataa

							def       __A (  self:   Dict   )    ->				Dict:
												_A								= self.in_channels if self.out_channels is None else self.out_channels

												_A								= nn.GroupNorm(num_groups=32		,     epsilon=1e-5   )
												_A								= nn.Conv(
												    __A		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )

												_A								= nn.Dense(__A		,     dtype=self.dtype   )

												_A								= nn.GroupNorm(num_groups=32		,     epsilon=1e-5   )
												_A								= nn.Dropout(self.dropout_prob   )
												_A								= nn.Conv(
												    __A		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )

												_A								= self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

												_A								= None
												if use_nin_shortcut:
																	_A								= nn.Conv(
																	    __A		,     kernel_size=(1, 1)		,     strides=(1, 1)		,     padding='''VALID'''		,     dtype=self.dtype		,     )



							def __call__(  self:   Dict		,     __A:   List[Any]		,     __A:   List[Any]		,     __A:   Any=True   )    ->				List[Any]:
												_A								= hidden_states
												_A								= self.norma(__A   )
												_A								= nn.swish(__A   )
												_A								= self.conva(__A   )

												_A								= self.time_emb_proj(nn.swish(__A   )   )
												_A								= jnp.expand_dims(jnp.expand_dims(__A		,     1   )		,     1   )
												_A								= hidden_states + temb

												_A								= self.norma(__A   )
												_A								= nn.swish(__A   )
												_A								= self.dropout(__A		,     __A   )
												_A								= self.conva(__A   )

												if self.conv_shortcut is not None:
																	_A								= self.conv_shortcut(__A   )

												return hidden_states + residual



