


"""simple docstring"""


import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
	import onnxruntime as ort



class 			lowerCamelCase__						(   UpperCAmelCase_							,    unittest.TestCase			):
						lowerCAmelCase       =     """hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""







						def __a						(							self   :    str			,			_lowercase   :    int=0		):
								A					       = floats_tensor((1, 3, 128, 128)			,			rng=random.Random(_lowercase		)		)
								A					       = np.random.RandomState(_lowercase		)
								A					       = {
								    'prompt': 'A painting of a squirrel eating a burger',
								    'image': image,
								    'generator': generator,
								    'num_inference_steps': 3,
								    'strength': 0.7_5,
								    'guidance_scale': 7.5,
								    'output_type': 'numpy',
								}
								return inputs







						def __a						(							self   :    Optional[int]		):
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,			provider='CPUExecutionProvider'		)
								pipe.set_progress_bar_config(disable=_lowercase		)

								A					       = self.get_dummy_inputs()
								A					       = pipe(**_lowercase		).images
								A					       = image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 128, 128, 3)
								A					       = np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]		)
								assert np.abs(image_slice - expected_slice		).max() < 1e-1







						def __a						(							self   :    List[Any]		):
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,			provider='CPUExecutionProvider'		)
								A					       = PNDMScheduler.from_config(pipe.scheduler.config			,			skip_prk_steps=_lowercase		)
								pipe.set_progress_bar_config(disable=_lowercase		)

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

								assert image.shape == (1, 128, 128, 3)
								A					       = np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]		)

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







						def __a						(							self   :    List[Any]		):
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,			provider='CPUExecutionProvider'		)
								A					       = LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_lowercase		)

								# warmup pass to apply optimizations
								A					       = pipe(**self.get_dummy_inputs()		)

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

								assert image.shape == (1, 128, 128, 3)
								A					       = np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]		)

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







						def __a						(							self   :    Tuple		):
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,			provider='CPUExecutionProvider'		)
								A					       = EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_lowercase		)

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

								assert image.shape == (1, 128, 128, 3)
								A					       = np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]		)

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







						def __a						(							self   :    Any		):
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,			provider='CPUExecutionProvider'		)
								A					       = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_lowercase		)

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

								assert image.shape == (1, 128, 128, 3)
								A					       = np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]		)

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







						def __a						(							self   :    List[Any]		):
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,			provider='CPUExecutionProvider'		)
								A					       = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
								pipe.set_progress_bar_config(disable=_lowercase		)

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

								assert image.shape == (1, 128, 128, 3)
								A					       = np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]		)

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



@nightly
@require_onnxruntime
@require_torch_gpu
class 			lowerCamelCase__						(   unittest.TestCase			):







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







						@property
						def __a						(							self   :    str		):
								A					       = ort.SessionOptions()
								A					       = False
								return options







						def __a						(							self   :    Union[str, Any]		):
								A					       = load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/img2img/sketch-mountains-input.jpg'		)
								A					       = init_image.resize((768, 512)		)
								# using the PNDM scheduler by default
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(
								    'CompVis/stable-diffusion-v1-4'			,			revision='onnx'			,			safety_checker=_lowercase			,			feature_extractor=_lowercase			,			provider=self.gpu_provider			,			sess_options=self.gpu_options			,			)
								pipe.set_progress_bar_config(disable=_lowercase		)

								A					       = 'A fantasy landscape, trending on artstation'

								A					       = np.random.RandomState(0		)
								A					       = pipe(
								    prompt=_lowercase			,			image=_lowercase			,			strength=0.7_5			,			guidance_scale=7.5			,			num_inference_steps=10			,			generator=_lowercase			,			output_type='np'			,			)
								A					       = output.images
								A					       = images[0, 255:258, 383:386, -1]

								assert images.shape == (1, 512, 768, 3)
								A					       = np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]		)
								# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







						def __a						(							self   :    int		):
								A					       = load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/img2img/sketch-mountains-input.jpg'		)
								A					       = init_image.resize((768, 512)		)
								A					       = LMSDiscreteScheduler.from_pretrained(
								    'runwayml/stable-diffusion-v1-5'			,			subfolder='scheduler'			,			revision='onnx'		)
								A					       = OnnxStableDiffusionImgaImgPipeline.from_pretrained(
								    'runwayml/stable-diffusion-v1-5'			,			revision='onnx'			,			scheduler=_lowercase			,			safety_checker=_lowercase			,			feature_extractor=_lowercase			,			provider=self.gpu_provider			,			sess_options=self.gpu_options			,			)
								pipe.set_progress_bar_config(disable=_lowercase		)

								A					       = 'A fantasy landscape, trending on artstation'

								A					       = np.random.RandomState(0		)
								A					       = pipe(
								    prompt=_lowercase			,			image=_lowercase			,			strength=0.7_5			,			guidance_scale=7.5			,			num_inference_steps=20			,			generator=_lowercase			,			output_type='np'			,			)
								A					       = output.images
								A					       = images[0, 255:258, 383:386, -1]

								assert images.shape == (1, 512, 768, 3)
								A					       = np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]		)
								# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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









"""simple docstring"""


from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
	import tensorflow as tf


UpperCamelCase						:				Optional[Any]             =      logging.get_logger(__name__)



@dataclass
class 			lowerCamelCase__						(   UpperCAmelCase_			):
						lowerCAmelCase       =     [
						    """no_inference""",
						    """no_cuda""",
						    """no_tpu""",
						    """no_speed""",
						    """no_memory""",
						    """no_env_print""",
						    """no_multi_process""",
						]







						def __init__(							self   :    str			,			**_lowercase   :    Union[str, Any]		):
								for deprecated_arg in self.deprecated_args:
										if deprecated_arg in kwargs:
												A					       = deprecated_arg[3:]
												A					       = not kwargs.pop(_lowercase		)
												logger.warning(
												    f'{deprecated_arg} is depreciated. Please use --no-{positive_arg} or'
												    f' {positive_arg}={kwargs[positive_arg]}'		)
								A					       = kwargs.pop('tpu_name'			,			self.tpu_name		)
								A					       = kwargs.pop('device_idx'			,			self.device_idx		)
								A					       = kwargs.pop('eager_mode'			,			self.eager_mode		)
								A					       = kwargs.pop('use_xla'			,			self.use_xla		)
								super().__init__(**_lowercase		)

						lowerCAmelCase       =     field(
						    default=UpperCAmelCase_							,    metadata={"""help""": """Name of TPU"""}							,    )
						lowerCAmelCase       =     field(
						    default=0							,    metadata={"""help""": """CPU / GPU device index. Defaults to 0."""}							,    )
						lowerCAmelCase       =     field(default=UpperCAmelCase_							,    metadata={"""help""": """Benchmark models in eager model."""}			)
						lowerCAmelCase       =     field(
						    default=UpperCAmelCase_							,    metadata={
						        """help""": """Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`."""
						    }							,    )







						@cached_property
						def __a						(							self   :    Optional[Any]		):
								requires_backends(self			,			['tf']		)
								A					       = None
								if self.tpu:
										try:
												if self.tpu_name:
														A					       = tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name		)
												else:
														A					       = tf.distribute.cluster_resolver.TPUClusterResolver()
										except ValueError:
												A					       = None
								return tpu







						@cached_property
						def __a						(							self   :    Dict		):
								requires_backends(self			,			['tf']		)
								if self.is_tpu:
										tf.config.experimental_connect_to_cluster(self._setup_tpu		)
										tf.tpu.experimental.initialize_tpu_system(self._setup_tpu		)

										A					       = tf.distribute.TPUStrategy(self._setup_tpu		)
								else:
										# currently no multi gpu is allowed
										if self.is_gpu:
												# TODO: Currently only single GPU is supported
												tf.config.set_visible_devices(self.gpu_list[self.device_idx]			,			'GPU'		)
												A					       = tf.distribute.OneDeviceStrategy(device=f'/gpu:{self.device_idx}'		)
										else:
												tf.config.set_visible_devices([]			,			'GPU'		)  # disable GPU
												A					       = tf.distribute.OneDeviceStrategy(device=f'/cpu:{self.device_idx}'		)

								return strategy







						@property
						def __a						(							self   :    List[Any]		):
								requires_backends(self			,			['tf']		)
								return self._setup_tpu is not None







						@property
						def __a						(							self   :    Optional[Any]		):
								requires_backends(self			,			['tf']		)
								return self._setup_strategy







						@property
						def __a						(							self   :    str		):
								requires_backends(self			,			['tf']		)
								return tf.config.list_physical_devices('GPU'		)







						@property
						def __a						(							self   :    Any		):
								requires_backends(self			,			['tf']		)
								if self.cuda:
										return len(self.gpu_list		)
								return 0







						@property
						def __a						(							self   :    Dict		):
								return self.n_gpu > 0





