




# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import doctest
import sys
import warnings
from os.path import abspath, dirname, join

import _pytest

from transformers.testing_utils import HfDoctestModule, HfDocTestParser


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
UpperCAmelCase__				     = abspath(join(dirname(__file__), "src"))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action="ignore", category=FutureWarning)







def 		_A(       UpperCamelCase__							:			List[Any]				)    -> Dict:
							'''simple docstring'''


							config.addinivalue_line(
							    '''markers'''				,       '''is_pt_tf_cross_test: mark test to run only when PT and TF interactions are tested'''				)
							config.addinivalue_line(
							    '''markers'''				,       '''is_pt_flax_cross_test: mark test to run only when PT and FLAX interactions are tested'''				)
							config.addinivalue_line('''markers'''				,       '''is_pipeline_test: mark test to run only when pipelines are tested'''				)
							config.addinivalue_line('''markers'''				,       '''is_staging_test: mark test to run only in the staging environment'''				)
							config.addinivalue_line('''markers'''				,       '''accelerate_tests: mark test that require accelerate'''				)
							config.addinivalue_line('''markers'''				,       '''tool_tests: mark the tool tests that are run on their specific schedule'''				)







def 		_A(       UpperCamelCase__							:			Union[str, Any]				)    -> Union[str, Any]:
							'''simple docstring'''


							from transformers.testing_utils import pytest_addoption_shared

							pytest_addoption_shared(UpperCamelCase__				)







def 		_A(       UpperCamelCase__							:			List[Any]				)    -> Any:
							'''simple docstring'''


							from transformers.testing_utils import pytest_terminal_summary_main

							__lowercase      =    terminalreporter.config.getoption('''--make-reports'''				)
							if make_reports:
														pytest_terminal_summary_main(UpperCamelCase__				,       id=UpperCamelCase__				)







def 		_A(       UpperCamelCase__							:			List[Any]				,       UpperCamelCase__							:			Union[str, Any]				)    -> List[str]:
							'''simple docstring'''


							if exitstatus == 5:
														__lowercase      =    0


# Doctest custom flag to ignore output.
UpperCAmelCase__				     = doctest.register_optionflag("IGNORE_RESULT")

UpperCAmelCase__				     = doctest.OutputChecker







class a (     __SCREAMING_SNAKE_CASE  ):
						"""simple docstring"""

						def     UpperCAmelCase_      ( self :     Tuple				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     Optional[int]				,							lowerCamelCase__ :     Union[str, Any]				)    ->						Union[str, Any]:
													"""simple docstring"""


													if IGNORE_RESULT & optionflags:
																				return True
													return OutputChecker.check_output(self				,							lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				)





UpperCAmelCase__				     = CustomOutputChecker
UpperCAmelCase__				     = HfDoctestModule



UpperCAmelCase__				     = HfDocTestParser





import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest







