
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__  (	__lowercase   ,unittest.TestCase   ):
   # TODO: is there an appropriate internal test set?
   _SCREAMING_SNAKE_CASE				:							Optional[int]     	=				"ssube/stable-diffusion-x4-upscaler-onnx"
   def 				lowerCAmelCase      (self			: str ,    snake_case_			: Union[str, Any]=0       ):
    __a			:  Optional[Any]		    =  floats_tensor((1, 3, 1_2_8, 1_2_8) ,    rng=random.Random(snake_case_       )       )
    __a			:  Tuple		    =  torch.manual_seed(snake_case_       )
    __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 				lowerCAmelCase      (self			: Tuple       ):
    __a			:  int		    =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint ,    provider='''CPUExecutionProvider'''       )
    pipe.set_progress_bar_config(disable=snake_case_       )

    __a			:  Dict		    =  self.get_dummy_inputs()
    __a			:  Dict		    =  pipe(**snake_case_       ).images
    __a			:  List[Any]		    =  image[0, -3:, -3:, -1].flatten()

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

    __a			:  Tuple		    =  self.get_dummy_inputs()
    __a			:  Any		    =  pipe(**snake_case_       ).images
    __a			:  Dict		    =  image[0, -3:, -3:, -1]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a			:  Optional[int]		    =  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			: Tuple       ):
    __a			:  Any		    =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint ,    provider='''CPUExecutionProvider'''       )
    __a			:  str		    =  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
    pipe.set_progress_bar_config(disable=snake_case_       )

    __a			:  int		    =  self.get_dummy_inputs()
    __a			:  Tuple		    =  pipe(**snake_case_       ).images
    __a			:  List[str]		    =  image[0, -3:, -3:, -1]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a			:  List[str]		    =  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			: str       ):
    __a			:  Tuple		    =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint ,    provider='''CPUExecutionProvider'''       )
    __a			:  Tuple		    =  EulerDiscreteScheduler.from_config(pipe.scheduler.config       )
    pipe.set_progress_bar_config(disable=snake_case_       )

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

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a			:  str		    =  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			: Union[str, Any]       ):
    __a			:  Optional[Any]		    =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint ,    provider='''CPUExecutionProvider'''       )
    __a			:  int		    =  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config       )
    pipe.set_progress_bar_config(disable=snake_case_       )

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

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a			:  int		    =  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 		UpperCamelCase__  (	unittest.TestCase   ):
   @property
   def 				lowerCAmelCase      (self			: Dict       ):
    return (
        "CUDAExecutionProvider",
        {
            "gpu_mem_limit": "15000000000",  # 15GB
            "arena_extend_strategy": "kSameAsRequested",
        },
    )
   @property
   def 				lowerCAmelCase      (self			: Union[str, Any]       ):
    __a			:  Union[str, Any]		    =  ort.SessionOptions()
    __a			:  Tuple		    =  False
    return options
   def 				lowerCAmelCase      (self			: Tuple       ):
    __a			:  List[Any]		    =  load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/img2img/sketch-mountains-input.jpg'''       )
    __a			:  List[Any]		    =  init_image.resize((1_2_8, 1_2_8)       )
    # using the PNDM scheduler by default
    __a			:  str		    =  OnnxStableDiffusionUpscalePipeline.from_pretrained(
        '''ssube/stable-diffusion-x4-upscaler-onnx''' ,    provider=self.gpu_provider ,    sess_options=self.gpu_options ,    )
    pipe.set_progress_bar_config(disable=snake_case_       )

    __a			:  Optional[int]		    =  '''A fantasy landscape, trending on artstation'''

    __a			:  Optional[int]		    =  torch.manual_seed(0       )
    __a			:  Tuple		    =  pipe(
        prompt=snake_case_ ,    image=snake_case_ ,    guidance_scale=7.5 ,    num_inference_steps=1_0 ,    generator=snake_case_ ,    output_type='''np''' ,    )
    __a			:  str		    =  output.images
    __a			:  str		    =  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)
    __a			:  Optional[Any]		    =  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			: Any       ):
    __a			:  int		    =  load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/img2img/sketch-mountains-input.jpg'''       )
    __a			:  Optional[int]		    =  init_image.resize((1_2_8, 1_2_8)       )
    __a			:  Tuple		    =  LMSDiscreteScheduler.from_pretrained(
        '''ssube/stable-diffusion-x4-upscaler-onnx''' ,    subfolder='''scheduler'''       )
    __a			:  int		    =  OnnxStableDiffusionUpscalePipeline.from_pretrained(
        '''ssube/stable-diffusion-x4-upscaler-onnx''' ,    scheduler=snake_case_ ,    provider=self.gpu_provider ,    sess_options=self.gpu_options ,    )
    pipe.set_progress_bar_config(disable=snake_case_       )

    __a			:  Union[str, Any]		    =  '''A fantasy landscape, trending on artstation'''

    __a			:  Any		    =  torch.manual_seed(0       )
    __a			:  Any		    =  pipe(
        prompt=snake_case_ ,    image=snake_case_ ,    guidance_scale=7.5 ,    num_inference_steps=2_0 ,    generator=snake_case_ ,    output_type='''np''' ,    )
    __a			:  str		    =  output.images
    __a			:  List[str]		    =  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)
    __a			:  Dict		    =  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_tf_available, is_torch_available


