





"""simple docstring"""





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


						def     snake_case_   (self   ,							__a=0		) ->   str:
								UpperCamelCase   =	np.random.RandomState(_a		)
								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     snake_case_   (self		) ->   Optional[Any]:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 1_28, 1_28, 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     snake_case_   (self		) ->   str:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								UpperCamelCase   =	PNDMScheduler.from_config(pipe.scheduler.config   ,							skip_prk_steps=_a		)
								pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 1_28, 1_28, 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     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								UpperCamelCase   =	LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 1_28, 1_28, 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     snake_case_   (self		) ->   List[Any]:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								UpperCamelCase   =	EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 1_28, 1_28, 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     snake_case_   (self		) ->   List[Any]:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								UpperCamelCase   =	EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 1_28, 1_28, 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     snake_case_   (self		) ->   int:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								UpperCamelCase   =	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 1_28, 1_28, 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     snake_case_   (self		) ->   int:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								pipe.set_progress_bar_config(disable=_a		)

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

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

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

								UpperCamelCase   =	pipe.tokenizer(
								    _a   ,							padding="max_length"   ,							max_length=pipe.tokenizer.model_max_length   ,							truncation=_a   ,							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(**_a		)
								UpperCamelCase   =	output.images[0, -3:, -3:, -1]

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




						def     snake_case_   (self		) ->   Any:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider"		)
								pipe.set_progress_bar_config(disable=_a		)

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

								# forward
								UpperCamelCase   =	pipe(**_a		)
								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(
										    _a   ,							padding="max_length"   ,							max_length=pipe.tokenizer.model_max_length   ,							truncation=_a   ,							return_tensors="np"   ,							)
										UpperCamelCase   =	text_inputs["""input_ids"""]

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

								UpperCamelCase   =	embeds

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


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


						@property
						def     snake_case_   (self		) ->   int:
								UpperCamelCase   =	ort.SessionOptions()
								UpperCamelCase   =	False
								return options


						def     snake_case_   (self		) ->   Optional[int]:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(
								    "CompVis/stable-diffusion-v1-4"   ,							revision="onnx"   ,							safety_checker=_a   ,							feature_extractor=_a   ,							provider=self.gpu_provider   ,							sess_options=self.gpu_options   ,							)
								sd_pipe.set_progress_bar_config(disable=_a		)

								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, 5_12, 5_12, 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     snake_case_   (self		) ->   Dict:
								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=_a   ,							safety_checker=_a   ,							feature_extractor=_a   ,							provider=self.gpu_provider   ,							sess_options=self.gpu_options   ,							)
								sd_pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 5_12, 5_12, 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     snake_case_   (self		) ->   Dict:
								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=_a   ,							safety_checker=_a   ,							feature_extractor=_a   ,							provider=self.gpu_provider   ,							sess_options=self.gpu_options   ,							)
								sd_pipe.set_progress_bar_config(disable=_a		)

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

								assert image.shape == (1, 5_12, 5_12, 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     snake_case_   (self		) ->   Any:
								UpperCamelCase   =	0

								def test_callback_fn(__a   ,							__a   ,							__a		) -> 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=_a   ,							feature_extractor=_a   ,							provider=self.gpu_provider   ,							sess_options=self.gpu_options   ,							)
								pipe.set_progress_bar_config(disable=_a		)

								UpperCamelCase   =	"""Andromeda galaxy in a bottle"""

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




						def     snake_case_   (self		) ->   Union[str, Any]:
								UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(
								    "runwayml/stable-diffusion-v1-5"   ,							revision="onnx"   ,							safety_checker=_a   ,							feature_extractor=_a   ,							provider=self.gpu_provider   ,							sess_options=self.gpu_options   ,							)
								assert isinstance(_a   ,							_a		)
								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(_a		)
										UpperCamelCase   =	OnnxStableDiffusionPipeline.from_pretrained(_a		)

								# 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













"""simple docstring"""





import importlib
import os
import sys


