
"""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 				A__ (			_lowerCamelCase ,  unittest.TestCase):
				A_		: Tuple            =			'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'

				def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE=0				):
					__lowerCAmelCase    :		Union[str, Any]		    =   np.random.RandomState(_SCREAMING_SNAKE_CASE				)
					__lowerCAmelCase    :		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 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		str		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					assert image.shape == (1, 1_28, 1_28, 3)
					__lowerCAmelCase    :		str		    =   np.array([0.6_5072, 0.5_8492, 0.4_8219, 0.5_5521, 0.5_3180, 0.5_5939, 0.5_0697, 0.3_9800, 0.4_6455]				)

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

				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Optional[int]		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					__lowerCAmelCase    :		str		    =   PNDMScheduler.from_config(pipe.scheduler.config ,	skip_prk_steps=_SCREAMING_SNAKE_CASE				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					assert image.shape == (1, 1_28, 1_28, 3)
					__lowerCAmelCase    :		Union[str, Any]		    =   np.array([0.6_5863, 0.5_9425, 0.4_9326, 0.5_6313, 0.5_3875, 0.5_6627, 0.5_1065, 0.3_9777, 0.4_6330]				)

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

				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		int		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					__lowerCAmelCase    :		int		    =   LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					assert image.shape == (1, 1_28, 1_28, 3)
					__lowerCAmelCase    :		Optional[int]		    =   np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]				)

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

				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Dict		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					__lowerCAmelCase    :		Any		    =   EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					assert image.shape == (1, 1_28, 1_28, 3)
					__lowerCAmelCase    :		str		    =   np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]				)

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

				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Optional[int]		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					__lowerCAmelCase    :		int		    =   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					assert image.shape == (1, 1_28, 1_28, 3)
					__lowerCAmelCase    :		Optional[int]		    =   np.array([0.5_3817, 0.6_0812, 0.4_7384, 0.4_9530, 0.5_1894, 0.4_9814, 0.4_7984, 0.3_8958, 0.4_4271]				)

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

				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Union[str, Any]		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					__lowerCAmelCase    :		int		    =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					assert image.shape == (1, 1_28, 1_28, 3)
					__lowerCAmelCase    :		List[Any]		    =   np.array([0.5_3895, 0.6_0808, 0.4_7933, 0.4_9608, 0.5_1886, 0.4_9950, 0.4_8053, 0.3_8957, 0.4_4200]				)

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

				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Optional[Any]		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					# forward
					__lowerCAmelCase    :		Any		    =   pipe(**_SCREAMING_SNAKE_CASE				)
					__lowerCAmelCase    :		Any		    =   output.images[0, -3:, -3:, -1]

					__lowerCAmelCase    :		List[Any]		    =   self.get_dummy_inputs()
					__lowerCAmelCase    :		Union[str, Any]		    =   3 * [inputs.pop('prompt'				)]

					__lowerCAmelCase    :		str		    =   pipe.tokenizer(
					    _SCREAMING_SNAKE_CASE ,	padding='max_length' ,	max_length=pipe.tokenizer.model_max_length ,	truncation=_SCREAMING_SNAKE_CASE ,	return_tensors='np' ,	)
					__lowerCAmelCase    :		List[Any]		    =   text_inputs['input_ids']

					__lowerCAmelCase    :		List[Any]		    =   pipe.text_encoder(input_ids=text_inputs.astype(np.intaa				)				)[0]

					__lowerCAmelCase    :		Union[str, Any]		    =   prompt_embeds

					# forward
					__lowerCAmelCase    :		List[Any]		    =   pipe(**_SCREAMING_SNAKE_CASE				)
					__lowerCAmelCase    :		List[str]		    =   output.images[0, -3:, -3:, -1]

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







				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		List[Any]		    =   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,	provider='CPUExecutionProvider'				)
					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

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

					# forward
					__lowerCAmelCase    :		List[str]		    =   pipe(**_SCREAMING_SNAKE_CASE				)
					__lowerCAmelCase    :		Dict		    =   output.images[0, -3:, -3:, -1]

					__lowerCAmelCase    :		List[Any]		    =   self.get_dummy_inputs()
					__lowerCAmelCase    :		List[str]		    =   3 * [inputs.pop('prompt'				)]

					__lowerCAmelCase    :		List[str]		    =   []
					for p in [prompt, negative_prompt]:
						__lowerCAmelCase    :		int		    =   pipe.tokenizer(
						    _SCREAMING_SNAKE_CASE ,	padding='max_length' ,	max_length=pipe.tokenizer.model_max_length ,	truncation=_SCREAMING_SNAKE_CASE ,	return_tensors='np' ,	)
						__lowerCAmelCase    :		List[str]		    =   text_inputs['input_ids']

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

					__lowerCAmelCase    :		Any		    =   embeds

					# forward
					__lowerCAmelCase    :		Union[str, Any]		    =   pipe(**_SCREAMING_SNAKE_CASE				)
					__lowerCAmelCase    :		List[str]		    =   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 				A__ (			unittest.TestCase):

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

				@property
				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Optional[Any]		    =   ort.SessionOptions()
					__lowerCAmelCase    :		List[str]		    =   False
					return options

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

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

					__lowerCAmelCase    :		Tuple		    =   image[0, -3:, -3:, -1]

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

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

					assert image.shape == (1, 5_12, 5_12, 3)
					__lowerCAmelCase    :		Tuple		    =   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 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		str		    =   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=_SCREAMING_SNAKE_CASE ,	safety_checker=_SCREAMING_SNAKE_CASE ,	feature_extractor=_SCREAMING_SNAKE_CASE ,	provider=self.gpu_provider ,	sess_options=self.gpu_options ,	)
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE				)

					__lowerCAmelCase    :		List[str]		    =   'open neural network exchange'
					__lowerCAmelCase    :		Tuple		    =   np.random.RandomState(0				)
					__lowerCAmelCase    :		List[str]		    =   sd_pipe([prompt] ,	guidance_scale=7.5 ,	num_inference_steps=10 ,	generator=_SCREAMING_SNAKE_CASE ,	output_type='np'				)
					__lowerCAmelCase    :		List[str]		    =   output.images
					__lowerCAmelCase    :		Optional[int]		    =   image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					__lowerCAmelCase    :		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 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Dict		    =   0

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

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

					__lowerCAmelCase    :		int		    =   'Andromeda galaxy in a bottle'

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







				def 							__lowerCamelCase							(   self				):
					__lowerCAmelCase    :		Union[str, Any]		    =   OnnxStableDiffusionPipeline.from_pretrained(
					    'runwayml/stable-diffusion-v1-5' ,	revision='onnx' ,	safety_checker=_SCREAMING_SNAKE_CASE ,	feature_extractor=_SCREAMING_SNAKE_CASE ,	provider=self.gpu_provider ,	sess_options=self.gpu_options ,	)
					assert isinstance(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)
					assert pipe.safety_checker is None

					__lowerCAmelCase    :		Optional[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(_SCREAMING_SNAKE_CASE				)
						__lowerCAmelCase    :		Optional[Any]		    =   OnnxStableDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE				)

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

"""simple docstring"""







from transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments
def 					__lowerCAmelCase       ():
	__lowerCAmelCase    :		Tuple		    =   HfArgumentParser(_UpperCamelCase	)
	__lowerCAmelCase    :		Optional[Any]		    =   parser.parse_args_into_dataclasses()[0]
	__lowerCAmelCase    :		int		    =   TensorFlowBenchmark(args=_UpperCamelCase	)
	try:
		__lowerCAmelCase    :		Dict		    =   parser.parse_args_into_dataclasses()[0]
	except ValueError as e:
		__lowerCAmelCase    :		List[Any]		    =   'Arg --no_{0} is no longer used, please use --no-{0} instead.'
		__lowerCAmelCase    :		Union[str, Any]		    =   ' '.join(str(_UpperCamelCase	).split(' '	)[:-1]	)
		__lowerCAmelCase    :		List[str]		    =   ''
		__lowerCAmelCase    :		Any		    =   eval(str(_UpperCamelCase	).split(' '	)[-1]	)
		__lowerCAmelCase    :		Dict		    =   []
		for arg in depreciated_args:
			# arg[2:] removes '--'
			if arg[2:] in TensorFlowBenchmark.deprecated_args:
				# arg[5:] removes '--no_'
				full_error_msg += arg_error_msg.format(arg[5:]	)
			else:
				wrong_args.append(_UpperCamelCase	)
		if len(_UpperCamelCase	) > 0:
			__lowerCAmelCase    :		Any		    =   full_error_msg + begin_error_msg + str(_UpperCamelCase	)
		raise ValueError(_UpperCamelCase	)
	benchmark.run()


if __name__ == "__main__":
							main()