





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_  (  snake_case__  ,     unittest.TestCase							):
					_a      :						Dict						    =					'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'



					def   __a     (		self     :     Union[str, Any]       ,      lowerCamelCase     :     List[Any]=0						):
							lowerCamelCase_	:     List[str] 						=      np.random.RandomState(lowerCamelCase						)
							lowerCamelCase_	:     Tuple 						=      {
							    '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   __a     (		self     :     Optional[Any]						):
							lowerCamelCase_	:     str 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							lowerCamelCase_	:     Tuple 						=      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   __a     (		self     :     Union[str, Any]						):
							lowerCamelCase_	:     Union[str, Any] 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							lowerCamelCase_	:     Any 						=      PNDMScheduler.from_config(pipe.scheduler.config       ,      skip_prk_steps=lowerCamelCase						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							lowerCamelCase_	:     Dict 						=      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   __a     (		self     :     str						):
							lowerCamelCase_	:     Dict 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							lowerCamelCase_	:     int 						=      LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							lowerCamelCase_	:     int 						=      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   __a     (		self     :     List[Any]						):
							lowerCamelCase_	:     List[Any] 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							lowerCamelCase_	:     Tuple 						=      EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							lowerCamelCase_	:     Optional[Any] 						=      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   __a     (		self     :     Any						):
							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=lowerCamelCase						)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							lowerCamelCase_	:     Union[str, Any] 						=      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   __a     (		self     :     Tuple						):
							lowerCamelCase_	:     str 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							lowerCamelCase_	:     Dict 						=      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							lowerCamelCase_	:     List[str] 						=      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   __a     (		self     :     List[str]						):
							lowerCamelCase_	:     List[Any] 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

							lowerCamelCase_	:     str 						=      self.get_dummy_inputs()
							lowerCamelCase_	:     str 						=      3 * [inputs['prompt']]

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

							lowerCamelCase_	:     Optional[int] 						=      self.get_dummy_inputs()
							lowerCamelCase_	:     int 						=      3 * [inputs.pop('prompt'						)]

							lowerCamelCase_	:     Any 						=      pipe.tokenizer(
							    lowerCamelCase       ,      padding='max_length'       ,      max_length=pipe.tokenizer.model_max_length       ,      truncation=lowerCamelCase       ,      return_tensors='np'       ,      )
							lowerCamelCase_	:     Tuple 						=      text_inputs['input_ids']

							lowerCamelCase_	:     str 						=      pipe.text_encoder(input_ids=text_inputs.astype(np.intaa						)						)[0]

							lowerCamelCase_	:     Optional[Any] 						=      prompt_embeds

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

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





					def   __a     (		self     :     Any						):
							lowerCamelCase_	:     str 						=      OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,      provider='CPUExecutionProvider'						)
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

							lowerCamelCase_	:     Dict 						=      self.get_dummy_inputs()
							lowerCamelCase_	:     List[str] 						=      3 * ['this is a negative prompt']
							lowerCamelCase_	:     Union[str, Any] 						=      negative_prompt
							lowerCamelCase_	:     int 						=      3 * [inputs['prompt']]

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

							lowerCamelCase_	:     Dict 						=      self.get_dummy_inputs()
							lowerCamelCase_	:     Optional[Any] 						=      3 * [inputs.pop('prompt'						)]

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

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

							lowerCamelCase_						,      lowerCamelCase_	:     Dict 						=      embeds

							# forward
							lowerCamelCase_	:     Any 						=      pipe(**lowerCamelCase						)
							lowerCamelCase_	:     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							):



					@property
					def   __a     (		self     :     List[Any]						):
							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)



					@property
					def   __a     (		self     :     Optional[int]						):
							lowerCamelCase_	:     List[Any] 						=      ort.SessionOptions()
							lowerCamelCase_	:     Dict 						=      False
							return options



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

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

							lowerCamelCase_	:     int 						=      image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							lowerCamelCase_	:     Dict 						=      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   __a     (		self     :     Any						):
							lowerCamelCase_	:     Optional[Any] 						=      DDIMScheduler.from_pretrained(
							    'runwayml/stable-diffusion-v1-5'       ,      subfolder='scheduler'       ,      revision='onnx'						)
							lowerCamelCase_	:     List[Any] 						=      OnnxStableDiffusionPipeline.from_pretrained(
							    'runwayml/stable-diffusion-v1-5'       ,      revision='onnx'       ,      scheduler=lowerCamelCase       ,      safety_checker=lowerCamelCase       ,      feature_extractor=lowerCamelCase       ,      provider=self.gpu_provider       ,      sess_options=self.gpu_options       ,      )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 5_12, 5_12, 3)
							lowerCamelCase_	:     Optional[int] 						=      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   __a     (		self     :     Dict						):
							lowerCamelCase_	:     Optional[int] 						=      LMSDiscreteScheduler.from_pretrained(
							    'runwayml/stable-diffusion-v1-5'       ,      subfolder='scheduler'       ,      revision='onnx'						)
							lowerCamelCase_	:     str 						=      OnnxStableDiffusionPipeline.from_pretrained(
							    'runwayml/stable-diffusion-v1-5'       ,      revision='onnx'       ,      scheduler=lowerCamelCase       ,      safety_checker=lowerCamelCase       ,      feature_extractor=lowerCamelCase       ,      provider=self.gpu_provider       ,      sess_options=self.gpu_options       ,      )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

							assert image.shape == (1, 5_12, 5_12, 3)
							lowerCamelCase_	:     Optional[Any] 						=      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   __a     (		self     :     str						):
							lowerCamelCase_	:     Optional[int] 						=      0

							def test_callback_fn(lowerCamelCase     :     int       ,      lowerCamelCase     :     int       ,      lowerCamelCase     :     np.ndarray						) -> None:
									lowerCamelCase_	:     Dict 						=      True
									nonlocal number_of_steps
									number_of_steps += 1
									if step == 0:
											assert latents.shape == (1, 4, 64, 64)
											lowerCamelCase_	:     str 						=      latents[0, -3:, -3:, -1]
											lowerCamelCase_	:     List[Any] 						=      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, 64, 64)
											lowerCamelCase_	:     List[str] 						=      latents[0, -3:, -3:, -1]
											lowerCamelCase_	:     Dict 						=      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

							lowerCamelCase_	:     int 						=      False

							lowerCamelCase_	:     Optional[Any] 						=      OnnxStableDiffusionPipeline.from_pretrained(
							    'runwayml/stable-diffusion-v1-5'       ,      revision='onnx'       ,      safety_checker=lowerCamelCase       ,      feature_extractor=lowerCamelCase       ,      provider=self.gpu_provider       ,      sess_options=self.gpu_options       ,      )
							pipe.set_progress_bar_config(disable=lowerCamelCase						)

							lowerCamelCase_	:     int 						=      'Andromeda galaxy in a bottle'

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





					def   __a     (		self     :     Tuple						):
							lowerCamelCase_	:     Any 						=      OnnxStableDiffusionPipeline.from_pretrained(
							    'runwayml/stable-diffusion-v1-5'       ,      revision='onnx'       ,      safety_checker=lowerCamelCase       ,      feature_extractor=lowerCamelCase       ,      provider=self.gpu_provider       ,      sess_options=self.gpu_options       ,      )
							assert isinstance(lowerCamelCase       ,      lowerCamelCase						)
							assert pipe.safety_checker is None

							lowerCamelCase_	:     Any 						=      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(lowerCamelCase						)
									lowerCamelCase_	:     Union[str, Any] 						=      OnnxStableDiffusionPipeline.from_pretrained(lowerCamelCase						)

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







