




'''simple docstring'''







import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class 					lowerCamelCase_							(	UpperCAmelCase_							):




					'''simple docstring'''






					lowercase__   =			(HeunDiscreteScheduler,)
					lowercase__   =			1_0






					def     __SCREAMING_SNAKE_CASE     (   self,  **__UpperCamelCase			):
							_snake_case     	=	{
							    "num_train_timesteps": 1_100,
							    "beta_start": 0.00_01,
							    "beta_end": 0.02,
							    "beta_schedule": "linear",
							}

							config.update(**__UpperCamelCase			)
							return config






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for timesteps in [10, 50, 100, 1_000]:
									self.check_over_configs(num_train_timesteps=__UpperCamelCase			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01],  [0.00_02, 0.0_02, 0.02]			):
									self.check_over_configs(beta_start=__UpperCamelCase,  beta_end=__UpperCamelCase			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for schedule in ["linear", "scaled_linear", "exp"]:
									self.check_over_configs(beta_schedule=__UpperCamelCase			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for clip_sample_range in [1.0, 2.0, 3.0]:
									self.check_over_configs(clip_sample_range=__UpperCamelCase,  clip_sample=__UpperCamelCase			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							for prediction_type in ["epsilon", "v_prediction", "sample"]:
									self.check_over_configs(prediction_type=__UpperCamelCase			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.scheduler_classes[0]
							_snake_case     	=	self.get_scheduler_config()
							_snake_case     	=	scheduler_class(**__UpperCamelCase			)

							scheduler.set_timesteps(self.num_inference_steps			)

							_snake_case     	=	self.dummy_model()
							_snake_case     	=	self.dummy_sample_deter * scheduler.init_noise_sigma
							_snake_case     	=	sample.to(__UpperCamelCase			)

							for i, t in enumerate(scheduler.timesteps			):
									_snake_case     	=	scheduler.scale_model_input(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	model(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	scheduler.step(__UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)
									_snake_case     	=	output.prev_sample

							_snake_case     	=	torch.sum(torch.abs(__UpperCamelCase			)			)
							_snake_case     	=	torch.mean(torch.abs(__UpperCamelCase			)			)

							if torch_device in ["cpu", "mps"]:
									assert abs(result_sum.item() - 0.12_33			) < 1E-2
									assert abs(result_mean.item() - 0.00_02			) < 1E-3
							else:
									# CUDA
									assert abs(result_sum.item() - 0.12_33			) < 1E-2
									assert abs(result_mean.item() - 0.00_02			) < 1E-3






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.scheduler_classes[0]
							_snake_case     	=	self.get_scheduler_config(prediction_type="v_prediction"			)
							_snake_case     	=	scheduler_class(**__UpperCamelCase			)

							scheduler.set_timesteps(self.num_inference_steps			)

							_snake_case     	=	self.dummy_model()
							_snake_case     	=	self.dummy_sample_deter * scheduler.init_noise_sigma
							_snake_case     	=	sample.to(__UpperCamelCase			)

							for i, t in enumerate(scheduler.timesteps			):
									_snake_case     	=	scheduler.scale_model_input(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	model(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	scheduler.step(__UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)
									_snake_case     	=	output.prev_sample

							_snake_case     	=	torch.sum(torch.abs(__UpperCamelCase			)			)
							_snake_case     	=	torch.mean(torch.abs(__UpperCamelCase			)			)

							if torch_device in ["cpu", "mps"]:
									assert abs(result_sum.item() - 4.6934E-07			) < 1E-2
									assert abs(result_mean.item() - 6.1112E-10			) < 1E-3
							else:
									# CUDA
									assert abs(result_sum.item() - 4.693428650170972E-07			) < 1E-2
									assert abs(result_mean.item() - 0.00_02			) < 1E-3






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.scheduler_classes[0]
							_snake_case     	=	self.get_scheduler_config()
							_snake_case     	=	scheduler_class(**__UpperCamelCase			)

							scheduler.set_timesteps(self.num_inference_steps,  device=__UpperCamelCase			)

							_snake_case     	=	self.dummy_model()
							_snake_case     	=	self.dummy_sample_deter.to(__UpperCamelCase			) * scheduler.init_noise_sigma

							for t in scheduler.timesteps:
									_snake_case     	=	scheduler.scale_model_input(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	model(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	scheduler.step(__UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)
									_snake_case     	=	output.prev_sample

							_snake_case     	=	torch.sum(torch.abs(__UpperCamelCase			)			)
							_snake_case     	=	torch.mean(torch.abs(__UpperCamelCase			)			)

							if str(__UpperCamelCase			).startswith("cpu"			):
									# The following sum varies between 148 and 156 on mps. Why?
									assert abs(result_sum.item() - 0.12_33			) < 1E-2
									assert abs(result_mean.item() - 0.00_02			) < 1E-3
							elif str(__UpperCamelCase			).startswith("mps"			):
									# Larger tolerance on mps
									assert abs(result_mean.item() - 0.00_02			) < 1E-2
							else:
									# CUDA
									assert abs(result_sum.item() - 0.12_33			) < 1E-2
									assert abs(result_mean.item() - 0.00_02			) < 1E-3






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.scheduler_classes[0]
							_snake_case     	=	self.get_scheduler_config()
							_snake_case     	=	scheduler_class(**__UpperCamelCase,  use_karras_sigmas=__UpperCamelCase			)

							scheduler.set_timesteps(self.num_inference_steps,  device=__UpperCamelCase			)

							_snake_case     	=	self.dummy_model()
							_snake_case     	=	self.dummy_sample_deter.to(__UpperCamelCase			) * scheduler.init_noise_sigma
							_snake_case     	=	sample.to(__UpperCamelCase			)

							for t in scheduler.timesteps:
									_snake_case     	=	scheduler.scale_model_input(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	model(__UpperCamelCase,  __UpperCamelCase			)

									_snake_case     	=	scheduler.step(__UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)
									_snake_case     	=	output.prev_sample

							_snake_case     	=	torch.sum(torch.abs(__UpperCamelCase			)			)
							_snake_case     	=	torch.mean(torch.abs(__UpperCamelCase			)			)

							assert abs(result_sum.item() - 0.0_00_15			) < 1E-2
							assert abs(result_mean.item() - 1.9869554535034695E-07			) < 1E-2











'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_lowerCamelCase							:						List[Any]							     =  {
    'configuration_mobilenet_v1': [
        'MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'MobileNetV1Config',
        'MobileNetV1OnnxConfig',
    ],
}

try:
						if not is_vision_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:						str							     =  ['MobileNetV1FeatureExtractor']
						_lowerCamelCase							:						Tuple							     =  ['MobileNetV1ImageProcessor']

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:						List[Any]							     =  [
						    'MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'MobileNetV1ForImageClassification',
						    'MobileNetV1Model',
						    'MobileNetV1PreTrainedModel',
						    'load_tf_weights_in_mobilenet_v1',
						]


if TYPE_CHECKING:
						from .configuration_mobilenet_va import (
						    MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    MobileNetVaConfig,
						    MobileNetVaOnnxConfig,
						)

						try:
												if not is_vision_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .feature_extraction_mobilenet_va import MobileNetVaFeatureExtractor
												from .image_processing_mobilenet_va import MobileNetVaImageProcessor

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_mobilenet_va import (
												    MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST,
												    MobileNetVaForImageClassification,
												    MobileNetVaModel,
												    MobileNetVaPreTrainedModel,
												    load_tf_weights_in_mobilenet_va,
												)


else:
						import sys

						_lowerCamelCase							:						Union[str, Any]							     =  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





