





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

			'''simple docstring'''







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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase=0							):



				'''simple docstring'''







				lowerCamelCase__              =	floats_tensor((1, 3, 1_2_8, 1_2_8) ,						rng=random.Random(__A							)							)
				lowerCamelCase__              =	torch.manual_seed(__A							)
				lowerCamelCase__              =	{
				    '''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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint ,						provider='''CPUExecutionProvider'''							)
				pipe.set_progress_bar_config(disable=__A							)

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

				# started as 128, should now be 512
				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				lowerCamelCase__              =	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       __A					(					unittest.TestCase				):

			'''simple docstring'''







			@property
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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





			@property
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

				lowerCamelCase__              =	torch.manual_seed(0							)
				lowerCamelCase__              =	pipe(
				    prompt=__A ,						image=__A ,						guidance_scale=7.5 ,						num_inference_steps=1_0 ,						generator=__A ,						output_type='''np''' ,						)
				lowerCamelCase__              =	output.images
				lowerCamelCase__              =	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)
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

				lowerCamelCase__              =	torch.manual_seed(0							)
				lowerCamelCase__              =	pipe(
				    prompt=__A ,						image=__A ,						guidance_scale=7.5 ,						num_inference_steps=2_0 ,						generator=__A ,						output_type='''np''' ,						)
				lowerCamelCase__              =	output.images
				lowerCamelCase__              =	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)
				lowerCamelCase__              =	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










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


_a						=   {
    "configuration_gpt_neo": ["GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoConfig", "GPTNeoOnnxConfig"],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_a						=   [
			    "GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "GPTNeoForCausalLM",
			    "GPTNeoForQuestionAnswering",
			    "GPTNeoForSequenceClassification",
			    "GPTNeoForTokenClassification",
			    "GPTNeoModel",
			    "GPTNeoPreTrainedModel",
			    "load_tf_weights_in_gpt_neo",
			]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_a						=   [
			    "FlaxGPTNeoForCausalLM",
			    "FlaxGPTNeoModel",
			    "FlaxGPTNeoPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig, GPTNeoOnnxConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_gpt_neo import (
						    GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
						    GPTNeoForCausalLM,
						    GPTNeoForQuestionAnswering,
						    GPTNeoForSequenceClassification,
						    GPTNeoForTokenClassification,
						    GPTNeoModel,
						    GPTNeoPreTrainedModel,
						    load_tf_weights_in_gpt_neo,
						)

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel


else:
			import sys

			_a						=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



