"""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      ):






					'''simple docstring'''


					a   	=		"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"



					def 			lowerCAmelCase_					(						self			:				Dict   , _snake_case			:				Union[str, Any]=0		)							->       Dict:
									SCREAMING_SNAKE_CASE__		  =				floats_tensor((1, 3, 128, 128)   , rng=random.Random(lowercase__		)		)
									SCREAMING_SNAKE_CASE__		  =				np.random.RandomState(lowercase__		)
									SCREAMING_SNAKE_CASE__		  =				{
									    "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 			lowerCAmelCase_					(						self			:				Any		)							->       Union[str, Any]:
									SCREAMING_SNAKE_CASE__		  =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   , provider="CPUExecutionProvider"		)
									pipe.set_progress_bar_config(disable=lowercase__		)

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

									assert image.shape == (1, 128, 128, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.6_9643, 0.5_8484, 0.5_0314, 0.5_8760, 0.5_5368, 0.5_9643, 0.5_1529, 0.4_1217, 0.4_9087]		)
									assert np.abs(image_slice - expected_slice		).max() < 1e-1



					def 			lowerCAmelCase_					(						self			:				Optional[int]		)							->       List[Any]:
									SCREAMING_SNAKE_CASE__		  =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   , provider="CPUExecutionProvider"		)
									SCREAMING_SNAKE_CASE__		  =				PNDMScheduler.from_config(pipe.scheduler.config   , skip_prk_steps=lowercase__		)
									pipe.set_progress_bar_config(disable=lowercase__		)

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

									assert image.shape == (1, 128, 128, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.6_1737, 0.5_4642, 0.5_3183, 0.5_4465, 0.5_2742, 0.6_0525, 0.4_9969, 0.4_0655, 0.4_8154]		)

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



					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       Optional[int]:
									SCREAMING_SNAKE_CASE__		  =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   , provider="CPUExecutionProvider"		)
									SCREAMING_SNAKE_CASE__		  =				LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
									pipe.set_progress_bar_config(disable=lowercase__		)

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

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

									assert image.shape == (1, 128, 128, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.5_2761, 0.5_9977, 0.4_9033, 0.4_9619, 0.5_4282, 0.5_0311, 0.4_7600, 0.4_0918, 0.4_5203]		)

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



					def 			lowerCAmelCase_					(						self			:				str		)							->       Optional[Any]:
									SCREAMING_SNAKE_CASE__		  =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   , provider="CPUExecutionProvider"		)
									SCREAMING_SNAKE_CASE__		  =				EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
									pipe.set_progress_bar_config(disable=lowercase__		)

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

									assert image.shape == (1, 128, 128, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]		)

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



					def 			lowerCAmelCase_					(						self			:				Dict		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   , provider="CPUExecutionProvider"		)
									SCREAMING_SNAKE_CASE__		  =				EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
									pipe.set_progress_bar_config(disable=lowercase__		)

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

									assert image.shape == (1, 128, 128, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]		)

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



					def 			lowerCAmelCase_					(						self			:				int		)							->       Optional[Any]:
									SCREAMING_SNAKE_CASE__		  =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   , provider="CPUExecutionProvider"		)
									SCREAMING_SNAKE_CASE__		  =				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
									pipe.set_progress_bar_config(disable=lowercase__		)

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

									assert image.shape == (1, 128, 128, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.6_5331, 0.5_8277, 0.4_8204, 0.5_6059, 0.5_3665, 0.5_6235, 0.5_0969, 0.4_0009, 0.4_6552]		)

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



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






					'''simple docstring'''


					@property
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       str:
									return (
									    "CUDAExecutionProvider",
									    {
									        "gpu_mem_limit": "15000000000",  # 15GB
									        "arena_extend_strategy": "kSameAsRequested",
									    },
									)



					@property
					def 			lowerCAmelCase_					(						self			:				Dict		)							->       Optional[Any]:
									SCREAMING_SNAKE_CASE__		  =				ort.SessionOptions()
									SCREAMING_SNAKE_CASE__		  =				False
									return options



					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/img2img/sketch-mountains-input.jpg"		)
									SCREAMING_SNAKE_CASE__		  =				init_image.resize((768, 512)		)
									# using the PNDM scheduler by default
									SCREAMING_SNAKE_CASE__		  =				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__		)

									SCREAMING_SNAKE_CASE__		  =				"A fantasy landscape, trending on artstation"

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

									assert images.shape == (1, 512, 768, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]		)
									# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/img2img/sketch-mountains-input.jpg"		)
									SCREAMING_SNAKE_CASE__		  =				init_image.resize((768, 512)		)
									SCREAMING_SNAKE_CASE__		  =				LMSDiscreteScheduler.from_pretrained(
									    "runwayml/stable-diffusion-v1-5"   , subfolder="scheduler"   , revision="onnx"		)
									SCREAMING_SNAKE_CASE__		  =				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__		)

									SCREAMING_SNAKE_CASE__		  =				"A fantasy landscape, trending on artstation"

									SCREAMING_SNAKE_CASE__		  =				np.random.RandomState(0		)
									SCREAMING_SNAKE_CASE__		  =				pipe(
									    prompt=lowercase__   , image=lowercase__   , strength=0.75   , guidance_scale=7.5   , num_inference_steps=20   , generator=lowercase__   , output_type="np"   , )
									SCREAMING_SNAKE_CASE__		  =				output.images
									SCREAMING_SNAKE_CASE__		  =				images[0, 255:258, 383:386, -1]

									assert images.shape == (1, 512, 768, 3)
									SCREAMING_SNAKE_CASE__		  =				np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]		)
									# 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 typing import Callable, Optional

from .. import Features
from ..packaged_modules.generator.generator import Generator
from .abc import AbstractDatasetInputStream

class lowerCamelCase		(_SCREAMING_SNAKE_CASE      ):






					'''simple docstring'''


					def __init__(						self			:				List[str]   , _snake_case			:				Callable   , _snake_case			:				Optional[Features] = None   , _snake_case			:				str = None   , _snake_case			:				bool = False   , _snake_case			:				bool = False   , _snake_case			:				Optional[dict] = None   , _snake_case			:				Optional[int] = None   , **_snake_case			:				Dict   , )							->       Any:
									super().__init__(
									    features=_snake_case   , cache_dir=_snake_case   , keep_in_memory=_snake_case   , streaming=_snake_case   , num_proc=_snake_case   , **_snake_case   , )
									SCREAMING_SNAKE_CASE__		  =				Generator(
									    cache_dir=_snake_case   , features=_snake_case   , generator=_snake_case   , gen_kwargs=_snake_case   , **_snake_case   , )



					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       Optional[Any]:
									# Build iterable dataset
									if self.streaming:
													SCREAMING_SNAKE_CASE__		  =				self.builder.as_streaming_dataset(split="train"		)
									# Build regular (map-style) dataset
									else:
													SCREAMING_SNAKE_CASE__		  =				None
													SCREAMING_SNAKE_CASE__		  =				None
													SCREAMING_SNAKE_CASE__		  =				None
													SCREAMING_SNAKE_CASE__		  =				None

													self.builder.download_and_prepare(
													    download_config=_snake_case   , download_mode=_snake_case   , verification_mode=_snake_case   , base_path=_snake_case   , num_proc=self.num_proc   , )
													SCREAMING_SNAKE_CASE__		  =				self.builder.as_dataset(
													    split="train"   , verification_mode=_snake_case   , in_memory=self.keep_in_memory		)
									return dataset