_lowercase		:  Any			       ={"""a""": ["""c""", """b"""], """b""": ["""d""", """e"""], """c""": [], """d""": [], """e""": []}
_lowercase		:  Union[str, Any]			       =["""a""", """b""", """c""", """d""", """e"""]


def 		_SCREAMING_SNAKE_CASE     ( lowerCAmelCase__							,lowerCAmelCase__							,lowerCAmelCase__						):
		lowerCamelCase_	:     int 						=      start
		# add current to visited
		visited.append(lowerCAmelCase__						)
		lowerCamelCase_	:     Optional[int] 						=      edges[current]
		for neighbor in neighbors:
				# if neighbor not in visited, visit
				if neighbor not in visited:
						lowerCamelCase_	:     Dict 						=      topological_sort(lowerCAmelCase__							,lowerCAmelCase__							,lowerCAmelCase__						)
    # if all neighbors visited add current to sort
		sort.append(lowerCAmelCase__						)
		# if all vertices haven't been visited select a new one to visit
		if len(lowerCAmelCase__						) != len(lowerCAmelCase__						):
				for vertice in vertices:
						if vertice not in visited:
								lowerCamelCase_	:     Any 						=      topological_sort(lowerCAmelCase__							,lowerCAmelCase__							,lowerCAmelCase__						)
    # return sort
		return sort


if __name__ == "__main__":
		_lowercase		:  str			       =topological_sort("""a""", [], [])
		print(sort)