lowercase__          ={
    'configuration_data2vec_audio': ['DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Data2VecAudioConfig'],
    'configuration_data2vec_text': [
        'DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'Data2VecTextConfig',
        'Data2VecTextOnnxConfig',
    ],
    'configuration_data2vec_vision': [
        'DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'Data2VecVisionConfig',
        'Data2VecVisionOnnxConfig',
    ],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase__          =[
          'DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST',
          'Data2VecAudioForAudioFrameClassification',
          'Data2VecAudioForCTC',
          'Data2VecAudioForSequenceClassification',
          'Data2VecAudioForXVector',
          'Data2VecAudioModel',
          'Data2VecAudioPreTrainedModel',
      ]
      lowercase__          =[
          'DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'Data2VecTextForCausalLM',
          'Data2VecTextForMaskedLM',
          'Data2VecTextForMultipleChoice',
          'Data2VecTextForQuestionAnswering',
          'Data2VecTextForSequenceClassification',
          'Data2VecTextForTokenClassification',
          'Data2VecTextModel',
          'Data2VecTextPreTrainedModel',
      ]
      lowercase__          =[
          'DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST',
          'Data2VecVisionForImageClassification',
          'Data2VecVisionForMaskedImageModeling',
          'Data2VecVisionForSemanticSegmentation',
          'Data2VecVisionModel',
          'Data2VecVisionPreTrainedModel',
      ]

if is_tf_available():
      lowercase__          =[
          'TFData2VecVisionForImageClassification',
          'TFData2VecVisionForSemanticSegmentation',
          'TFData2VecVisionModel',
          'TFData2VecVisionPreTrainedModel',
      ]

if TYPE_CHECKING:
      from .configuration_dataavec_audio import DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP, DataaVecAudioConfig
      from .configuration_dataavec_text import (
          DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP,
          DataaVecTextConfig,
          DataaVecTextOnnxConfig,
      )
      from .configuration_dataavec_vision import (
          DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP,
          DataaVecVisionConfig,
          DataaVecVisionOnnxConfig,
      )

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_dataavec_audio import (
                DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST,
                DataaVecAudioForAudioFrameClassification,
                DataaVecAudioForCTC,
                DataaVecAudioForSequenceClassification,
                DataaVecAudioForXVector,
                DataaVecAudioModel,
                DataaVecAudioPreTrainedModel,
            )
            from .modeling_dataavec_text import (
                DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
                DataaVecTextForCausalLM,
                DataaVecTextForMaskedLM,
                DataaVecTextForMultipleChoice,
                DataaVecTextForQuestionAnswering,
                DataaVecTextForSequenceClassification,
                DataaVecTextForTokenClassification,
                DataaVecTextModel,
                DataaVecTextPreTrainedModel,
            )
            from .modeling_dataavec_vision import (
                DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST,
                DataaVecVisionForImageClassification,
                DataaVecVisionForMaskedImageModeling,
                DataaVecVisionForSemanticSegmentation,
                DataaVecVisionModel,
                DataaVecVisionPreTrainedModel,
            )
      if is_tf_available():
            from .modeling_tf_dataavec_vision import (
                TFDataaVecVisionForImageClassification,
                TFDataaVecVisionForSemanticSegmentation,
                TFDataaVecVisionModel,
                TFDataaVecVisionPreTrainedModel,
            )

else:
      import sys

      lowercase__          =_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
