



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						):
					__A     :    str           =						'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'







					def  _lowercase			(       self      :     Dict			,		_A      :     Union[str, Any]=0):
											A__	:							Union[str, Any]						    = floats_tensor((1, 3, 128, 128)			,		rng=random.Random(_A))
											A__	:							Any						    = np.random.RandomState(_A)
											A__	:							Optional[int]						    = {
											    "prompt": "A painting of a squirrel eating a burger",
											    "image": image,
											    "generator": generator,
											    "num_inference_steps": 3,
											    "strength": 0.75,
											    "guidance_scale": 7.5,
											    "output_type": "numpy",
											}
											return inputs







					def  _lowercase			(       self      :     List[str]):
											A__	:							List[str]						    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,		provider="CPUExecutionProvider")
											pipe.set_progress_bar_config(disable=_A)

											A__	:							Any						    = self.get_dummy_inputs()
											A__	:							Optional[Any]						    = pipe(**_A).images
											A__	:							Optional[Any]						    = image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 128, 128, 3)
											A__	:							Union[str, Any]						    = np.array([0.6_96_43, 0.5_84_84, 0.5_03_14, 0.5_87_60, 0.5_53_68, 0.5_96_43, 0.5_15_29, 0.4_12_17, 0.4_90_87])
											assert np.abs(image_slice - expected_slice).max() < 1e-1







					def  _lowercase			(       self      :     Dict):
											A__	:							Union[str, Any]						    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,		provider="CPUExecutionProvider")
											A__	:							Union[str, Any]						    = PNDMScheduler.from_config(pipe.scheduler.config			,		skip_prk_steps=_A)
											pipe.set_progress_bar_config(disable=_A)

											A__	:							Union[str, Any]						    = self.get_dummy_inputs()
											A__	:							Optional[int]						    = pipe(**_A).images
											A__	:							Union[str, Any]						    = image[0, -3:, -3:, -1]

											assert image.shape == (1, 128, 128, 3)
											A__	:							List[str]						    = np.array([0.6_17_37, 0.5_46_42, 0.5_31_83, 0.5_44_65, 0.5_27_42, 0.6_05_25, 0.4_99_69, 0.4_06_55, 0.4_81_54])

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







					def  _lowercase			(       self      :     str):
											A__	:							str						    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,		provider="CPUExecutionProvider")
											A__	:							List[Any]						    = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
											pipe.set_progress_bar_config(disable=_A)

											# warmup pass to apply optimizations
											A__	:							Optional[Any]						    = pipe(**self.get_dummy_inputs())

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

											assert image.shape == (1, 128, 128, 3)
											A__	:							Optional[Any]						    = np.array([0.5_27_61, 0.5_99_77, 0.4_90_33, 0.4_96_19, 0.5_42_82, 0.5_03_11, 0.4_76_00, 0.4_09_18, 0.4_52_03])

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







					def  _lowercase			(       self      :     Optional[Any]):
											A__	:							Tuple						    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,		provider="CPUExecutionProvider")
											A__	:							str						    = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
											pipe.set_progress_bar_config(disable=_A)

											A__	:							Dict						    = self.get_dummy_inputs()
											A__	:							int						    = pipe(**_A).images
											A__	:							Optional[int]						    = image[0, -3:, -3:, -1]

											assert image.shape == (1, 128, 128, 3)
											A__	:							int						    = np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04])

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







					def  _lowercase			(       self      :     List[Any]):
											A__	:							int						    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,		provider="CPUExecutionProvider")
											A__	:							str						    = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
											pipe.set_progress_bar_config(disable=_A)

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

											assert image.shape == (1, 128, 128, 3)
											A__	:							List[str]						    = np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04])

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







					def  _lowercase			(       self      :     List[Any]):
											A__	:							int						    = OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint			,		provider="CPUExecutionProvider")
											A__	:							Tuple						    = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
											pipe.set_progress_bar_config(disable=_A)

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

											assert image.shape == (1, 128, 128, 3)
											A__	:							Tuple						    = np.array([0.6_53_31, 0.5_82_77, 0.4_82_04, 0.5_60_59, 0.5_36_65, 0.5_62_35, 0.5_09_69, 0.4_00_09, 0.4_65_52])

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




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







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







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







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

											A__	:							List[Any]						    = "A fantasy landscape, trending on artstation"

											A__	:							List[Any]						    = np.random.RandomState(0)
											A__	:							Dict						    = pipe(
											    prompt=_A			,		image=_A			,		strength=0.75			,		guidance_scale=7.5			,		num_inference_steps=10			,		generator=_A			,		output_type="np"			,		)
											A__	:							List[str]						    = output.images
											A__	:							str						    = images[0, 255:258, 383:386, -1]

											assert images.shape == (1, 512, 768, 3)
											A__	:							str						    = np.array([0.49_09, 0.50_59, 0.53_72, 0.46_23, 0.48_76, 0.50_49, 0.48_20, 0.49_56, 0.50_19])
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







					def  _lowercase			(       self      :     str):
											A__	:							Union[str, Any]						    = load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg")
											A__	:							Any						    = init_image.resize((768, 512))
											A__	:							Optional[int]						    = LMSDiscreteScheduler.from_pretrained(
											    "runwayml/stable-diffusion-v1-5"			,		subfolder="scheduler"			,		revision="onnx")
											A__	:							Any						    = OnnxStableDiffusionImgaImgPipeline.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			,		)
											pipe.set_progress_bar_config(disable=_A)

											A__	:							Optional[int]						    = "A fantasy landscape, trending on artstation"

											A__	:							Optional[Any]						    = np.random.RandomState(0)
											A__	:							Tuple						    = pipe(
											    prompt=_A			,		image=_A			,		strength=0.75			,		guidance_scale=7.5			,		num_inference_steps=20			,		generator=_A			,		output_type="np"			,		)
											A__	:							Union[str, Any]						    = output.images
											A__	:							Optional[Any]						    = images[0, 255:258, 383:386, -1]

											assert images.shape == (1, 512, 768, 3)
											A__	:							int						    = np.array([0.80_43, 0.9_26, 0.95_81, 0.81_19, 0.89_54, 0.9_13, 0.72_09, 0.74_63, 0.74_31])
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
		import jax
		import jaxlib

