



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 	lowerCamelCase      (						a__,     unittest.TestCase  ):







					'''simple docstring'''







					lowerCAmelCase_				:    str									=  "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"



					def    A__		(			self    ,	lowerCAmelCase=0     ):
							UpperCAmelCase_						=				np.random.RandomState(lowerCamelCase_     )
							UpperCAmelCase_						=				{
							    "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     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

							assert image.shape == (1, 128, 128, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							UpperCAmelCase_						=				PNDMScheduler.from_config(pipe.scheduler.config    ,	skip_prk_steps=lowerCamelCase_     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

							assert image.shape == (1, 128, 128, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							UpperCAmelCase_						=				LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

							assert image.shape == (1, 128, 128, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							UpperCAmelCase_						=				EulerDiscreteScheduler.from_config(pipe.scheduler.config     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

							assert image.shape == (1, 128, 128, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							UpperCAmelCase_						=				EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

							assert image.shape == (1, 128, 128, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							UpperCAmelCase_						=				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

							assert image.shape == (1, 128, 128, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

							UpperCAmelCase_						=				self.get_dummy_inputs()
							UpperCAmelCase_						=				3 * [inputs["prompt"]]

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

							UpperCAmelCase_						=				self.get_dummy_inputs()
							UpperCAmelCase_						=				3 * [inputs.pop("prompt"     )]

							UpperCAmelCase_						=				pipe.tokenizer(
							    lowerCamelCase_    ,	padding="max_length"    ,	max_length=pipe.tokenizer.model_max_length    ,	truncation=lowerCamelCase_    ,	return_tensors="np"    ,	)
							UpperCAmelCase_						=				text_inputs["input_ids"]

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

							UpperCAmelCase_						=				prompt_embeds

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

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



					def    A__		(			self     ):
							UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,	provider="CPUExecutionProvider"     )
							pipe.set_progress_bar_config(disable=lowerCamelCase_     )

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

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

							UpperCAmelCase_						=				self.get_dummy_inputs()
							UpperCAmelCase_						=				3 * [inputs.pop("prompt"     )]

							UpperCAmelCase_						=				[]
							for p in [prompt, negative_prompt]:
									UpperCAmelCase_						=				pipe.tokenizer(
									    lowerCamelCase_    ,	padding="max_length"    ,	max_length=pipe.tokenizer.model_max_length    ,	truncation=lowerCamelCase_    ,	return_tensors="np"    ,	)
									UpperCAmelCase_						=				text_inputs["input_ids"]

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

							UpperCAmelCase_					,						UpperCAmelCase_						=				embeds

							# forward
							UpperCAmelCase_						=				pipe(**lowerCamelCase_     )
							UpperCAmelCase_						=				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 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







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



					@property
					def    A__		(			self     ):
							UpperCAmelCase_						=				ort.SessionOptions()
							UpperCAmelCase_						=				False
							return options



					def    A__		(			self     ):
							# using the PNDM scheduler by default
							UpperCAmelCase_						=				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_     )

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

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

							assert image.shape == (1, 512, 512, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				DDIMScheduler.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"    ,	subfolder="scheduler"    ,	revision="onnx"     )
							UpperCAmelCase_						=				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_     )

							UpperCAmelCase_						=				"open neural network exchange"
							UpperCAmelCase_						=				np.random.RandomState(0     )
							UpperCAmelCase_						=				sd_pipe([prompt]    ,	guidance_scale=7.5    ,	num_inference_steps=10    ,	generator=lowerCamelCase_    ,	output_type="np"     )
							UpperCAmelCase_						=				output.images
							UpperCAmelCase_						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				LMSDiscreteScheduler.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"    ,	subfolder="scheduler"    ,	revision="onnx"     )
							UpperCAmelCase_						=				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_     )

							UpperCAmelCase_						=				"open neural network exchange"
							UpperCAmelCase_						=				np.random.RandomState(0     )
							UpperCAmelCase_						=				sd_pipe([prompt]    ,	guidance_scale=7.5    ,	num_inference_steps=10    ,	generator=lowerCamelCase_    ,	output_type="np"     )
							UpperCAmelCase_						=				output.images
							UpperCAmelCase_						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				0

							def test_callback_fn(lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ) -> None:
									UpperCAmelCase_						=				True
									nonlocal number_of_steps
									number_of_steps += 1
									if step == 0:
											assert latents.shape == (1, 4, 64, 64)
											UpperCAmelCase_						=				latents[0, -3:, -3:, -1]
											UpperCAmelCase_						=				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)
											UpperCAmelCase_						=				latents[0, -3:, -3:, -1]
											UpperCAmelCase_						=				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

							UpperCAmelCase_						=				False

							UpperCAmelCase_						=				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_     )

							UpperCAmelCase_						=				"Andromeda galaxy in a bottle"

							UpperCAmelCase_						=				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     ):
							UpperCAmelCase_						=				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

							UpperCAmelCase_						=				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_     )
									UpperCAmelCase_						=				OnnxStableDiffusionPipeline.from_pretrained(lowerCamelCase_     )

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






def snake_case__  (			__SCREAMING_SNAKE_CASE )		->							int:
		UpperCAmelCase_						=				[1]

		UpperCAmelCase_					,						UpperCAmelCase_					,						UpperCAmelCase_						=				0, 0, 0
		UpperCAmelCase_						=				ugly_nums[ia] * 2
		UpperCAmelCase_						=				ugly_nums[ia] * 3
		UpperCAmelCase_						=				ugly_nums[ia] * 5

		for _ in range(1			,     __SCREAMING_SNAKE_CASE ):
				UpperCAmelCase_						=				min(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE )
				ugly_nums.append(__SCREAMING_SNAKE_CASE )
				if next_num == next_a:
						ia += 1
						UpperCAmelCase_						=				ugly_nums[ia] * 2
				if next_num == next_a:
						ia += 1
						UpperCAmelCase_						=				ugly_nums[ia] * 3
				if next_num == next_a:
						ia += 1
						UpperCAmelCase_						=				ugly_nums[ia] * 5
		return ugly_nums[-1]


if __name__ == "__main__":
			from doctest import testmod

			testmod(verbose=True)
			print(f'''{ugly_numbers(200) = }''')

