

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 (     SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         ="hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"

    def UpperCamelCase			(			self :      Union[str, Any]							,    snake_case__ :      Optional[Any]=0    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			np.random.RandomState(snake_case__    )
     SCREAMING_SNAKE_CASE								=			{
         '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 UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     pipe.set_progress_bar_config(disable=snake_case__    )

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

     assert image.shape == (1, 1_2_8, 1_2_8, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.65_072, 0.58_492, 0.48_219, 0.55_521, 0.53_180, 0.55_939, 0.50_697, 0.39_800, 0.46_455]    )

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

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     SCREAMING_SNAKE_CASE								=			PNDMScheduler.from_config(pipe.scheduler.config							,    skip_prk_steps=snake_case__    )
     pipe.set_progress_bar_config(disable=snake_case__    )

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

     assert image.shape == (1, 1_2_8, 1_2_8, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.65_863, 0.59_425, 0.49_326, 0.56_313, 0.53_875, 0.56_627, 0.51_065, 0.39_777, 0.46_330]    )

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

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     SCREAMING_SNAKE_CASE								=			LMSDiscreteScheduler.from_config(pipe.scheduler.config    )
     pipe.set_progress_bar_config(disable=snake_case__    )

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

     assert image.shape == (1, 1_2_8, 1_2_8, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]    )

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

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     SCREAMING_SNAKE_CASE								=			EulerDiscreteScheduler.from_config(pipe.scheduler.config    )
     pipe.set_progress_bar_config(disable=snake_case__    )

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

     assert image.shape == (1, 1_2_8, 1_2_8, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]    )

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

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     SCREAMING_SNAKE_CASE								=			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config    )
     pipe.set_progress_bar_config(disable=snake_case__    )

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

     assert image.shape == (1, 1_2_8, 1_2_8, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.53_817, 0.60_812, 0.47_384, 0.49_530, 0.51_894, 0.49_814, 0.47_984, 0.38_958, 0.44_271]    )

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

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     SCREAMING_SNAKE_CASE								=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
     pipe.set_progress_bar_config(disable=snake_case__    )

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

     assert image.shape == (1, 1_2_8, 1_2_8, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.53_895, 0.60_808, 0.47_933, 0.49_608, 0.51_886, 0.49_950, 0.48_053, 0.38_957, 0.44_200]    )

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

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE								=			3 * [inputs['prompt']]

     # forward
     SCREAMING_SNAKE_CASE								=			pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.images[0, -3:, -3:, -1]

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE								=			3 * [inputs.pop('prompt'    )]

     SCREAMING_SNAKE_CASE								=			pipe.tokenizer(
         snake_case__							,    padding='max_length'							,    max_length=pipe.tokenizer.model_max_length							,    truncation=snake_case__							,    return_tensors='np'							,    )
     SCREAMING_SNAKE_CASE								=			text_inputs['input_ids']

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

     SCREAMING_SNAKE_CASE								=			prompt_embeds

     # forward
     SCREAMING_SNAKE_CASE								=			pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.images[0, -3:, -3:, -1]

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




    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider='CPUExecutionProvider'    )
     pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE								=			3 * ['this is a negative prompt']
     SCREAMING_SNAKE_CASE								=			negative_prompt
     SCREAMING_SNAKE_CASE								=			3 * [inputs['prompt']]

     # forward
     SCREAMING_SNAKE_CASE								=			pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.images[0, -3:, -3:, -1]

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs()
     SCREAMING_SNAKE_CASE								=			3 * [inputs.pop('prompt'    )]

     SCREAMING_SNAKE_CASE								=			[]
     for p in [prompt, negative_prompt]:
      SCREAMING_SNAKE_CASE								=			pipe.tokenizer(
          snake_case__							,    padding='max_length'							,    max_length=pipe.tokenizer.model_max_length							,    truncation=snake_case__							,    return_tensors='np'							,    )
      SCREAMING_SNAKE_CASE								=			text_inputs['input_ids']

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

     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			embeds

     # forward
     SCREAMING_SNAKE_CASE								=			pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			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   ):

    @property
    def UpperCamelCase			(			self :      Union[str, Any]    ):

     """simple docstring"""






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

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

     """simple docstring"""






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

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'							,    revision='onnx'							,    safety_checker=snake_case__							,    feature_extractor=snake_case__							,    provider=self.gpu_provider							,    sess_options=self.gpu_options							,    )
     sd_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			'A painting of a squirrel eating a burger'
     np.random.seed(0    )
     SCREAMING_SNAKE_CASE								=			sd_pipe([prompt]							,    guidance_scale=6.0							,    num_inference_steps=1_0							,    output_type='np'    )
     SCREAMING_SNAKE_CASE								=			output.images

     SCREAMING_SNAKE_CASE								=			image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.0_452, 0.0_390, 0.0_087, 0.0_350, 0.0_617, 0.0_364, 0.0_544, 0.0_523, 0.0_720]    )

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

    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			DDIMScheduler.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,    subfolder='scheduler'							,    revision='onnx'    )
     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,    revision='onnx'							,    scheduler=snake_case__							,    safety_checker=snake_case__							,    feature_extractor=snake_case__							,    provider=self.gpu_provider							,    sess_options=self.gpu_options							,    )
     sd_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			'open neural network exchange'
     SCREAMING_SNAKE_CASE								=			np.random.RandomState(0    )
     SCREAMING_SNAKE_CASE								=			sd_pipe([prompt]							,    guidance_scale=7.5							,    num_inference_steps=1_0							,    generator=snake_case__							,    output_type='np'    )
     SCREAMING_SNAKE_CASE								=			output.images
     SCREAMING_SNAKE_CASE								=			image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.2_867, 0.1_974, 0.1_481, 0.7_294, 0.7_251, 0.6_667, 0.4_194, 0.5_642, 0.6_486]    )

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

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			LMSDiscreteScheduler.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,    subfolder='scheduler'							,    revision='onnx'    )
     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,    revision='onnx'							,    scheduler=snake_case__							,    safety_checker=snake_case__							,    feature_extractor=snake_case__							,    provider=self.gpu_provider							,    sess_options=self.gpu_options							,    )
     sd_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			'open neural network exchange'
     SCREAMING_SNAKE_CASE								=			np.random.RandomState(0    )
     SCREAMING_SNAKE_CASE								=			sd_pipe([prompt]							,    guidance_scale=7.5							,    num_inference_steps=1_0							,    generator=snake_case__							,    output_type='np'    )
     SCREAMING_SNAKE_CASE								=			output.images
     SCREAMING_SNAKE_CASE								=			image[0, -3:, -3:, -1]

     assert image.shape == (1, 5_1_2, 5_1_2, 3)
     SCREAMING_SNAKE_CASE								=			np.array([0.2_306, 0.1_959, 0.1_593, 0.6_549, 0.6_394, 0.5_408, 0.5_065, 0.6_010, 0.6_161]    )

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

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			0

     def test_callback_fn(snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      np.ndarray    ) -> None:
      SCREAMING_SNAKE_CASE								=			True
      nonlocal number_of_steps
      number_of_steps += 1
      if step == 0:
       assert latents.shape == (1, 4, 6_4, 6_4)
       SCREAMING_SNAKE_CASE								=			latents[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE								=			np.array(
           [-0.6_772, -0.3_835, -1.2_456, 0.1_905, -1.0_974, 0.6_967, -1.9_353, 0.0_178, 1.0_167]    )

       assert np.abs(latents_slice.flatten() - expected_slice    ).max() < 1E-3
      elif step == 5:
       assert latents.shape == (1, 4, 6_4, 6_4)
       SCREAMING_SNAKE_CASE								=			latents[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE								=			np.array(
           [-0.3_351, 0.2_241, -0.1_837, -0.2_325, -0.6_577, 0.3_393, -0.0_241, 0.5_899, 1.3_875]    )

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

     SCREAMING_SNAKE_CASE								=			False

     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,    revision='onnx'							,    safety_checker=snake_case__							,    feature_extractor=snake_case__							,    provider=self.gpu_provider							,    sess_options=self.gpu_options							,    )
     pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			'Andromeda galaxy in a bottle'

     SCREAMING_SNAKE_CASE								=			np.random.RandomState(0    )
     pipe(
         prompt=snake_case__							,    num_inference_steps=5							,    guidance_scale=7.5							,    generator=snake_case__							,    callback=snake_case__							,    callback_steps=1							,    )
     assert test_callback_fn.has_been_called
     assert number_of_steps == 6




    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(
         'runwayml/stable-diffusion-v1-5'							,    revision='onnx'							,    safety_checker=snake_case__							,    feature_extractor=snake_case__							,    provider=self.gpu_provider							,    sess_options=self.gpu_options							,    )
     assert isinstance(snake_case__							,    snake_case__    )
     assert pipe.safety_checker is None

     SCREAMING_SNAKE_CASE								=			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(snake_case__    )
      SCREAMING_SNAKE_CASE								=			OnnxStableDiffusionPipeline.from_pretrained(snake_case__    )

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




import random


def        __lowerCAmelCase    (       _UpperCamelCase		:  int		,							_UpperCamelCase		:  float		,							_UpperCamelCase		:  bool = False	)		->							dict:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			{i: [] for i in range(_UpperCamelCase	)}

 # if probability is greater or equal than 1, then generate a complete graph
 if probability >= 1:
  return complete_graph(_UpperCamelCase	)
 # if probability is lower or equal than 0, then return a graph without edges
 if probability <= 0:
  return graph

 # for each couple of nodes, add an edge from u to v
 # if the number randomly generated is greater than probability probability
 for i in range(_UpperCamelCase	):
  for j in range(i + 1		,							_UpperCamelCase	):
   if random.random() < probability:
    graph[i].append(_UpperCamelCase	)
    if not directed:
     # if the graph is undirected, add an edge in from j to i, either
     graph[j].append(_UpperCamelCase	)
 return graph


def        __lowerCAmelCase    (       _UpperCamelCase		:  int	)		->							dict:
 '''simple docstring'''


 return {
     i: [j for j in range(_UpperCamelCase	) if i != j] for i in range(_UpperCamelCase	)
 }


if __name__ == "__main__":
      import doctest

      doctest.testmod()