snake_case    :						Optional[int]  =						get_logger()

snake_case    :						Optional[dict]  =						None




class 							lowerCAmelCase__  (	TensorFormatter[Mapping, 'jax.Array', Mapping]						):







					def __init__(       self      :     Optional[Any]			,		_A      :     Optional[int]=None			,		_A      :     str=None			,		**_A      :     Tuple):
											super().__init__(features=_A)
											import jax
											from jaxlib.xla_client import Device

											if isinstance(_A			,		_A):
																	raise ValueError(
																	    F'Expected {device} to be a `str` not {type(_A)}, as `jaxlib.xla_extension.Device` '
																	    "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
																	    "the device with `str()` to get its string identifier that will be internally mapped "
																	    "to the actual `jaxlib.xla_extension.Device`.")
											A__	:							Union[str, Any]						    = device if isinstance(_A			,		_A) else str(jax.devices()[0])
											# using global variable since `jaxlib.xla_extension.Device` is not serializable neither
											# with `pickle` nor with `dill`, so we need to use a global variable instead
											global DEVICE_MAPPING
											if DEVICE_MAPPING is None:
																	A__	:							Any						    = self._map_devices_to_str()
											if self.device not in list(DEVICE_MAPPING.keys()):
																	logger.warning(
																	    F'Device with string identifier {self.device} not listed among the available '
																	    F'devices: {list(DEVICE_MAPPING.keys())}, so falling back to the default '
																	    F'device: {str(jax.devices()[0])}.')
																	A__	:							str						    = str(jax.devices()[0])
											A__	:							List[Any]						    = jnp_array_kwargs







					@staticmethod
					def  _lowercase			(       ):
											import jax

											return {str(_A): device for device in jax.devices()}







					def  _lowercase			(       self      :     Tuple			,		_A      :     int):
											import jax
											import jax.numpy as jnp

											if isinstance(_A			,		_A) and column:
																	if all(
																	    isinstance(_A			,		jax.Array) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column):
																							return jnp.stack(_A			,		axis=0)
											return column







					def  _lowercase			(       self      :     List[str]			,		_A      :     Union[str, Any]):
											import jax
											import jax.numpy as jnp

											if isinstance(_A			,		(str, bytes, type(_A))):
																	return value
											elif isinstance(_A			,		(np.character, np.ndarray)) and np.issubdtype(value.dtype			,		np.character):
																	return value.tolist()

											A__	:							Dict						    = {}

											if isinstance(_A			,		(np.number, np.ndarray)) and np.issubdtype(value.dtype			,		np.integer):
																	# the default int precision depends on the jax config
																	# see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
																	if jax.config.jax_enable_xaa:
																							A__	:							Union[str, Any]						    = {"dtype": jnp.intaa}
																	else:
																							A__	:							int						    = {"dtype": jnp.intaa}
											elif isinstance(_A			,		(np.number, np.ndarray)) and np.issubdtype(value.dtype			,		np.floating):
																	A__	:							str						    = {"dtype": jnp.floataa}
											elif config.PIL_AVAILABLE and "PIL" in sys.modules:
																	import PIL.Image

																	if isinstance(_A			,		PIL.Image.Image):
																							A__	:							Any						    = np.asarray(_A)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
											global DEVICE_MAPPING
											if DEVICE_MAPPING is None:
																	A__	:							Optional[Any]						    = self._map_devices_to_str()

											with jax.default_device(DEVICE_MAPPING[self.device]):
																	# calling jnp.array on a np.ndarray does copy the data
																	# see https://github.com/google/jax/issues/4486
																	return jnp.array(_A			,		**{**default_dtype, **self.jnp_array_kwargs})







					def  _lowercase			(       self      :     Dict			,		_A      :     Any):
											import jax

											# support for torch, tf, jax etc.
											if config.TORCH_AVAILABLE and "torch" in sys.modules:
																	import torch

																	if isinstance(_A			,		torch.Tensor):
																							return self._tensorize(data_struct.detach().cpu().numpy()[()])
											if hasattr(_A			,		"__array__") and not isinstance(_A			,		jax.Array):
																	A__	:							Any						    = data_struct.__array__()
											# support for nested types like struct of list of struct
											if isinstance(_A			,		np.ndarray):
																	if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
																							return self._consolidate([self.recursive_tensorize(_A) for substruct in data_struct])
											elif isinstance(_A			,		(list, tuple)):
																	return self._consolidate([self.recursive_tensorize(_A) for substruct in data_struct])
											return self._tensorize(_A)







					def  _lowercase			(       self      :     Optional[Any]			,		_A      :     dict):
											return map_nested(self._recursive_tensorize			,		_A			,		map_list=_A)







					def  _lowercase			(       self      :     List[str]			,		_A      :     pa.Table):
											A__	:							Dict						    = self.numpy_arrow_extractor().extract_row(_A)
											A__	:							int						    = self.python_features_decoder.decode_row(_A)
											return self.recursive_tensorize(_A)







					def  _lowercase			(       self      :     int			,		_A      :     pa.Table):
											A__	:							int						    = self.numpy_arrow_extractor().extract_column(_A)
											A__	:							Optional[int]						    = self.python_features_decoder.decode_column(_A			,		pa_table.column_names[0])
											A__	:							Dict						    = self.recursive_tensorize(_A)
											A__	:							List[Any]						    = self._consolidate(_A)
											return column







					def  _lowercase			(       self      :     Tuple			,		_A      :     pa.Table):
											A__	:							Any						    = self.numpy_arrow_extractor().extract_batch(_A)
											A__	:							str						    = self.python_features_decoder.decode_batch(_A)
											A__	:							Optional[Any]						    = self.recursive_tensorize(_A)
											for column_name in batch:
																	A__	:							Optional[int]						    = self._consolidate(batch[column_name])
											return batch