

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							):
						"""simple docstring"""






						_A    :				Dict       					=			'''ssube/stable-diffusion-x4-upscaler-onnx'''





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

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

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

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

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

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

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

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

									assert image.shape == (1, 512, 512, 3)
									a_						:						Optional[int] 							= 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 			UpperCamelCase__					(    self    )     -> Any:
									a_						:						int 							= OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,      provider="""CPUExecutionProvider"""    )
									a_						:						Any 							= EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config    )
									pipe.set_progress_bar_config(disable=_lowercase    )

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

									assert image.shape == (1, 512, 512, 3)
									a_						:						Optional[int] 							= 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        A__(unittest.TestCase							):
						"""simple docstring"""






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





						@property
						def 			UpperCamelCase__					(    self    )     -> Dict:
									a_						:						List[str] 							= ort.SessionOptions()
									a_						:						Optional[int] 							= False
									return options





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

									a_						:						str 							= """A fantasy landscape, trending on artstation"""

									a_						:						Optional[int] 							= torch.manual_seed(0    )
									a_						:						Optional[int] 							= pipe(
									    prompt=_lowercase  ,      image=_lowercase  ,      guidance_scale=7.5  ,      num_inference_steps=10  ,      generator=_lowercase  ,      output_type="""np"""  ,      )
									a_						:						int 							= output.images
									a_						:						Dict 							= images[0, 255:258, 383:386, -1]

									assert images.shape == (1, 512, 512, 3)
									a_						:						Union[str, Any] 							= 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 			UpperCamelCase__					(    self    )     -> List[str]:
									a_						:						str 							= load_image(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
									    """/img2img/sketch-mountains-input.jpg"""    )
									a_						:						List[Any] 							= init_image.resize((128, 128)    )
									a_						:						List[Any] 							= LMSDiscreteScheduler.from_pretrained(
									    """ssube/stable-diffusion-x4-upscaler-onnx"""  ,      subfolder="""scheduler"""    )
									a_						:						int 							= 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    )

									a_						:						Optional[int] 							= """A fantasy landscape, trending on artstation"""

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

									assert images.shape == (1, 512, 512, 3)
									a_						:						Union[str, Any] 							= 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 typing import TYPE_CHECKING

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


__snake_case				:      Optional[int]   		=			{
    """configuration_xlm""": ["""XLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XLMConfig""", """XLMOnnxConfig"""],
    """tokenization_xlm""": ["""XLMTokenizer"""],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      Optional[Any]   		=			[
	    """XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """XLMForMultipleChoice""",
	    """XLMForQuestionAnswering""",
	    """XLMForQuestionAnsweringSimple""",
	    """XLMForSequenceClassification""",
	    """XLMForTokenClassification""",
	    """XLMModel""",
	    """XLMPreTrainedModel""",
	    """XLMWithLMHeadModel""",
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      Optional[Any]   		=			[
	    """TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """TFXLMForMultipleChoice""",
	    """TFXLMForQuestionAnsweringSimple""",
	    """TFXLMForSequenceClassification""",
	    """TFXLMForTokenClassification""",
	    """TFXLMMainLayer""",
	    """TFXLMModel""",
	    """TFXLMPreTrainedModel""",
	    """TFXLMWithLMHeadModel""",
	]


if TYPE_CHECKING:
	from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
	from .tokenization_xlm import XLMTokenizer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_xlm import (
		    XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    XLMForMultipleChoice,
		    XLMForQuestionAnswering,
		    XLMForQuestionAnsweringSimple,
		    XLMForSequenceClassification,
		    XLMForTokenClassification,
		    XLMModel,
		    XLMPreTrainedModel,
		    XLMWithLMHeadModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_xlm import (
		    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFXLMForMultipleChoice,
		    TFXLMForQuestionAnsweringSimple,
		    TFXLMForSequenceClassification,
		    TFXLMForTokenClassification,
		    TFXLMMainLayer,
		    TFXLMModel,
		    TFXLMPreTrainedModel,
		    TFXLMWithLMHeadModel,
		)

else:
	import sys

	__snake_case				:      Union[str, Any]   		=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


