



"""simple docstring"""
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__       (							__UpperCAmelCase					,     unittest.TestCase      ):
    lowercase__          =      "ssube/stable-diffusion-x4-upscaler-onnx"

    def     _UpperCAmelCase		( self	: Optional[Any]       ,     lowerCAmelCase_	: Optional[int]=0):







      """simple docstring"""







      lowercase_	    =       floats_tensor((1, 3, 1_2_8, 1_2_8)       ,     rng=random.Random(lowerCamelCase__))
      lowercase_	    =       torch.manual_seed(lowerCamelCase__)
      lowercase_	    =       {
          """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	: str):







      """simple docstring"""







      lowercase_	    =       OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,     provider="""CPUExecutionProvider""")
      pipe.set_progress_bar_config(disable=lowerCamelCase__)

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

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







      """simple docstring"""







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

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

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







      """simple docstring"""







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

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

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      lowercase_	    =       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     _UpperCAmelCase		( self	: Optional[Any]):







      """simple docstring"""







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

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

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      lowercase_	    =       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     _UpperCAmelCase		( self	: Union[str, Any]):







      """simple docstring"""







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

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

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

    @property
    def     _UpperCAmelCase		( self	: Optional[Any]):







      """simple docstring"""







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

    @property
    def     _UpperCAmelCase		( self	: Tuple):







      """simple docstring"""







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

    def     _UpperCAmelCase		( self	: str):







      """simple docstring"""







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

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

      lowercase_	    =       torch.manual_seed(0)
      lowercase_	    =       pipe(
          prompt=lowerCamelCase__       ,     image=lowerCamelCase__       ,     guidance_scale=7.5       ,     num_inference_steps=1_0       ,     generator=lowerCamelCase__       ,     output_type="""np"""       ,     )
      lowercase_	    =       output.images
      lowercase_	    =       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)
      lowercase_	    =       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     _UpperCAmelCase		( self	: Optional[int]):







      """simple docstring"""







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

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

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







class     __lowerCAmelCase :

			"""simple docstring"""







			def __init__(      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    0
							__lowerCamelCase		       =    0
							__lowerCamelCase		       =    {}


			def       lowercase_    (      self			,						lowerCamelCase__							)    ->      Tuple:



							'''simple docstring'''
							if vertex not in self.adjacency:
											__lowerCamelCase		       =    {}
											self.num_vertices += 1


			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      str:



							'''simple docstring'''
							self.add_vertex(lowerCamelCase__							)
							self.add_vertex(lowerCamelCase__							)

							if head == tail:
											return

							__lowerCamelCase		       =    weight
							__lowerCamelCase		       =    weight


			def       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''
							__lowerCamelCase		       =    self.get_edges()
							for edge in edges:
											__lowerCamelCase       ,						__lowerCamelCase       ,						__lowerCamelCase		       =    edge
											edges.remove((tail, head, weight)							)
							for i in range(len(lowerCamelCase__							)							):
											__lowerCamelCase		       =    list(edges[i]							)

							edges.sort(key=lambda lowerCamelCase__							: e[2]							)
							for i in range(len(lowerCamelCase__							) - 1							):
											if edges[i][2] >= edges[i + 1][2]:
															__lowerCamelCase		       =    edges[i][2] + 1
							for edge in edges:
											__lowerCamelCase       ,						__lowerCamelCase       ,						__lowerCamelCase		       =    edge
											__lowerCamelCase		       =    weight
											__lowerCamelCase		       =    weight


			def __str__(      self							)    ->      Union[str, Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    ''
							for tail in self.adjacency:
											for head in self.adjacency[tail]:
															__lowerCamelCase		       =    self.adjacency[head][tail]
															string += f"""{head} -> {tail} == {weight}\n"""
							return string.rstrip('\n'							)


			def       lowercase_    (      self							)    ->      Optional[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    []
							for tail in self.adjacency:
											for head in self.adjacency[tail]:
															output.append((tail, head, self.adjacency[head][tail])							)
							return output


			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							return self.adjacency.keys()


			@staticmethod
			def       lowercase_    (      lowerCamelCase__=None			,						lowerCamelCase__=None							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase		       =    Graph()
							if vertices is None:
											__lowerCamelCase		       =    []
							if edges is None:
											__lowerCamelCase		       =    []
							for vertex in vertices:
											g.add_vertex(lowerCamelCase__							)
							for edge in edges:
											g.add_edge(*lowerCamelCase__							)
							return g







			class     __lowerCAmelCase :

						"""simple docstring"""







						def __init__(      self							)    ->      Tuple:



										'''simple docstring'''
										__lowerCamelCase		       =    {}
										__lowerCamelCase		       =    {}


						def __len__(      self							)    ->      Tuple:



										'''simple docstring'''
										return len(self.parent							)


						def       lowercase_    (      self			,						lowerCamelCase__							)    ->      List[Any]:



										'''simple docstring'''
										if item in self.parent:
														return self.find(lowerCamelCase__							)

										__lowerCamelCase		       =    item
										__lowerCamelCase		       =    0
										return item


						def       lowercase_    (      self			,						lowerCamelCase__							)    ->      List[str]:



										'''simple docstring'''
										if item not in self.parent:
														return self.make_set(lowerCamelCase__							)
										if item != self.parent[item]:
														__lowerCamelCase		       =    self.find(self.parent[item]							)
										return self.parent[item]






						def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      int:



										'''simple docstring'''
										__lowerCamelCase		       =    self.find(lowerCamelCase__							)
										__lowerCamelCase		       =    self.find(lowerCamelCase__							)

										if roota == roota:
														return roota

										if self.rank[roota] > self.rank[roota]:
														__lowerCamelCase		       =    roota
														return roota

										if self.rank[roota] < self.rank[roota]:
														__lowerCamelCase		       =    roota
														return roota

										if self.rank[roota] == self.rank[roota]:
														self.rank[roota] += 1
														__lowerCamelCase		       =    roota
														return roota
										return None




			@staticmethod
			def       lowercase_    (      lowerCamelCase__							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase		       =    graph.num_vertices

							__lowerCamelCase		       =    Graph.UnionFind()
							__lowerCamelCase		       =    []
							while num_components > 1:
											__lowerCamelCase		       =    {}
											for vertex in graph.get_vertices():
															__lowerCamelCase		       =    -1

											__lowerCamelCase		       =    graph.get_edges()
											for edge in edges:
															__lowerCamelCase       ,						__lowerCamelCase       ,						__lowerCamelCase		       =    edge
															edges.remove((tail, head, weight)							)
											for edge in edges:
															__lowerCamelCase       ,						__lowerCamelCase       ,						__lowerCamelCase		       =    edge
															__lowerCamelCase		       =    union_find.find(lowerCamelCase__							)
															__lowerCamelCase		       =    union_find.find(lowerCamelCase__							)
															if seta != seta:
																			if cheap_edge[seta] == -1 or cheap_edge[seta][2] > weight:
																							__lowerCamelCase		       =    [head, tail, weight]

																			if cheap_edge[seta] == -1 or cheap_edge[seta][2] > weight:
																							__lowerCamelCase		       =    [head, tail, weight]
											for vertex in cheap_edge:
															if cheap_edge[vertex] != -1:
																			__lowerCamelCase       ,						__lowerCamelCase       ,						__lowerCamelCase		       =    cheap_edge[vertex]
																			if union_find.find(lowerCamelCase__							) != union_find.find(lowerCamelCase__							):
																							union_find.union(lowerCamelCase__			,						lowerCamelCase__							)
																							mst_edges.append(cheap_edge[vertex]							)
																							__lowerCamelCase		       =    num_components - 1
							__lowerCamelCase		       =    Graph.build(edges=lowerCamelCase__							)
							return mst


