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


  '''simple docstring'''







  a__           =		"""ssube/stable-diffusion-x4-upscaler-onnx"""
  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: Dict=0    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       floats_tensor((1, 3, 128, 128)			,							rng=random.Random(UpperCamelCase__    )    )
    __magic_name__				     =       torch.manual_seed(UpperCamelCase__    )
    __magic_name__				     =       {
        """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			: Optional[Any]    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,							provider="""CPUExecutionProvider"""    )
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

    # started as 128, should now be 512
    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array(
        [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]    )
    assert np.abs(image_slice - expected_slice    ).max() < 1E-1
  def 							_lowercase							( self			: int    )			->	Optional[int]:







    """simple docstring"""






    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,							provider="""CPUExecutionProvider"""    )
    __magic_name__				     =       PNDMScheduler.from_config(pipe.scheduler.config			,							skip_prk_steps=UpperCamelCase__    )
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array(
        [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964]    )
    assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-1
  def 							_lowercase							( self			: Dict    )			->	Dict:







    """simple docstring"""






    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,							provider="""CPUExecutionProvider"""    )
    __magic_name__				     =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array(
        [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515]    )

    assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-1
  def 							_lowercase							( self			: Any    )			->	Optional[Any]:







    """simple docstring"""






    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,							provider="""CPUExecutionProvider"""    )
    __magic_name__				     =       EulerDiscreteScheduler.from_config(pipe.scheduler.config    )
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array(
        [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]    )
    assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-1





  def 							_lowercase							( self			: int    )			->	Dict:







    """simple docstring"""






    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,							provider="""CPUExecutionProvider"""    )
    __magic_name__				     =       EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config    )
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array(
        [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043]    )

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







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


  '''simple docstring'''







  @property
  def 							_lowercase							( self			: List[Any]    )			->	Optional[Any]:







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__				     =       ort.SessionOptions()
    __magic_name__				     =       False
    return options
  def 							_lowercase							( self			: List[Any]    )			->	Dict:







    """simple docstring"""






    __magic_name__				     =       load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
        """/img2img/sketch-mountains-input.jpg"""    )
    __magic_name__				     =       init_image.resize((128, 128)    )
    # using the PNDM scheduler by default
    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(
        """ssube/stable-diffusion-x4-upscaler-onnx"""			,							provider=self.gpu_provider			,							sess_options=self.gpu_options			,							)
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

    __magic_name__				     =       """A fantasy landscape, trending on artstation"""

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

    assert images.shape == (1, 512, 512, 3)
    __magic_name__				     =       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 							_lowercase							( self			: Tuple    )			->	int:







    """simple docstring"""






    __magic_name__				     =       load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
        """/img2img/sketch-mountains-input.jpg"""    )
    __magic_name__				     =       init_image.resize((128, 128)    )
    __magic_name__				     =       LMSDiscreteScheduler.from_pretrained(
        """ssube/stable-diffusion-x4-upscaler-onnx"""			,							subfolder="""scheduler"""    )
    __magic_name__				     =       OnnxStableDiffusionUpscalePipeline.from_pretrained(
        """ssube/stable-diffusion-x4-upscaler-onnx"""			,							scheduler=UpperCamelCase__			,							provider=self.gpu_provider			,							sess_options=self.gpu_options			,							)
    pipe.set_progress_bar_config(disable=UpperCamelCase__    )

    __magic_name__				     =       """A fantasy landscape, trending on artstation"""

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

    assert images.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array(
        [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566]    )
    # 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


__lowerCAmelCase			:     List[str]           =   {
    'configuration_xlm': ['XLM_PRETRAINED_CONFIG_ARCHIVE_MAP', 'XLMConfig', 'XLMOnnxConfig'],
    'tokenization_xlm': ['XLMTokenizer'],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCAmelCase			:     str           =   [
        '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:
    __lowerCAmelCase			:     Dict           =   [
        '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

    __lowerCAmelCase			:     Any           =   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

