






"""simple docstring"""



import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class        __lowercase(       _UpperCamelCase  ):
      lowerCAmelCase_		:							List[str]=       (HeunDiscreteScheduler,)
      lowerCAmelCase_		:							List[Any]=       10






      def 		lowerCAmelCase   (					self    ,						**_SCREAMING_SNAKE_CASE)			->  List[Any]:




           """simple docstring"""





           lowerCamelCase__				={
               'num_train_timesteps': 11_00,
               'beta_start': 0.0001,
               'beta_end': 0.02,
               'beta_schedule': 'linear',
           }

           config.update(**_SCREAMING_SNAKE_CASE)
           return config






      def 		lowerCAmelCase   (					self)			->  str:




           """simple docstring"""





           for timesteps in [10, 50, 1_00, 10_00]:
                self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE)






      def 		lowerCAmelCase   (					self)			->  List[str]:




           """simple docstring"""





           for beta_start, beta_end in zip([0.00001, 0.0001, 0.001]    ,						[0.0002, 0.002, 0.02]):
                self.check_over_configs(beta_start=_SCREAMING_SNAKE_CASE    ,						beta_end=_SCREAMING_SNAKE_CASE)






      def 		lowerCAmelCase   (					self)			->  Tuple:




           """simple docstring"""





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






      def 		lowerCAmelCase   (					self)			->  str:




           """simple docstring"""





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






      def 		lowerCAmelCase   (					self)			->  Optional[int]:




           """simple docstring"""





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






      def 		lowerCAmelCase   (					self)			->  Union[str, Any]:




           """simple docstring"""





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

           scheduler.set_timesteps(self.num_inference_steps)

           lowerCamelCase__				=self.dummy_model()
           lowerCamelCase__				=self.dummy_sample_deter * scheduler.init_noise_sigma
           lowerCamelCase__				=sample.to(_SCREAMING_SNAKE_CASE)

           for i, t in enumerate(scheduler.timesteps):
                lowerCamelCase__				=scheduler.scale_model_input(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=model(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=scheduler.step(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)
                lowerCamelCase__				=output.prev_sample

           lowerCamelCase__				=torch.sum(torch.abs(_SCREAMING_SNAKE_CASE))
           lowerCamelCase__				=torch.mean(torch.abs(_SCREAMING_SNAKE_CASE))

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






      def 		lowerCAmelCase   (					self)			->  Any:




           """simple docstring"""





           lowerCamelCase__				=self.scheduler_classes[0]
           lowerCamelCase__				=self.get_scheduler_config(prediction_type='v_prediction')
           lowerCamelCase__				=scheduler_class(**_SCREAMING_SNAKE_CASE)

           scheduler.set_timesteps(self.num_inference_steps)

           lowerCamelCase__				=self.dummy_model()
           lowerCamelCase__				=self.dummy_sample_deter * scheduler.init_noise_sigma
           lowerCamelCase__				=sample.to(_SCREAMING_SNAKE_CASE)

           for i, t in enumerate(scheduler.timesteps):
                lowerCamelCase__				=scheduler.scale_model_input(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=model(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=scheduler.step(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)
                lowerCamelCase__				=output.prev_sample

           lowerCamelCase__				=torch.sum(torch.abs(_SCREAMING_SNAKE_CASE))
           lowerCamelCase__				=torch.mean(torch.abs(_SCREAMING_SNAKE_CASE))

           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.0002) < 1E-3






      def 		lowerCAmelCase   (					self)			->  Union[str, Any]:




           """simple docstring"""





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

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

           lowerCamelCase__				=self.dummy_model()
           lowerCamelCase__				=self.dummy_sample_deter.to(_SCREAMING_SNAKE_CASE) * scheduler.init_noise_sigma

           for t in scheduler.timesteps:
                lowerCamelCase__				=scheduler.scale_model_input(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=model(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=scheduler.step(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)
                lowerCamelCase__				=output.prev_sample

           lowerCamelCase__				=torch.sum(torch.abs(_SCREAMING_SNAKE_CASE))
           lowerCamelCase__				=torch.mean(torch.abs(_SCREAMING_SNAKE_CASE))

           if str(_SCREAMING_SNAKE_CASE).startswith('cpu'):
                # The following sum varies between 148 and 156 on mps. Why?
                assert abs(result_sum.item() - 0.1233) < 1E-2
                assert abs(result_mean.item() - 0.0002) < 1E-3
           elif str(_SCREAMING_SNAKE_CASE).startswith('mps'):
                # Larger tolerance on mps
                assert abs(result_mean.item() - 0.0002) < 1E-2
           else:
                # CUDA
                assert abs(result_sum.item() - 0.1233) < 1E-2
                assert abs(result_mean.item() - 0.0002) < 1E-3






      def 		lowerCAmelCase   (					self)			->  int:




           """simple docstring"""





           lowerCamelCase__				=self.scheduler_classes[0]
           lowerCamelCase__				=self.get_scheduler_config()
           lowerCamelCase__				=scheduler_class(**_SCREAMING_SNAKE_CASE    ,						use_karras_sigmas=_SCREAMING_SNAKE_CASE)

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

           lowerCamelCase__				=self.dummy_model()
           lowerCamelCase__				=self.dummy_sample_deter.to(_SCREAMING_SNAKE_CASE) * scheduler.init_noise_sigma
           lowerCamelCase__				=sample.to(_SCREAMING_SNAKE_CASE)

           for t in scheduler.timesteps:
                lowerCamelCase__				=scheduler.scale_model_input(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=model(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

                lowerCamelCase__				=scheduler.step(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)
                lowerCamelCase__				=output.prev_sample

           lowerCamelCase__				=torch.sum(torch.abs(_SCREAMING_SNAKE_CASE))
           lowerCamelCase__				=torch.mean(torch.abs(_SCREAMING_SNAKE_CASE))

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







"""simple docstring"""



from doctest import testmod
from math import sqrt




def 						lowerCamelCase_				(  __A  :		int			):
					lowerCamelCase__				=[]
					if num < 1:
										return facs
					facs.append(1			)
					if num == 1:
										return facs
					facs.append(__A			)
					for i in range(2			,	int(sqrt(__A			)			) + 1			):
										if num % i == 0:  # If i is a factor of num
															facs.append(__A			)
															lowerCamelCase__				=num // i  # num//i is the other factor of num
															if d != i:  # If d and i are distinct
																				facs.append(__A			)  # we have found another factor
					facs.sort()
					return facs


if __name__ == "__main__":
						testmod(name="factors_of_a_number", verbose=True)