






"""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    lowercase_    (lowercase__,    unittest.TestCase     ):
      # TODO: is there an appropriate internal test set?
      A__	:		List[Any]       			=				"""ssube/stable-diffusion-x4-upscaler-onnx"""






      def 			lowerCamelCase__    (    self     ,     a_=0	)	->Dict:







        '''simple docstring'''


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







        '''simple docstring'''


        _a							=	OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,     provider="CPUExecutionProvider"	)
        pipe.set_progress_bar_config(disable=__lowercase	)

        _a							=	self.get_dummy_inputs()
        _a							=	pipe(**__lowercase	).images
        _a							=	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							=	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 			lowerCamelCase__    (    self	)	->Union[str, Any]:







        '''simple docstring'''


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

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

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







        '''simple docstring'''


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

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

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







        '''simple docstring'''


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

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

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







        '''simple docstring'''


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

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

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






      @property
      def 			lowerCamelCase__    (    self	)	->Tuple:







        '''simple docstring'''


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






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







        '''simple docstring'''


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






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







        '''simple docstring'''


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

        _a							=	torch.manual_seed(0	)
        _a							=	pipe(
            prompt=__lowercase     ,     image=__lowercase     ,     guidance_scale=7.5     ,     num_inference_steps=1_0     ,     generator=__lowercase     ,     output_type="np"     ,     )
        _a							=	output.images
        _a							=	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							=	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 			lowerCamelCase__    (    self	)	->Union[str, Any]:







        '''simple docstring'''


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

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









"""simple docstring"""




from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

UpperCamelCase							       =    TypeVar("""T""")
UpperCamelCase							       =    TypeVar("""U""")




class    lowercase_    (Generic[T, U]     ):






      def __init__(    self     ,     a_     ,     a_	)	->Any:







        '''simple docstring'''


        _a							=	key
        _a							=	val
        _a							=	None
        _a							=	None






      def __repr__(    self	)	->str:







        '''simple docstring'''


        return (
            f'''Node: key: {self.key}, val: {self.val}, '''
            f'''has next: {bool(self.next	)}, has prev: {bool(self.prev	)}'''
        )







class    lowercase_    (Generic[T, U]     ):






      def __init__(    self	)	->None:







        '''simple docstring'''


        _a							=	DoubleLinkedListNode(a_     ,     a_	)
        _a							=	DoubleLinkedListNode(a_     ,     a_	)
        _a			,	_a							=	self.rear, self.head






      def __repr__(    self	)	->str:







        '''simple docstring'''


        _a							=	["DoubleLinkedList"]
        _a							=	self.head
        while node.next is not None:
          rep.append(str(a_	)	)
          _a							=	node.next
        rep.append(str(self.rear	)	)
        return ",\n    ".join(a_	)






      def 			lowerCamelCase__    (    self     ,     a_	)	->None:







        '''simple docstring'''


        _a							=	self.rear.prev

        # All nodes other than self.head are guaranteed to have non-None previous
        assert previous is not None

        _a							=	node
        _a							=	previous
        _a							=	node
        _a							=	self.rear






      def 			lowerCamelCase__    (    self     ,     a_	)	->DoubleLinkedListNode[T, U] | None:







        '''simple docstring'''


        if node.prev is None or node.next is None:
          return None

        _a							=	node.next
        _a							=	node.prev
        _a							=	None
        _a							=	None
        return node







class    lowercase_    (Generic[T, U]     ):
      A__	:		dict[Callable[[T], U], LRUCache[T, U]]       			=				{}






      def __init__(    self     ,     a_	)	->int:







        '''simple docstring'''


        _a							=	DoubleLinkedList()
        _a							=	capacity
        _a							=	0
        _a							=	0
        _a							=	0
        _a							=	{}






      def __repr__(    self	)	->str:







        '''simple docstring'''


        return (
            f'''CacheInfo(hits={self.hits}, misses={self.miss}, '''
            f'''capacity={self.capacity}, current size={self.num_keys})'''
        )






      def __contains__(    self     ,     a_	)	->bool:







        '''simple docstring'''


        return key in self.cache






      def 			lowerCamelCase__    (    self     ,     a_	)	->U | None:







        '''simple docstring'''


        if key in self.cache:
          self.hits += 1
          _a							=	self.cache[key]
          _a							=	self.list.remove(self.cache[key]	)
          assert node == value_node

          # node is guaranteed not None because it is in self.cache
          assert node is not None
          self.list.add(a_	)
          return node.val
        self.miss += 1
        return None






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_	)	->None:







        '''simple docstring'''


        if key not in self.cache:
          if self.num_keys >= self.capacity:
            # delete first node (oldest) when over capacity
            _a							=	self.list.head.next

            # guaranteed to have a non-None first node when num_keys > 0
            # explain to type checker via assertions
            assert first_node is not None
            assert first_node.key is not None
            assert (
                self.list.remove(a_	) is not None
            )  # node guaranteed to be in list assert node.key is not None

            del self.cache[first_node.key]
            self.num_keys -= 1
          _a							=	DoubleLinkedListNode(a_     ,     a_	)
          self.list.add(self.cache[key]	)
          self.num_keys += 1

        else:
          # bump node to the end of the list, update value
          _a							=	self.list.remove(self.cache[key]	)
          assert node is not None  # node guaranteed to be in list
          _a							=	value
          self.list.add(a_	)






      @classmethod
      def 			lowerCamelCase__    (    cls     ,     a_ = 1_2_8	)	->Callable[[Callable[[T], U]], Callable[..., U]]:







        '''simple docstring'''


        def cache_decorator_inner(a_	) -> Callable[..., U]:
          def cache_decorator_wrapper(*a_	) -> U:
            if func not in cls.decorator_function_to_instance_map:
              _a							=	LRUCache(a_	)

            _a							=	cls.decorator_function_to_instance_map[func].get(args[0]	)
            if result is None:
              _a							=	func(*a_	)
              cls.decorator_function_to_instance_map[func].put(args[0]     ,     a_	)
            return result

          def cache_info() -> LRUCache[T, U]:
            return cls.decorator_function_to_instance_map[func]

          setattr(a_     ,     "cache_info"     ,     a_	)  # noqa: B010

          return cache_decorator_wrapper

        return cache_decorator_inner


if __name__ == "__main__":
     import doctest

     doctest.testmod()


