

"""simple docstring"""






import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

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


if is_torch_available():
	import torch

if is_vision_available():
	import PIL


snake_case_   :   Optional[int]        =       logging.get_logger(__name__)



def   lowercase_	(    _lowercase       :		np.ndarray							,      _lowercase       :		Union[int, Iterable[int]]							,      _lowercase       :		bool							,      _lowercase       :		int	):




					'''simple docstring'''






					def constraint_to_multiple_of(_lowercase       :		Optional[Any]							,      _lowercase       :		Union[str, Any]							,      _lowercase       :		Union[str, Any]=0							,      _lowercase       :		Optional[int]=None	):
										UpperCAmelCase			:   Tuple      =				round(val / multiple	) * multiple

										if max_val is not None and x > max_val:
															UpperCAmelCase			:   List[str]      =				math.floor(val / multiple	) * multiple

										if x < min_val:
															UpperCAmelCase			:   List[str]      =				math.ceil(val / multiple	) * multiple

										return x

					UpperCAmelCase			:   Tuple      =				(output_size, output_size) if isinstance(_lowercase							,      _lowercase	) else output_size

					UpperCAmelCase      ,							UpperCAmelCase			:   Optional[Any]      =				get_image_size(_lowercase	)
					UpperCAmelCase      ,							UpperCAmelCase			:   Dict      =				output_size

					# determine new height and width
					UpperCAmelCase			:   Any      =				output_height / input_height
					UpperCAmelCase			:   Tuple      =				output_width / input_width

					if keep_aspect_ratio:
										# scale as little as possible
										if abs(1 - scale_width	) < abs(1 - scale_height	):
															# fit width
															UpperCAmelCase			:   List[Any]      =				scale_width
										else:
															# fit height
															UpperCAmelCase			:   Any      =				scale_height

					UpperCAmelCase			:   List[str]      =				constraint_to_multiple_of(scale_height * input_height							,      multiple=_lowercase	)
					UpperCAmelCase			:   Dict      =				constraint_to_multiple_of(scale_width * input_width							,      multiple=_lowercase	)

					return (new_height, new_width)