# This is required to make the module import works (when the python process is running from the root of the repo)
sys.path.append('''.''')


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	test_file.split(os.path.sep    )
		if components[0:2] != ["tests", "models"]:
				raise ValueError(
				    "`test_file` should start with `tests/models/` (with `/` being the OS specific path separator). Got "
				    F"{test_file} instead."    )
		UpperCamelCase   =	components[-1]
		if not test_fn.endswith("py"    ):
				raise ValueError(F"`test_file` should be a python file. Got {test_fn} instead."    )
		if not test_fn.startswith("test_modeling_"    ):
				raise ValueError(
				    F"`test_file` should point to a file name of the form `test_modeling_*.py`. Got {test_fn} instead."    )

		UpperCamelCase   =	components[:-1] + [test_fn.replace(".py"  ,	""    )]
		UpperCamelCase   =	".".join(_SCREAMING_SNAKE_CASE    )

		return test_module_path


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_module_path(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	importlib.import_module(_SCREAMING_SNAKE_CASE    )

		return test_module


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	[]
		UpperCamelCase   =	get_test_module(_SCREAMING_SNAKE_CASE    )
		for attr in dir(_SCREAMING_SNAKE_CASE    ):
				if attr.endswith("ModelTester"    ):
						tester_classes.append(getattr(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	[]
		UpperCamelCase   =	get_test_module(_SCREAMING_SNAKE_CASE    )
		for attr in dir(_SCREAMING_SNAKE_CASE    ):
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
				# (TF/Flax)ModelTesterMixin is also an attribute in specific model test module. Let's exclude them by checking
				# `all_model_classes` is not empty (which also excludes other special classes).
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	"all_model_classes"  ,	[]    )
				if len(_SCREAMING_SNAKE_CASE    ) > 0:
						test_classes.append(_SCREAMING_SNAKE_CASE    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	set()
		for test_class in test_classes:
				model_classes.update(test_class.all_model_classes    )

		# sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	test_class()
		if hasattr(_SCREAMING_SNAKE_CASE  ,	"setUp"    ):
				test.setUp()

		UpperCamelCase   =	None
		if hasattr(_SCREAMING_SNAKE_CASE  ,	"model_tester"    ):
				# `(TF/Flax)ModelTesterMixin` has this attribute default to `None`. Let's skip this case.
				if test.model_tester is not None:
						UpperCamelCase   =	test.model_tester.__class__

		return model_tester


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes(_SCREAMING_SNAKE_CASE    )

		UpperCamelCase   =	[]
		for test_class in test_classes:
				if model_class in test_class.all_model_classes:
						target_test_classes.append(_SCREAMING_SNAKE_CASE    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes_for_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

		UpperCamelCase   =	[]
		for test_class in test_classes:
				UpperCamelCase   =	get_model_tester_from_test_class(_SCREAMING_SNAKE_CASE    )
				if tester_class is not None:
						tester_classes.append(_SCREAMING_SNAKE_CASE    )

    # sort with class names
		return sorted(_SCREAMING_SNAKE_CASE  ,	key=lambda _SCREAMING_SNAKE_CASE    : x.__name__    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_test_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	{test_class: get_model_tester_from_test_class(_SCREAMING_SNAKE_CASE    ) for test_class in test_classes}
		return test_tester_mapping


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_model_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	{
		    model_class: get_test_classes_for_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ) for model_class in model_classes
		}
		return model_test_mapping


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_model_classes(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	{
		    model_class: get_tester_classes_for_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ) for model_class in model_classes
		}
		return model_to_tester_mapping






def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		if isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):
				return o
		elif isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):
				return o.__name__
		elif isinstance(_SCREAMING_SNAKE_CASE  ,	(list, tuple)    ):
				return [to_json(_SCREAMING_SNAKE_CASE    ) for x in o]
		elif isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):
				return {to_json(_SCREAMING_SNAKE_CASE    ): to_json(_SCREAMING_SNAKE_CASE    ) for k, v in o.items()}
		else:
				return o






