


'''simple docstring'''



import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest





class 							__UpperCamelCase    (		lowerCAmelCase_ ):
	A_           =   (DEISMultistepScheduler,)
	A_           =   (("num_inference_steps", 25),)




	def 				__UpperCAmelCase							(       self						,    **__a							):

						'''simple docstring'''





						__a      :	Union[str, Any]    =  {
						    'num_train_timesteps': 1000,
						    'beta_start': 0.0001,
						    'beta_end': 0.02,
						    'beta_schedule': 'linear',
						    'solver_order': 2,
						}

						config.update(**__a							)
						return config




	def 				__UpperCAmelCase							(       self						,    __a=0						,    **__a							):

						'''simple docstring'''





						__a      :	Optional[int]    =  dict(self.forward_default_kwargs							)
						__a      :	Dict    =  kwargs.pop('num_inference_steps'						,    __a							)
						__a      :	List[Any]    =  self.dummy_sample
						__a      :	Tuple    =  0.1 * sample
						__a      :	List[Any]    =  [residual + 0.2, residual + 0.15, residual + 0.10]

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

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

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

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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						pass




	def 				__UpperCAmelCase							(       self						,    __a=0						,    **__a							):

						'''simple docstring'''





						__a      :	Dict    =  dict(self.forward_default_kwargs							)
						__a      :	Tuple    =  kwargs.pop('num_inference_steps'						,    __a							)
						__a      :	Optional[Any]    =  self.dummy_sample
						__a      :	Any    =  0.1 * sample
						__a      :	Tuple    =  [residual + 0.2, residual + 0.15, residual + 0.10]

						for scheduler_class in self.scheduler_classes:
											__a      :	Any    =  self.get_scheduler_config()
											__a      :	int    =  scheduler_class(**__a							)
											scheduler.set_timesteps(__a							)

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

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

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

											__a      :	List[str]    =  scheduler.step(__a						,    __a						,    __a						,    **__a							).prev_sample
											__a      :	Tuple    =  new_scheduler.step(__a						,    __a						,    __a						,    **__a							).prev_sample

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




	def 				__UpperCAmelCase							(       self						,    __a=None						,    **__a							):

						'''simple docstring'''





						if scheduler is None:
											__a      :	str    =  self.scheduler_classes[0]
											__a      :	Optional[Any]    =  self.get_scheduler_config(**__a							)
											__a      :	str    =  scheduler_class(**__a							)

						__a      :	Optional[int]    =  self.scheduler_classes[0]
						__a      :	Any    =  self.get_scheduler_config(**__a							)
						__a      :	Optional[int]    =  scheduler_class(**__a							)

						__a      :	List[Any]    =  10
						__a      :	Any    =  self.dummy_model()
						__a      :	Optional[Any]    =  self.dummy_sample_deter
						scheduler.set_timesteps(__a							)

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

						return sample




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	int    =  dict(self.forward_default_kwargs							)

						__a      :	Optional[int]    =  kwargs.pop('num_inference_steps'						,    __a							)

						for scheduler_class in self.scheduler_classes:
											__a      :	Any    =  self.get_scheduler_config()
											__a      :	Optional[int]    =  scheduler_class(**__a							)

											__a      :	List[Any]    =  self.dummy_sample
											__a      :	List[Any]    =  0.1 * sample

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

											# copy over dummy past residuals (must be done after set_timesteps)
											__a      :	Dict    =  [residual + 0.2, residual + 0.15, residual + 0.10]
											__a      :	List[Any]    =  dummy_past_residuals[: scheduler.config.solver_order]

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

											__a      :	Optional[Any]    =  scheduler.step(__a						,    __a						,    __a						,    **__a							).prev_sample
											__a      :	str    =  scheduler.step(__a						,    __a						,    __a						,    **__a							).prev_sample

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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Dict    =  DEISMultistepScheduler(**self.get_scheduler_config()							)
						__a      :	Optional[int]    =  self.full_loop(scheduler=__a							)
						__a      :	Any    =  torch.mean(torch.abs(__a							)							)

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

						__a      :	Dict    =  DPMSolverSinglestepScheduler.from_config(scheduler.config							)
						__a      :	Dict    =  DPMSolverMultistepScheduler.from_config(scheduler.config							)
						__a      :	int    =  UniPCMultistepScheduler.from_config(scheduler.config							)
						__a      :	Optional[int]    =  DEISMultistepScheduler.from_config(scheduler.config							)

						__a      :	str    =  self.full_loop(scheduler=__a							)
						__a      :	Any    =  torch.mean(torch.abs(__a							)							)

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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						for timesteps in [25, 50, 100, 999, 1000]:
											self.check_over_configs(num_train_timesteps=__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						self.check_over_configs(thresholding=__a							)
						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=__a						,    prediction_type=__a						,    sample_max_value=__a						,    algorithm_type='deis'						,    solver_order=__a						,    solver_type=__a						,    )




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





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




	def 				__UpperCAmelCase							(       self							):

						'''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=__a						,    solver_type=__a						,    prediction_type=__a						,    algorithm_type=__a						,    )
																										__a      :	int    =  self.full_loop(
																										    solver_order=__a						,    solver_type=__a						,    prediction_type=__a						,    algorithm_type=__a						,    )
																										assert not torch.isnan(__a							).any(), "Samples have nan numbers"




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						self.check_over_configs(lower_order_final=__a							)
						self.check_over_configs(lower_order_final=__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[int]    =  self.full_loop()
						__a      :	Optional[Any]    =  torch.mean(torch.abs(__a							)							)

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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	List[Any]    =  self.full_loop(prediction_type='v_prediction'							)
						__a      :	Any    =  torch.mean(torch.abs(__a							)							)

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





	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	List[str]    =  self.scheduler_classes[0]
						__a      :	str    =  self.get_scheduler_config(thresholding=__a						,    dynamic_thresholding_ratio=0							)
						__a      :	Dict    =  scheduler_class(**__a							)

						__a      :	Any    =  10
						__a      :	Tuple    =  self.dummy_model()
						__a      :	Tuple    =  self.dummy_sample_deter.half()
						scheduler.set_timesteps(__a							)

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

						assert sample.dtype == torch.floataa







'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowercase     :						Tuple	=     {
    'configuration_distilbert': [
        'DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'DistilBertConfig',
        'DistilBertOnnxConfig',
    ],
    'tokenization_distilbert': ['DistilBertTokenizer'],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						str	=     ['DistilBertTokenizerFast']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						Any	=     [
	    'DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'DistilBertForMaskedLM',
	    'DistilBertForMultipleChoice',
	    'DistilBertForQuestionAnswering',
	    'DistilBertForSequenceClassification',
	    'DistilBertForTokenClassification',
	    'DistilBertModel',
	    'DistilBertPreTrainedModel',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						List[str]	=     [
	    'TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'TFDistilBertForMaskedLM',
	    'TFDistilBertForMultipleChoice',
	    'TFDistilBertForQuestionAnswering',
	    'TFDistilBertForSequenceClassification',
	    'TFDistilBertForTokenClassification',
	    'TFDistilBertMainLayer',
	    'TFDistilBertModel',
	    'TFDistilBertPreTrainedModel',
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						List[str]	=     [
	    'FlaxDistilBertForMaskedLM',
	    'FlaxDistilBertForMultipleChoice',
	    'FlaxDistilBertForQuestionAnswering',
	    'FlaxDistilBertForSequenceClassification',
	    'FlaxDistilBertForTokenClassification',
	    'FlaxDistilBertModel',
	    'FlaxDistilBertPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_distilbert import (
	    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    DistilBertConfig,
	    DistilBertOnnxConfig,
	)
	from .tokenization_distilbert import DistilBertTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_distilbert_fast import DistilBertTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_distilbert import (
		    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DistilBertForMaskedLM,
		    DistilBertForMultipleChoice,
		    DistilBertForQuestionAnswering,
		    DistilBertForSequenceClassification,
		    DistilBertForTokenClassification,
		    DistilBertModel,
		    DistilBertPreTrainedModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_distilbert import (
		    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFDistilBertForMaskedLM,
		    TFDistilBertForMultipleChoice,
		    TFDistilBertForQuestionAnswering,
		    TFDistilBertForSequenceClassification,
		    TFDistilBertForTokenClassification,
		    TFDistilBertMainLayer,
		    TFDistilBertModel,
		    TFDistilBertPreTrainedModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_distilbert import (
		    FlaxDistilBertForMaskedLM,
		    FlaxDistilBertForMultipleChoice,
		    FlaxDistilBertForQuestionAnswering,
		    FlaxDistilBertForSequenceClassification,
		    FlaxDistilBertForTokenClassification,
		    FlaxDistilBertModel,
		    FlaxDistilBertPreTrainedModel,
		)

else:
	import sys

	__lowercase     :						int	=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



