




import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
       import onnxruntime as ort




class 		UpperCamelCase__	(lowerCAmelCase__    ,				unittest.TestCase     ):

    '''simple docstring'''



    lowerCamelCase_      :	Union[str, Any]	     =      """hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""







    def     _lowercase     (    self     ,  UpperCamelCase__=0     )     ->      Dict:
           lowerCamelCase   :      List[Any]			 =					np.random.RandomState(UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					{
               "prompt": "A painting of a squirrel eating a burger",
               "generator": generator,
               "num_inference_steps": 2,
               "guidance_scale": 7.5,
               "output_type": "numpy",
           }
           return inputs







    def     _lowercase     (    self     )     ->      Union[str, Any]:
           lowerCamelCase   :      List[str]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

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

           assert image.shape == (1, 128, 128, 3)
           lowerCamelCase   :      str			 =					np.array([0.65072, 0.58492, 0.48219, 0.55521, 0.53180, 0.55939, 0.50697, 0.39800, 0.46455]     )

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







    def     _lowercase     (    self     )     ->      Optional[int]:
           lowerCamelCase   :      Union[str, Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           lowerCamelCase   :      Any			 =					PNDMScheduler.from_config(pipe.scheduler.config     ,  skip_prk_steps=UpperCamelCase__     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					self.get_dummy_inputs()
           lowerCamelCase   :      int			 =					pipe(**UpperCamelCase__     ).images
           lowerCamelCase   :      Dict			 =					image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           lowerCamelCase   :      Any			 =					np.array([0.65863, 0.59425, 0.49326, 0.56313, 0.53875, 0.56627, 0.51065, 0.39777, 0.46330]     )

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







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase   :      Optional[Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           lowerCamelCase   :      Union[str, Any]			 =					LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Dict			 =					self.get_dummy_inputs()
           lowerCamelCase   :      Any			 =					pipe(**UpperCamelCase__     ).images
           lowerCamelCase   :      Optional[int]			 =					image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           lowerCamelCase   :      List[Any]			 =					np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279]     )

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







    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase   :      Union[str, Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           lowerCamelCase   :      List[str]			 =					EulerDiscreteScheduler.from_config(pipe.scheduler.config     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

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

           assert image.shape == (1, 128, 128, 3)
           lowerCamelCase   :      Optional[Any]			 =					np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279]     )

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







    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase   :      Optional[Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           lowerCamelCase   :      Union[str, Any]			 =					EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

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

           assert image.shape == (1, 128, 128, 3)
           lowerCamelCase   :      Optional[Any]			 =					np.array([0.53817, 0.60812, 0.47384, 0.49530, 0.51894, 0.49814, 0.47984, 0.38958, 0.44271]     )

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







    def     _lowercase     (    self     )     ->      Union[str, Any]:
           lowerCamelCase   :      Dict			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           lowerCamelCase   :      List[str]			 =					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      int			 =					self.get_dummy_inputs()
           lowerCamelCase   :      Union[str, Any]			 =					pipe(**UpperCamelCase__     ).images
           lowerCamelCase   :      int			 =					image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           lowerCamelCase   :      str			 =					np.array([0.53895, 0.60808, 0.47933, 0.49608, 0.51886, 0.49950, 0.48053, 0.38957, 0.44200]     )

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







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase   :      List[Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					self.get_dummy_inputs()
           lowerCamelCase   :      Optional[int]			 =					3 * [inputs["prompt"]]

           # forward
           lowerCamelCase   :      str			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					output.images[0, -3:, -3:, -1]

           lowerCamelCase   :      Optional[Any]			 =					self.get_dummy_inputs()
           lowerCamelCase   :      Optional[Any]			 =					3 * [inputs.pop("prompt"     )]

           lowerCamelCase   :      Any			 =					pipe.tokenizer(
               UpperCamelCase__     ,  padding="max_length"     ,  max_length=pipe.tokenizer.model_max_length     ,  truncation=UpperCamelCase__     ,  return_tensors="np"     ,  )
           lowerCamelCase   :      Optional[int]			 =					text_inputs["input_ids"]

           lowerCamelCase   :      Dict			 =					pipe.text_encoder(input_ids=text_inputs.astype(np.intaa     )     )[0]

           lowerCamelCase   :      str			 =					prompt_embeds

           # forward
           lowerCamelCase   :      Tuple			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      Tuple			 =					output.images[0, -3:, -3:, -1]

           assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()     ).max() < 1e-4







    def     _lowercase     (    self     )     ->      Tuple:
           lowerCamelCase   :      List[Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,  provider="CPUExecutionProvider"     )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					self.get_dummy_inputs()
           lowerCamelCase   :      Any			 =					3 * ["this is a negative prompt"]
           lowerCamelCase   :      str			 =					negative_prompt
           lowerCamelCase   :      Tuple			 =					3 * [inputs["prompt"]]

           # forward
           lowerCamelCase   :      Dict			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      int			 =					output.images[0, -3:, -3:, -1]

           lowerCamelCase   :      List[Any]			 =					self.get_dummy_inputs()
           lowerCamelCase   :      Tuple			 =					3 * [inputs.pop("prompt"     )]

           lowerCamelCase   :      Tuple			 =					[]
           for p in [prompt, negative_prompt]:
                  lowerCamelCase   :      int			 =					pipe.tokenizer(
                      UpperCamelCase__     ,  padding="max_length"     ,  max_length=pipe.tokenizer.model_max_length     ,  truncation=UpperCamelCase__     ,  return_tensors="np"     ,  )
                  lowerCamelCase   :      Optional[int]			 =					text_inputs["input_ids"]

                  embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa     )     )[0]     )

           lowerCamelCase ,      lowerCamelCase   :      List[str]			 =					embeds

           # forward
           lowerCamelCase   :      Optional[int]			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					output.images[0, -3:, -3:, -1]

           assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()     ).max() < 1e-4




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

    '''simple docstring'''



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







    @property
    def     _lowercase     (    self     )     ->      Optional[int]:
           lowerCamelCase   :      str			 =					ort.SessionOptions()
           lowerCamelCase   :      str			 =					False
           return options







    def     _lowercase     (    self     )     ->      Optional[Any]:
           # using the PNDM scheduler by default
           lowerCamelCase   :      Optional[int]			 =					OnnxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"     ,  revision="onnx"     ,  safety_checker=UpperCamelCase__     ,  feature_extractor=UpperCamelCase__     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					"A painting of a squirrel eating a burger"
           np.random.seed(0     )
           lowerCamelCase   :      Tuple			 =					sd_pipe([prompt]     ,  guidance_scale=6.0     ,  num_inference_steps=10     ,  output_type="np"     )
           lowerCamelCase   :      str			 =					output.images

           lowerCamelCase   :      Union[str, Any]			 =					image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           lowerCamelCase   :      Tuple			 =					np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720]     )

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







    def     _lowercase     (    self     )     ->      Optional[int]:
           lowerCamelCase   :      List[Any]			 =					DDIMScheduler.from_pretrained(
               "runwayml/stable-diffusion-v1-5"     ,  subfolder="scheduler"     ,  revision="onnx"     )
           lowerCamelCase   :      Optional[Any]			 =					OnnxStableDiffusionPipeline.from_pretrained(
               "runwayml/stable-diffusion-v1-5"     ,  revision="onnx"     ,  scheduler=UpperCamelCase__     ,  safety_checker=UpperCamelCase__     ,  feature_extractor=UpperCamelCase__     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Union[str, Any]			 =					"open neural network exchange"
           lowerCamelCase   :      Tuple			 =					np.random.RandomState(0     )
           lowerCamelCase   :      str			 =					sd_pipe([prompt]     ,  guidance_scale=7.5     ,  num_inference_steps=10     ,  generator=UpperCamelCase__     ,  output_type="np"     )
           lowerCamelCase   :      Tuple			 =					output.images
           lowerCamelCase   :      Tuple			 =					image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           lowerCamelCase   :      Dict			 =					np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486]     )

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







    def     _lowercase     (    self     )     ->      Any:
           lowerCamelCase   :      Any			 =					LMSDiscreteScheduler.from_pretrained(
               "runwayml/stable-diffusion-v1-5"     ,  subfolder="scheduler"     ,  revision="onnx"     )
           lowerCamelCase   :      Any			 =					OnnxStableDiffusionPipeline.from_pretrained(
               "runwayml/stable-diffusion-v1-5"     ,  revision="onnx"     ,  scheduler=UpperCamelCase__     ,  safety_checker=UpperCamelCase__     ,  feature_extractor=UpperCamelCase__     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
           sd_pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      int			 =					"open neural network exchange"
           lowerCamelCase   :      Optional[Any]			 =					np.random.RandomState(0     )
           lowerCamelCase   :      int			 =					sd_pipe([prompt]     ,  guidance_scale=7.5     ,  num_inference_steps=10     ,  generator=UpperCamelCase__     ,  output_type="np"     )
           lowerCamelCase   :      Optional[Any]			 =					output.images
           lowerCamelCase   :      Tuple			 =					image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           lowerCamelCase   :      List[Any]			 =					np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161]     )

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







    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase   :      List[str]			 =					0

           def test_callback_fn(UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ) -> None:
                  lowerCamelCase   :      Dict			 =					True
                  nonlocal number_of_steps
                  number_of_steps += 1
                  if step == 0:
                         assert latents.shape == (1, 4, 64, 64)
                         lowerCamelCase   :      int			 =					latents[0, -3:, -3:, -1]
                         lowerCamelCase   :      Optional[int]			 =					np.array(
                             [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167]     )

                         assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 1e-3
                  elif step == 5:
                         assert latents.shape == (1, 4, 64, 64)
                         lowerCamelCase   :      Union[str, Any]			 =					latents[0, -3:, -3:, -1]
                         lowerCamelCase   :      List[str]			 =					np.array(
                             [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875]     )

                         assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 1e-3

           lowerCamelCase   :      List[str]			 =					False

           lowerCamelCase   :      str			 =					OnnxStableDiffusionPipeline.from_pretrained(
               "runwayml/stable-diffusion-v1-5"     ,  revision="onnx"     ,  safety_checker=UpperCamelCase__     ,  feature_extractor=UpperCamelCase__     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Union[str, Any]			 =					"Andromeda galaxy in a bottle"

           lowerCamelCase   :      List[str]			 =					np.random.RandomState(0     )
           pipe(
               prompt=UpperCamelCase__     ,  num_inference_steps=5     ,  guidance_scale=7.5     ,  generator=UpperCamelCase__     ,  callback=UpperCamelCase__     ,  callback_steps=1     ,  )
           assert test_callback_fn.has_been_called
           assert number_of_steps == 6







    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      int			 =					OnnxStableDiffusionPipeline.from_pretrained(
               "runwayml/stable-diffusion-v1-5"     ,  revision="onnx"     ,  safety_checker=UpperCamelCase__     ,  feature_extractor=UpperCamelCase__     ,  provider=self.gpu_provider     ,  sess_options=self.gpu_options     ,  )
           assert isinstance(UpperCamelCase__     ,  UpperCamelCase__     )
           assert pipe.safety_checker is None

           lowerCamelCase   :      Optional[int]			 =					pipe("example prompt"     ,  num_inference_steps=2     ).images[0]
           assert image is not None

           # check that there's no error when saving a pipeline with one of the models being None
           with tempfile.TemporaryDirectory() as tmpdirname:
                  pipe.save_pretrained(UpperCamelCase__     )
                  lowerCamelCase   :      int			 =					OnnxStableDiffusionPipeline.from_pretrained(UpperCamelCase__     )

           # sanity check that the pipeline still works
           assert pipe.safety_checker is None
           lowerCamelCase   :      Tuple			 =					pipe("example prompt"     ,  num_inference_steps=2     ).images[0]
           assert image is not None







from functools import lru_cache




@lru_cache
def       _snake_case	(   lowerCAmelCase							:   int	):
		"""simple docstring"""







		if num < 0:
				raise ValueError("Number should not be negative."	)

		return 1 if num in (0, 1) else num * factorial(num - 1	)


if __name__ == "__main__":
		import doctest

		doctest.testmod()