class       snake_case__  ( lowerCAmelCase_     ):
			SCREAMING_SNAKE_CASE__     				=						['''pixel_values''']







			def __init__(			self			: List[Any]			,  lowercase			: bool = True			,  lowercase			: Dict[str, int] = None			,  lowercase			: PILImageResampling = PILImageResampling.BILINEAR			,  lowercase			: bool = False			,  lowercase			: int = 1			,  lowercase			: bool = True			,  lowercase			: Union[int, float] = 1 / 2_55			,  lowercase			: bool = True			,  lowercase			: Optional[Union[float, List[float]]] = None			,  lowercase			: Optional[Union[float, List[float]]] = None			,  **lowercase			: int			,  ):





								'''simple docstring'''



								super().__init__(**lowercase  )
								UpperCAmelCase			:   Tuple      =				size if size is not None else {"height": 3_84, "width": 3_84}
								UpperCAmelCase			:   Optional[Any]      =				get_size_dict(lowercase  )
								UpperCAmelCase			:   Optional[Any]      =				do_resize
								UpperCAmelCase			:   List[Any]      =				size
								UpperCAmelCase			:   Union[str, Any]      =				keep_aspect_ratio
								UpperCAmelCase			:   List[Any]      =				ensure_multiple_of
								UpperCAmelCase			:   Union[str, Any]      =				resample
								UpperCAmelCase			:   Tuple      =				do_rescale
								UpperCAmelCase			:   Tuple      =				rescale_factor
								UpperCAmelCase			:   str      =				do_normalize
								UpperCAmelCase			:   List[str]      =				image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
								UpperCAmelCase			:   str      =				image_std if image_std is not None else IMAGENET_STANDARD_STD







			def 			__lowerCAmelCase				(			self			: Optional[Any]			,  lowercase			: np.ndarray			,  lowercase			: Dict[str, int]			,  lowercase			: bool = False			,  lowercase			: int = 1			,  lowercase			: PILImageResampling = PILImageResampling.BICUBIC			,  lowercase			: Optional[Union[str, ChannelDimension]] = None			,  **lowercase			: Union[str, Any]			,  ):





								'''simple docstring'''



								UpperCAmelCase			:   Optional[int]      =				get_size_dict(lowercase  )
								if "height" not in size or "width" not in size:
													raise ValueError(f"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"""  )
								UpperCAmelCase			:   Union[str, Any]      =				get_resize_output_image_size(
								    lowercase			,  output_size=(size["height"], size["width"])			,  keep_aspect_ratio=lowercase			,  multiple=lowercase			,  )
								return resize(lowercase			,  size=lowercase			,  resample=lowercase			,  data_format=lowercase			,  **lowercase  )







			def 			__lowerCAmelCase				(			self			: Optional[Any]			,  lowercase			: np.ndarray			,  lowercase			: Union[int, float]			,  lowercase			: Optional[Union[str, ChannelDimension]] = None			,  **lowercase			: Union[str, Any]			,  ):





								'''simple docstring'''



								return rescale(lowercase			,  scale=lowercase			,  data_format=lowercase			,  **lowercase  )







			def 			__lowerCAmelCase				(			self			: int			,  lowercase			: np.ndarray			,  lowercase			: Union[float, List[float]]			,  lowercase			: Union[float, List[float]]			,  lowercase			: Optional[Union[str, ChannelDimension]] = None			,  **lowercase			: Any			,  ):





								'''simple docstring'''



								return normalize(lowercase			,  mean=lowercase			,  std=lowercase			,  data_format=lowercase			,  **lowercase  )







			def 			__lowerCAmelCase				(			self			: List[str]			,  lowercase			: ImageInput			,  lowercase			: bool = None			,  lowercase			: int = None			,  lowercase			: bool = None			,  lowercase			: int = None			,  lowercase			: PILImageResampling = None			,  lowercase			: bool = None			,  lowercase			: float = None			,  lowercase			: bool = None			,  lowercase			: Optional[Union[float, List[float]]] = None			,  lowercase			: Optional[Union[float, List[float]]] = None			,  lowercase			: Optional[Union[str, TensorType]] = None			,  lowercase			: ChannelDimension = ChannelDimension.FIRST			,  **lowercase			: Union[str, Any]			,  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[Any]      =				do_resize if do_resize is not None else self.do_resize
								UpperCAmelCase			:   List[str]      =				size if size is not None else self.size
								UpperCAmelCase			:   Tuple      =				get_size_dict(lowercase  )
								UpperCAmelCase			:   str      =				keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
								UpperCAmelCase			:   Optional[Any]      =				ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
								UpperCAmelCase			:   List[Any]      =				resample if resample is not None else self.resample
								UpperCAmelCase			:   Any      =				do_rescale if do_rescale is not None else self.do_rescale
								UpperCAmelCase			:   Optional[int]      =				rescale_factor if rescale_factor is not None else self.rescale_factor
								UpperCAmelCase			:   Optional[int]      =				do_normalize if do_normalize is not None else self.do_normalize
								UpperCAmelCase			:   Union[str, Any]      =				image_mean if image_mean is not None else self.image_mean
								UpperCAmelCase			:   Union[str, Any]      =				image_std if image_std is not None else self.image_std

								UpperCAmelCase			:   str      =				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_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.
								UpperCAmelCase			:   Union[str, Any]      =				[to_numpy_array(lowercase  ) for image in images]

								if do_resize:
													UpperCAmelCase			:   Any      =				[self.resize(image=lowercase			,  size=lowercase			,  resample=lowercase  ) for image in images]

								if do_rescale:
													UpperCAmelCase			:   int      =				[self.rescale(image=lowercase			,  scale=lowercase  ) for image in images]

								if do_normalize:
													UpperCAmelCase			:   str      =				[self.normalize(image=lowercase			,  mean=lowercase			,  std=lowercase  ) for image in images]

								UpperCAmelCase			:   Optional[Any]      =				[to_channel_dimension_format(lowercase			,  lowercase  ) for image in images]

								UpperCAmelCase			:   Any      =				{"pixel_values": images}
								return BatchFeature(data=lowercase			,  tensor_type=lowercase  )







			def 			__lowerCAmelCase				(			self			: int			,  lowercase			: List[Any]			,  lowercase			: List[Tuple] = None  ):





								'''simple docstring'''



								UpperCAmelCase			:   int      =				outputs.logits

								# Resize logits and compute semantic segmentation maps
								if target_sizes is not None:
													if len(lowercase  ) != len(lowercase  ):
																		raise ValueError(
																		    "Make sure that you pass in as many target sizes as the batch dimension of the logits"  )

													if is_torch_tensor(lowercase  ):
																		UpperCAmelCase			:   Optional[Any]      =				target_sizes.numpy()

													UpperCAmelCase			:   Optional[Any]      =				[]

													for idx in range(len(lowercase  )  ):
																		UpperCAmelCase			:   Any      =				torch.nn.functional.interpolate(
																		    logits[idx].unsqueeze(dim=0  )			,  size=target_sizes[idx]			,  mode="bilinear"			,  align_corners=lowercase  )
																		UpperCAmelCase			:   Optional[Any]      =				resized_logits[0].argmax(dim=0  )
																		semantic_segmentation.append(lowercase  )
								else:
													UpperCAmelCase			:   Any      =				logits.argmax(dim=1  )
													UpperCAmelCase			:   Optional[Any]      =				[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]  )]

								return semantic_segmentation