class a (     __SCREAMING_SNAKE_CASE  ):
						"""simple docstring"""

						UpperCamelCase_	:       List[Any]       	=		(DEISMultistepScheduler,)
						UpperCamelCase_	:       Any       	=		(('num_inference_steps', 25),)






						def     UpperCAmelCase_      ( self :     List[str]				,							**lowerCamelCase__ :     Union[str, Any]				)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    {
													    '''num_train_timesteps''': 1_000,
													    '''beta_start''': 0.0_0_0_1,
													    '''beta_end''': 0.0_2,
													    '''beta_schedule''': '''linear''',
													    '''solver_order''': 2,
													}

													config.update(**lowerCamelCase__				)
													return config






						def     UpperCAmelCase_      ( self :     List[Any]				,							lowerCamelCase__ :     Union[str, Any]=0				,							**lowerCamelCase__ :     Dict				)    ->						Any:
													"""simple docstring"""


													__lowercase      =    dict(self.forward_default_kwargs				)
													__lowercase      =    kwargs.pop('''num_inference_steps'''				,							lowerCamelCase__				)
													__lowercase      =    self.dummy_sample
													__lowercase      =    0.1 * sample
													__lowercase      =    [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

													for scheduler_class in self.scheduler_classes:
																				__lowercase      =    self.get_scheduler_config(**lowerCamelCase__				)
																				__lowercase      =    scheduler_class(**lowerCamelCase__				)
																				scheduler.set_timesteps(lowerCamelCase__				)
																				# copy over dummy past residuals
																				__lowercase      =    dummy_past_residuals[: scheduler.config.solver_order]

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

																				__lowercase     ,     __lowercase      =    sample, sample
																				for t in range(lowerCamelCase__				,							time_step + scheduler.config.solver_order + 1				):
																											__lowercase      =    scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							**lowerCamelCase__				).prev_sample
																											__lowercase      =    new_scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							**lowerCamelCase__				).prev_sample

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






						def     UpperCAmelCase_      ( self :     Any				)    ->						Tuple:
													"""simple docstring"""


													pass






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Any=0				,							**lowerCamelCase__ :     Any				)    ->						Any:
													"""simple docstring"""


													__lowercase      =    dict(self.forward_default_kwargs				)
													__lowercase      =    kwargs.pop('''num_inference_steps'''				,							lowerCamelCase__				)
													__lowercase      =    self.dummy_sample
													__lowercase      =    0.1 * sample
													__lowercase      =    [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

													for scheduler_class in self.scheduler_classes:
																				__lowercase      =    self.get_scheduler_config()
																				__lowercase      =    scheduler_class(**lowerCamelCase__				)
																				scheduler.set_timesteps(lowerCamelCase__				)

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

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

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

																				__lowercase      =    scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							**lowerCamelCase__				).prev_sample
																				__lowercase      =    new_scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							**lowerCamelCase__				).prev_sample

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






						def     UpperCAmelCase_      ( self :     Union[str, Any]				,							lowerCamelCase__ :     Optional[int]=None				,							**lowerCamelCase__ :     Dict				)    ->						Optional[Any]:
													"""simple docstring"""


													if scheduler is None:
																				__lowercase      =    self.scheduler_classes[0]
																				__lowercase      =    self.get_scheduler_config(**lowerCamelCase__				)
																				__lowercase      =    scheduler_class(**lowerCamelCase__				)

													__lowercase      =    self.scheduler_classes[0]
													__lowercase      =    self.get_scheduler_config(**lowerCamelCase__				)
													__lowercase      =    scheduler_class(**lowerCamelCase__				)

													__lowercase      =    10
													__lowercase      =    self.dummy_model()
													__lowercase      =    self.dummy_sample_deter
													scheduler.set_timesteps(lowerCamelCase__				)

													for i, t in enumerate(scheduler.timesteps				):
																				__lowercase      =    model(lowerCamelCase__				,							lowerCamelCase__				)
																				__lowercase      =    scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				).prev_sample

													return sample






						def     UpperCAmelCase_      ( self :     Tuple				)    ->						Optional[Any]:
													"""simple docstring"""


													__lowercase      =    dict(self.forward_default_kwargs				)

													__lowercase      =    kwargs.pop('''num_inference_steps'''				,							lowerCamelCase__				)

													for scheduler_class in self.scheduler_classes:
																				__lowercase      =    self.get_scheduler_config()
																				__lowercase      =    scheduler_class(**lowerCamelCase__				)

																				__lowercase      =    self.dummy_sample
																				__lowercase      =    0.1 * sample

																				if num_inference_steps is not None and hasattr(lowerCamelCase__				,							'''set_timesteps'''				):
																											scheduler.set_timesteps(lowerCamelCase__				)
																				elif num_inference_steps is not None and not hasattr(lowerCamelCase__				,							'''set_timesteps'''				):
																											__lowercase      =    num_inference_steps

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

																				__lowercase      =    scheduler.timesteps[5]
																				__lowercase      =    scheduler.timesteps[6]

																				__lowercase      =    scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							**lowerCamelCase__				).prev_sample
																				__lowercase      =    scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							**lowerCamelCase__				).prev_sample

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






						def     UpperCAmelCase_      ( self :     str				)    ->						List[Any]:
													"""simple docstring"""


													__lowercase      =    DEISMultistepScheduler(**self.get_scheduler_config()				)
													__lowercase      =    self.full_loop(scheduler=lowerCamelCase__				)
													__lowercase      =    torch.mean(torch.abs(lowerCamelCase__				)				)

													assert abs(result_mean.item() - 0.2_3_9_1_6				) < 1e-3

													__lowercase      =    DPMSolverSinglestepScheduler.from_config(scheduler.config				)
													__lowercase      =    DPMSolverMultistepScheduler.from_config(scheduler.config				)
													__lowercase      =    UniPCMultistepScheduler.from_config(scheduler.config				)
													__lowercase      =    DEISMultistepScheduler.from_config(scheduler.config				)

													__lowercase      =    self.full_loop(scheduler=lowerCamelCase__				)
													__lowercase      =    torch.mean(torch.abs(lowerCamelCase__				)				)

													assert abs(result_mean.item() - 0.2_3_9_1_6				) < 1e-3






						def     UpperCAmelCase_      ( self :     int				)    ->						List[Any]:
													"""simple docstring"""


													for timesteps in [25, 50, 100, 999, 1_000]:
																				self.check_over_configs(num_train_timesteps=lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     Union[str, Any]				)    ->						Any:
													"""simple docstring"""


													self.check_over_configs(thresholding=lowerCamelCase__				)
													for order in [1, 2, 3]:
																				for solver_type in ["logrho"]:
																											for threshold in [0.5, 1.0, 2.0]:
																																		for prediction_type in ["epsilon", "sample"]:
																																									self.check_over_configs(
																																									    thresholding=lowerCamelCase__				,							prediction_type=lowerCamelCase__				,							sample_max_value=lowerCamelCase__				,							algorithm_type='''deis'''				,							solver_order=lowerCamelCase__				,							solver_type=lowerCamelCase__				,							)






						def     UpperCAmelCase_      ( self :     List[str]				)    ->						Optional[Any]:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Optional[Any]				)    ->						int:
													"""simple docstring"""


													for algorithm_type in ["deis"]:
																				for solver_type in ["logrho"]:
																											for order in [1, 2, 3]:
																																		for prediction_type in ["epsilon", "sample"]:
																																									self.check_over_configs(
																																									    solver_order=lowerCamelCase__				,							solver_type=lowerCamelCase__				,							prediction_type=lowerCamelCase__				,							algorithm_type=lowerCamelCase__				,							)
																																									__lowercase      =    self.full_loop(
																																									    solver_order=lowerCamelCase__				,							solver_type=lowerCamelCase__				,							prediction_type=lowerCamelCase__				,							algorithm_type=lowerCamelCase__				,							)
																																									assert not torch.isnan(lowerCamelCase__				).any(), "Samples have nan numbers"






						def     UpperCAmelCase_      ( self :     Optional[Any]				)    ->						List[Any]:
													"""simple docstring"""


													self.check_over_configs(lower_order_final=lowerCamelCase__				)
													self.check_over_configs(lower_order_final=lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     List[Any]				)    ->						Dict:
													"""simple docstring"""


													for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1_000]:
																				self.check_over_forward(num_inference_steps=lowerCamelCase__				,							time_step=0				)






						def     UpperCAmelCase_      ( self :     Dict				)    ->						Tuple:
													"""simple docstring"""


													__lowercase      =    self.full_loop()
													__lowercase      =    torch.mean(torch.abs(lowerCamelCase__				)				)

													assert abs(result_mean.item() - 0.2_3_9_1_6				) < 1e-3






						def     UpperCAmelCase_      ( self :     Tuple				)    ->						Tuple:
													"""simple docstring"""


													__lowercase      =    self.full_loop(prediction_type='''v_prediction'''				)
													__lowercase      =    torch.mean(torch.abs(lowerCamelCase__				)				)

													assert abs(result_mean.item() - 0.0_9_1				) < 1e-3






						def     UpperCAmelCase_      ( self :     Optional[Any]				)    ->						Tuple:
													"""simple docstring"""


													__lowercase      =    self.scheduler_classes[0]
													__lowercase      =    self.get_scheduler_config(thresholding=lowerCamelCase__				,							dynamic_thresholding_ratio=0				)
													__lowercase      =    scheduler_class(**lowerCamelCase__				)

													__lowercase      =    10
													__lowercase      =    self.dummy_model()
													__lowercase      =    self.dummy_sample_deter.half()
													scheduler.set_timesteps(lowerCamelCase__				)

													for i, t in enumerate(scheduler.timesteps				):
																				__lowercase      =    model(lowerCamelCase__				,							lowerCamelCase__				)
																				__lowercase      =    scheduler.step(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				).prev_sample

													assert sample.dtype == torch.floataa






