


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 				snake_case__(							_UpperCAmelCase,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Dict					=							'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''




						def 						lowercase_ (				self   ,						__lowercase=0							)					->    int:
										lowerCAmelCase_ :      List[str]    				=				floats_tensor((1, 3, 1_2_8, 1_2_8)   ,						rng=random.Random(A_							)							)
										lowerCAmelCase_ :      Optional[int]    				=				np.random.RandomState(A_							)
										lowerCAmelCase_ :      List[Any]    				=				{
										    '''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[Any]:
										lowerCAmelCase_ :      Union[str, Any]    				=				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,						provider='''CPUExecutionProvider'''							)
										pipe.set_progress_bar_config(disable=A_							)

										lowerCAmelCase_ :      Any    				=				self.get_dummy_inputs()
										lowerCAmelCase_ :      Optional[Any]    				=				pipe(**A_							).images
										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 1_2_8, 1_2_8, 3)
										lowerCAmelCase_ :      Dict    				=				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							)					->    Optional[int]:
										lowerCAmelCase_ :      Optional[int]    				=				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,						provider='''CPUExecutionProvider'''							)
										lowerCAmelCase_ :      Optional[int]    				=				PNDMScheduler.from_config(pipe.scheduler.config   ,						skip_prk_steps=A_							)
										pipe.set_progress_bar_config(disable=A_							)

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

										assert image.shape == (1, 1_2_8, 1_2_8, 3)
										lowerCAmelCase_ :      Tuple    				=				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:
										lowerCAmelCase_ :      Union[str, Any]    				=				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,						provider='''CPUExecutionProvider'''							)
										lowerCAmelCase_ :      int    				=				LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
										pipe.set_progress_bar_config(disable=A_							)

										# warmup pass to apply optimizations
										lowerCAmelCase_ :      Any    				=				pipe(**self.get_dummy_inputs()							)

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

										assert image.shape == (1, 1_2_8, 1_2_8, 3)
										lowerCAmelCase_ :      List[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							)					->    List[str]:
										lowerCAmelCase_ :      Optional[Any]    				=				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,						provider='''CPUExecutionProvider'''							)
										lowerCAmelCase_ :      str    				=				EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
										pipe.set_progress_bar_config(disable=A_							)

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

										assert image.shape == (1, 1_2_8, 1_2_8, 3)
										lowerCAmelCase_ :      Optional[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							)					->    Union[str, Any]:
										lowerCAmelCase_ :      Union[str, Any]    				=				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,						provider='''CPUExecutionProvider'''							)
										lowerCAmelCase_ :      Tuple    				=				EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
										pipe.set_progress_bar_config(disable=A_							)

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

										assert image.shape == (1, 1_2_8, 1_2_8, 3)
										lowerCAmelCase_ :      Optional[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							)					->    Any:
										lowerCAmelCase_ :      List[str]    				=				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint   ,						provider='''CPUExecutionProvider'''							)
										lowerCAmelCase_ :      List[Any]    				=				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
										pipe.set_progress_bar_config(disable=A_							)

										lowerCAmelCase_ :      List[Any]    				=				self.get_dummy_inputs()
										lowerCAmelCase_ :      Tuple    				=				pipe(**A_							).images
										lowerCAmelCase_ :      int    				=				image[0, -3:, -3:, -1]

										assert image.shape == (1, 1_2_8, 1_2_8, 3)
										lowerCAmelCase_ :      Union[str, Any]    				=				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 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




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




						@property
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										lowerCAmelCase_ :      List[str]    				=				ort.SessionOptions()
										lowerCAmelCase_ :      List[Any]    				=				False
										return options




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										lowerCAmelCase_ :      Optional[Any]    				=				load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/img2img/sketch-mountains-input.jpg'''							)
										lowerCAmelCase_ :      Any    				=				init_image.resize((7_6_8, 5_1_2)							)
										# using the PNDM scheduler by default
										lowerCAmelCase_ :      Tuple    				=				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_							)

										lowerCAmelCase_ :      str    				=				'''A fantasy landscape, trending on artstation'''

										lowerCAmelCase_ :      List[Any]    				=				np.random.RandomState(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(
										    prompt=A_   ,						image=A_   ,						strength=0.75   ,						guidance_scale=7.5   ,						num_inference_steps=1_0   ,						generator=A_   ,						output_type='''np'''   ,						)
										lowerCAmelCase_ :      List[Any]    				=				output.images
										lowerCAmelCase_ :      Union[str, Any]    				=				images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

										assert images.shape == (1, 5_1_2, 7_6_8, 3)
										lowerCAmelCase_ :      Union[str, Any]    				=				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							)					->    Optional[int]:
										lowerCAmelCase_ :      List[Any]    				=				load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/img2img/sketch-mountains-input.jpg'''							)
										lowerCAmelCase_ :      Dict    				=				init_image.resize((7_6_8, 5_1_2)							)
										lowerCAmelCase_ :      Union[str, Any]    				=				LMSDiscreteScheduler.from_pretrained(
										    '''runwayml/stable-diffusion-v1-5'''   ,						subfolder='''scheduler'''   ,						revision='''onnx'''							)
										lowerCAmelCase_ :      Optional[int]    				=				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_							)

										lowerCAmelCase_ :      str    				=				'''A fantasy landscape, trending on artstation'''

										lowerCAmelCase_ :      List[str]    				=				np.random.RandomState(0							)
										lowerCAmelCase_ :      Tuple    				=				pipe(
										    prompt=A_   ,						image=A_   ,						strength=0.75   ,						guidance_scale=7.5   ,						num_inference_steps=2_0   ,						generator=A_   ,						output_type='''np'''   ,						)
										lowerCAmelCase_ :      Tuple    				=				output.images
										lowerCAmelCase_ :      int    				=				images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

										assert images.shape == (1, 5_1_2, 7_6_8, 3)
										lowerCAmelCase_ :      List[Any]    				=				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



from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
						import PIL


_UpperCAmelCase			:  Any     =logging.get_logger(__name__)







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Dict					=							["""pixel_values"""]




						def __init__(				self   ,						__lowercase = True   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = PILImageResampling.BILINEAR   ,						__lowercase = True   ,						__lowercase = 1 / 2_5_5   ,						__lowercase = True   ,						__lowercase = None   ,						__lowercase = None   ,						**__lowercase   ,						)					->    None:
										super().__init__(**__lowercase							)
										lowerCAmelCase_ :      Dict    				=				size if size is not None else {'''shortest_edge''': 3_8_4}
										lowerCAmelCase_ :      Optional[Any]    				=				get_size_dict(__lowercase   ,						default_to_square=__lowercase							)

										lowerCAmelCase_ :      List[Any]    				=				do_resize
										lowerCAmelCase_ :      Optional[int]    				=				size
										# Default value set here for backwards compatibility where the value in config is None
										lowerCAmelCase_ :      str    				=				crop_pct if crop_pct is not None else 2_2_4 / 2_5_6
										lowerCAmelCase_ :      Tuple    				=				resample
										lowerCAmelCase_ :      Optional[int]    				=				do_rescale
										lowerCAmelCase_ :      Any    				=				rescale_factor
										lowerCAmelCase_ :      List[str]    				=				do_normalize
										lowerCAmelCase_ :      str    				=				image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
										lowerCAmelCase_ :      Dict    				=				image_std if image_std is not None else IMAGENET_STANDARD_STD




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase = PILImageResampling.BICUBIC   ,						__lowercase = None   ,						**__lowercase   ,						)					->    np.ndarray:
										lowerCAmelCase_ :      Optional[Any]    				=				get_size_dict(__lowercase   ,						default_to_square=__lowercase							)
										if "shortest_edge" not in size:
														raise ValueError(f"""Size dictionary must contain 'shortest_edge' key. Got {size.keys()}"""							)
										lowerCAmelCase_ :      Optional[int]    				=				size['''shortest_edge''']

										if shortest_edge < 3_8_4:
														# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
														lowerCAmelCase_ :      Optional[Any]    				=				int(shortest_edge / crop_pct							)
														lowerCAmelCase_ :      List[str]    				=				get_resize_output_image_size(__lowercase   ,						size=__lowercase   ,						default_to_square=__lowercase							)
														lowerCAmelCase_ :      List[Any]    				=				resize(image=__lowercase   ,						size=__lowercase   ,						resample=__lowercase   ,						data_format=__lowercase   ,						**__lowercase							)
														# then crop to (shortest_edge, shortest_edge)
														return center_crop(image=__lowercase   ,						size=(shortest_edge, shortest_edge)   ,						data_format=__lowercase   ,						**__lowercase							)
										else:
														# warping (no cropping) when evaluated at 384 or larger
														return resize(
														    __lowercase   ,						size=(shortest_edge, shortest_edge)   ,						resample=__lowercase   ,						data_format=__lowercase   ,						**__lowercase							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase = None   ,						**__lowercase   ,						)					->    Any:
										return rescale(__lowercase   ,						scale=__lowercase   ,						data_format=__lowercase   ,						**__lowercase							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase = None   ,						**__lowercase   ,						)					->    np.ndarray:
										return normalize(__lowercase   ,						mean=__lowercase   ,						std=__lowercase   ,						data_format=__lowercase   ,						**__lowercase							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = None   ,						__lowercase = ChannelDimension.FIRST   ,						**__lowercase   ,						)					->    PIL.Image.Image:
										lowerCAmelCase_ :      Optional[int]    				=				do_resize if do_resize is not None else self.do_resize
										lowerCAmelCase_ :      Any    				=				crop_pct if crop_pct is not None else self.crop_pct
										lowerCAmelCase_ :      str    				=				resample if resample is not None else self.resample
										lowerCAmelCase_ :      Union[str, Any]    				=				do_rescale if do_rescale is not None else self.do_rescale
										lowerCAmelCase_ :      str    				=				rescale_factor if rescale_factor is not None else self.rescale_factor
										lowerCAmelCase_ :      Any    				=				do_normalize if do_normalize is not None else self.do_normalize
										lowerCAmelCase_ :      str    				=				image_mean if image_mean is not None else self.image_mean
										lowerCAmelCase_ :      int    				=				image_std if image_std is not None else self.image_std

										lowerCAmelCase_ :      int    				=				size if size is not None else self.size
										lowerCAmelCase_ :      List[str]    				=				get_size_dict(__lowercase   ,						default_to_square=__lowercase							)

										lowerCAmelCase_ :      Tuple    				=				make_list_of_images(__lowercase							)

										if not valid_images(__lowercase							):
														raise ValueError(
														    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
														    '''torch.Tensor, tf.Tensor or jax.ndarray.'''							)

										if do_resize and size is None or resample is None:
														raise ValueError('''Size and resample must be specified if do_resize is True.'''							)

										if do_resize and size["shortest_edge"] < 3_8_4 and crop_pct is None:
														raise ValueError('''crop_pct must be specified if size < 384.'''							)

										if do_rescale and rescale_factor is None:
														raise ValueError('''Rescale factor must be specified if do_rescale is True.'''							)

										if do_normalize and (image_mean is None or image_std is None):
														raise ValueError('''Image mean and std must be specified if do_normalize is True.'''							)

										# All transformations expect numpy arrays.
										lowerCAmelCase_ :      Optional[Any]    				=				[to_numpy_array(__lowercase							) for image in images]

										if do_resize:
														lowerCAmelCase_ :      Union[str, Any]    				=				[self.resize(image=__lowercase   ,						size=__lowercase   ,						crop_pct=__lowercase   ,						resample=__lowercase							) for image in images]

										if do_rescale:
														lowerCAmelCase_ :      Any    				=				[self.rescale(image=__lowercase   ,						scale=__lowercase							) for image in images]

										if do_normalize:
														lowerCAmelCase_ :      List[Any]    				=				[self.normalize(image=__lowercase   ,						mean=__lowercase   ,						std=__lowercase							) for image in images]

										lowerCAmelCase_ :      Optional[Any]    				=				[to_channel_dimension_format(__lowercase   ,						__lowercase							) for image in images]

										lowerCAmelCase_ :      Dict    				=				{'''pixel_values''': images}
										return BatchFeature(data=__lowercase   ,						tensor_type=__lowercase							)