"""simple docstring"""






import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest







class       snake_case__  ( lowerCAmelCase_     ):
			SCREAMING_SNAKE_CASE__     				=						(UniPCMultistepScheduler,)
			SCREAMING_SNAKE_CASE__     				=						(('''num_inference_steps''', 25),)







			def 			__lowerCAmelCase				(			self			: Tuple			,  **lowercase			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[Any]      =				{
								    "num_train_timesteps": 10_00,
								    "beta_start": 0.0_0_0_1,
								    "beta_end": 0.0_2,
								    "beta_schedule": "linear",
								    "solver_order": 2,
								    "solver_type": "bh2",
								}

								config.update(**lowercase  )
								return config







			def 			__lowerCAmelCase				(			self			: int			,  lowercase			: List[str]=0			,  **lowercase			: str  ):





								'''simple docstring'''



								UpperCAmelCase			:   Union[str, Any]      =				dict(self.forward_default_kwargs  )
								UpperCAmelCase			:   Dict      =				kwargs.pop("num_inference_steps"			,  lowercase  )
								UpperCAmelCase			:   str      =				self.dummy_sample
								UpperCAmelCase			:   Tuple      =				0.1 * sample
								UpperCAmelCase			:   int      =				[residual + 0.2, residual + 0.1_5, residual + 0.1_0]

								for scheduler_class in self.scheduler_classes:
													UpperCAmelCase			:   Optional[Any]      =				self.get_scheduler_config(**lowercase  )
													UpperCAmelCase			:   Tuple      =				scheduler_class(**lowercase  )
													scheduler.set_timesteps(lowercase  )
													# copy over dummy past residuals
													UpperCAmelCase			:   Tuple      =				dummy_past_residuals[: scheduler.config.solver_order]

													with tempfile.TemporaryDirectory() as tmpdirname:
																		scheduler.save_config(lowercase  )
																		UpperCAmelCase			:   Any      =				scheduler_class.from_pretrained(lowercase  )
																		new_scheduler.set_timesteps(lowercase  )
																		# copy over dummy past residuals
																		UpperCAmelCase			:   Tuple      =				dummy_past_residuals[: new_scheduler.config.solver_order]

													UpperCAmelCase      ,							UpperCAmelCase			:   List[Any]      =				sample, sample
													for t in range(lowercase			,  time_step + scheduler.config.solver_order + 1  ):
																		UpperCAmelCase			:   List[str]      =				scheduler.step(lowercase			,  lowercase			,  lowercase			,  **lowercase  ).prev_sample
																		UpperCAmelCase			:   List[Any]      =				new_scheduler.step(lowercase			,  lowercase			,  lowercase			,  **lowercase  ).prev_sample

																		assert torch.sum(torch.abs(output - new_output  )  ) < 1E-5, "Scheduler outputs are not identical"







			def 			__lowerCAmelCase				(			self			: int			,  lowercase			: int=0			,  **lowercase			: Tuple  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[str]      =				dict(self.forward_default_kwargs  )
								UpperCAmelCase			:   Union[str, Any]      =				kwargs.pop("num_inference_steps"			,  lowercase  )
								UpperCAmelCase			:   Tuple      =				self.dummy_sample
								UpperCAmelCase			:   Any      =				0.1 * sample
								UpperCAmelCase			:   int      =				[residual + 0.2, residual + 0.1_5, residual + 0.1_0]

								for scheduler_class in self.scheduler_classes:
													UpperCAmelCase			:   Optional[int]      =				self.get_scheduler_config()
													UpperCAmelCase			:   List[Any]      =				scheduler_class(**lowercase  )
													scheduler.set_timesteps(lowercase  )

													# copy over dummy past residuals (must be after setting timesteps)
													UpperCAmelCase			:   Dict      =				dummy_past_residuals[: scheduler.config.solver_order]

													with tempfile.TemporaryDirectory() as tmpdirname:
																		scheduler.save_config(lowercase  )
																		UpperCAmelCase			:   str      =				scheduler_class.from_pretrained(lowercase  )
																		# copy over dummy past residuals
																		new_scheduler.set_timesteps(lowercase  )

																		# copy over dummy past residual (must be after setting timesteps)
																		UpperCAmelCase			:   Tuple      =				dummy_past_residuals[: new_scheduler.config.solver_order]

													UpperCAmelCase			:   Dict      =				scheduler.step(lowercase			,  lowercase			,  lowercase			,  **lowercase  ).prev_sample
													UpperCAmelCase			:   Any      =				new_scheduler.step(lowercase			,  lowercase			,  lowercase			,  **lowercase  ).prev_sample

													assert torch.sum(torch.abs(output - new_output  )  ) < 1E-5, "Scheduler outputs are not identical"







			def 			__lowerCAmelCase				(			self			: Union[str, Any]			,  lowercase			: Optional[int]=None			,  **lowercase			: int  ):





								'''simple docstring'''



								if scheduler is None:
													UpperCAmelCase			:   Dict      =				self.scheduler_classes[0]
													UpperCAmelCase			:   int      =				self.get_scheduler_config(**lowercase  )
													UpperCAmelCase			:   str      =				scheduler_class(**lowercase  )

								UpperCAmelCase			:   Optional[int]      =				self.scheduler_classes[0]
								UpperCAmelCase			:   Dict      =				self.get_scheduler_config(**lowercase  )
								UpperCAmelCase			:   int      =				scheduler_class(**lowercase  )

								UpperCAmelCase			:   List[str]      =				10
								UpperCAmelCase			:   Optional[Any]      =				self.dummy_model()
								UpperCAmelCase			:   Union[str, Any]      =				self.dummy_sample_deter
								scheduler.set_timesteps(lowercase  )

								for i, t in enumerate(scheduler.timesteps  ):
													UpperCAmelCase			:   List[str]      =				model(lowercase			,  lowercase  )
													UpperCAmelCase			:   Optional[Any]      =				scheduler.step(lowercase			,  lowercase			,  lowercase  ).prev_sample

								return sample







			def 			__lowerCAmelCase				(			self			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Optional[Any]      =				dict(self.forward_default_kwargs  )

								UpperCAmelCase			:   List[Any]      =				kwargs.pop("num_inference_steps"			,  lowercase  )

								for scheduler_class in self.scheduler_classes:
													UpperCAmelCase			:   str      =				self.get_scheduler_config()
													UpperCAmelCase			:   Dict      =				scheduler_class(**lowercase  )

													UpperCAmelCase			:   Any      =				self.dummy_sample
													UpperCAmelCase			:   Any      =				0.1 * sample

													if num_inference_steps is not None and hasattr(lowercase			,  "set_timesteps"  ):
																		scheduler.set_timesteps(lowercase  )
													elif num_inference_steps is not None and not hasattr(lowercase			,  "set_timesteps"  ):
																		UpperCAmelCase			:   str      =				num_inference_steps

													# copy over dummy past residuals (must be done after set_timesteps)
													UpperCAmelCase			:   Union[str, Any]      =				[residual + 0.2, residual + 0.1_5, residual + 0.1_0]
													UpperCAmelCase			:   Dict      =				dummy_past_residuals[: scheduler.config.solver_order]

													UpperCAmelCase			:   str      =				scheduler.timesteps[5]
													UpperCAmelCase			:   List[Any]      =				scheduler.timesteps[6]

													UpperCAmelCase			:   Union[str, Any]      =				scheduler.step(lowercase			,  lowercase			,  lowercase			,  **lowercase  ).prev_sample
													UpperCAmelCase			:   int      =				scheduler.step(lowercase			,  lowercase			,  lowercase			,  **lowercase  ).prev_sample

													self.assertEqual(output_a.shape			,  sample.shape  )
													self.assertEqual(output_a.shape			,  output_a.shape  )







			def 			__lowerCAmelCase				(			self			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Union[str, Any]      =				UniPCMultistepScheduler(**self.get_scheduler_config()  )
								UpperCAmelCase			:   Tuple      =				self.full_loop(scheduler=lowercase  )
								UpperCAmelCase			:   Any      =				torch.mean(torch.abs(lowercase  )  )

								assert abs(result_mean.item() - 0.2_4_6_4  ) < 1E-3

								UpperCAmelCase			:   int      =				DPMSolverSinglestepScheduler.from_config(scheduler.config  )
								UpperCAmelCase			:   List[str]      =				DEISMultistepScheduler.from_config(scheduler.config  )
								UpperCAmelCase			:   Optional[int]      =				DPMSolverMultistepScheduler.from_config(scheduler.config  )
								UpperCAmelCase			:   int      =				UniPCMultistepScheduler.from_config(scheduler.config  )

								UpperCAmelCase			:   Dict      =				self.full_loop(scheduler=lowercase  )
								UpperCAmelCase			:   Any      =				torch.mean(torch.abs(lowercase  )  )

								assert abs(result_mean.item() - 0.2_4_6_4  ) < 1E-3







			def 			__lowerCAmelCase				(			self			: Any  ):





								'''simple docstring'''



								for timesteps in [25, 50, 1_00, 9_99, 10_00]:
													self.check_over_configs(num_train_timesteps=lowercase  )







			def 			__lowerCAmelCase				(			self			: Union[str, Any]  ):





								'''simple docstring'''



								self.check_over_configs(thresholding=lowercase  )
								for order in [1, 2, 3]:
													for solver_type in ["bh1", "bh2"]:
																		for threshold in [0.5, 1.0, 2.0]:
																							for prediction_type in ["epsilon", "sample"]:
																												self.check_over_configs(
																												    thresholding=lowercase			,  prediction_type=lowercase			,  sample_max_value=lowercase			,  solver_order=lowercase			,  solver_type=lowercase			,  )







			def 			__lowerCAmelCase				(			self			: Any  ):





								'''simple docstring'''



								for prediction_type in ["epsilon", "v_prediction"]:
													self.check_over_configs(prediction_type=lowercase  )







			def 			__lowerCAmelCase				(			self			: str  ):





								'''simple docstring'''



								for solver_type in ["bh1", "bh2"]:
													for order in [1, 2, 3]:
																		for prediction_type in ["epsilon", "sample"]:
																							self.check_over_configs(
																							    solver_order=lowercase			,  solver_type=lowercase			,  prediction_type=lowercase			,  )
																							UpperCAmelCase			:   Union[str, Any]      =				self.full_loop(
																							    solver_order=lowercase			,  solver_type=lowercase			,  prediction_type=lowercase			,  )
																							assert not torch.isnan(lowercase  ).any(), "Samples have nan numbers"







			def 			__lowerCAmelCase				(			self			: Optional[Any]  ):





								'''simple docstring'''



								self.check_over_configs(lower_order_final=lowercase  )
								self.check_over_configs(lower_order_final=lowercase  )







			def 			__lowerCAmelCase				(			self			: Tuple  ):





								'''simple docstring'''



								for num_inference_steps in [1, 2, 3, 5, 10, 50, 1_00, 9_99, 10_00]:
													self.check_over_forward(num_inference_steps=lowercase			,  time_step=0  )







			def 			__lowerCAmelCase				(			self			: Optional[Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Tuple      =				self.full_loop()
								UpperCAmelCase			:   Any      =				torch.mean(torch.abs(lowercase  )  )

								assert abs(result_mean.item() - 0.2_4_6_4  ) < 1E-3







			def 			__lowerCAmelCase				(			self			: Optional[int]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Tuple      =				self.full_loop(prediction_type="v_prediction"  )
								UpperCAmelCase			:   int      =				torch.mean(torch.abs(lowercase  )  )

								assert abs(result_mean.item() - 0.1_0_1_4  ) < 1E-3







			def 			__lowerCAmelCase				(			self			: Any  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[Any]      =				self.scheduler_classes[0]
								UpperCAmelCase			:   int      =				self.get_scheduler_config(thresholding=lowercase			,  dynamic_thresholding_ratio=0  )
								UpperCAmelCase			:   Tuple      =				scheduler_class(**lowercase  )

								UpperCAmelCase			:   Any      =				10
								UpperCAmelCase			:   Dict      =				self.dummy_model()
								UpperCAmelCase			:   List[str]      =				self.dummy_sample_deter.half()
								scheduler.set_timesteps(lowercase  )

								for i, t in enumerate(scheduler.timesteps  ):
													UpperCAmelCase			:   Any      =				model(lowercase			,  lowercase  )
													UpperCAmelCase			:   List[Any]      =				scheduler.step(lowercase			,  lowercase			,  lowercase  ).prev_sample

								assert sample.dtype == torch.floataa







			def 			__lowerCAmelCase				(			self			: List[str]			,  **lowercase			: Optional[int]  ):





								'''simple docstring'''



								for scheduler_class in self.scheduler_classes:
													UpperCAmelCase			:   Union[str, Any]      =				self.get_scheduler_config(**lowercase  )
													UpperCAmelCase			:   Union[str, Any]      =				scheduler_class(**lowercase  )

													scheduler.set_timesteps(scheduler.config.num_train_timesteps  )
													assert len(scheduler.timesteps.unique()  ) == scheduler.num_inference_steps


