





import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class 						_lowerCAmelCase(						_a):

  """simple docstring"""







  lowerCamelCase__      					=   (HeunDiscreteScheduler,)
  lowerCamelCase__      					=   10

  def 		SCREAMING_SNAKE_CASE__	( self  ,    **UpperCAmelCase     )-> int:
     __A				 =       {
         '''num_train_timesteps''': 11_00,
         '''beta_start''': 0.0001,
         '''beta_end''': 0.02,
         '''beta_schedule''': '''linear''',
     }

     config.update(**UpperCAmelCase     )
     return config

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[int]:
     for timesteps in [10, 50, 1_00, 10_00]:
        self.check_over_configs(num_train_timesteps=UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     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=UpperCAmelCase  ,    beta_end=UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     for schedule in ["linear", "scaled_linear", "exp"]:
        self.check_over_configs(beta_schedule=UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     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     )-> Optional[Any]:
     for prediction_type in ["epsilon", "v_prediction", "sample"]:
        self.check_over_configs(prediction_type=UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     __A				 =       self.scheduler_classes[0]
     __A				 =       self.get_scheduler_config()
     __A				 =       scheduler_class(**UpperCAmelCase     )

     scheduler.set_timesteps(self.num_inference_steps     )

     __A				 =       self.dummy_model()
     __A				 =       self.dummy_sample_deter * scheduler.init_noise_sigma
     __A				 =       sample.to(UpperCAmelCase     )

     for i, t in enumerate(scheduler.timesteps     ):
        __A				 =       scheduler.scale_model_input(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       model(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       scheduler.step(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
        __A				 =       output.prev_sample

     __A				 =       torch.sum(torch.abs(UpperCAmelCase     )     )
     __A				 =       torch.mean(torch.abs(UpperCAmelCase     )     )

     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 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     __A				 =       self.scheduler_classes[0]
     __A				 =       self.get_scheduler_config(prediction_type='''v_prediction'''     )
     __A				 =       scheduler_class(**UpperCAmelCase     )

     scheduler.set_timesteps(self.num_inference_steps     )

     __A				 =       self.dummy_model()
     __A				 =       self.dummy_sample_deter * scheduler.init_noise_sigma
     __A				 =       sample.to(UpperCAmelCase     )

     for i, t in enumerate(scheduler.timesteps     ):
        __A				 =       scheduler.scale_model_input(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       model(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       scheduler.step(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
        __A				 =       output.prev_sample

     __A				 =       torch.sum(torch.abs(UpperCAmelCase     )     )
     __A				 =       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.693_4286_5017_0972e-07     ) < 1e-2
        assert abs(result_mean.item() - 0.0002     ) < 1e-3

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __A				 =       self.scheduler_classes[0]
     __A				 =       self.get_scheduler_config()
     __A				 =       scheduler_class(**UpperCAmelCase     )

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

     __A				 =       self.dummy_model()
     __A				 =       self.dummy_sample_deter.to(UpperCAmelCase     ) * scheduler.init_noise_sigma

     for t in scheduler.timesteps:
        __A				 =       scheduler.scale_model_input(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       model(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       scheduler.step(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
        __A				 =       output.prev_sample

     __A				 =       torch.sum(torch.abs(UpperCAmelCase     )     )
     __A				 =       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.1233     ) < 1e-2
        assert abs(result_mean.item() - 0.0002     ) < 1e-3
     elif str(UpperCAmelCase     ).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 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __A				 =       self.scheduler_classes[0]
     __A				 =       self.get_scheduler_config()
     __A				 =       scheduler_class(**UpperCAmelCase  ,    use_karras_sigmas=UpperCAmelCase     )

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

     __A				 =       self.dummy_model()
     __A				 =       self.dummy_sample_deter.to(UpperCAmelCase     ) * scheduler.init_noise_sigma
     __A				 =       sample.to(UpperCAmelCase     )

     for t in scheduler.timesteps:
        __A				 =       scheduler.scale_model_input(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       model(UpperCAmelCase  ,    UpperCAmelCase     )

        __A				 =       scheduler.step(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
        __A				 =       output.prev_sample

     __A				 =       torch.sum(torch.abs(UpperCAmelCase     )     )
     __A				 =       torch.mean(torch.abs(UpperCAmelCase     )     )

     assert abs(result_sum.item() - 0.00015     ) < 1e-2
     assert abs(result_mean.item() - 1.9869_5545_3503_4695e-07     ) < 1e-2








from __future__ import annotations

import sys
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						tuple[int, int]:


   '''simple docstring'''







   if abs(snake_case     ) == 1:
      return a, 0
   elif abs(snake_case     ) == 1:
      return 0, b

   __A ,						__A				 =       a, b
   __A ,						__A				 =       1, 0
   __A ,						__A				 =       0, 1

   while remainder != 0:
      __A				 =       old_remainder // remainder
      __A ,						__A				 =       remainder, old_remainder - quotient * remainder
      __A ,						__A				 =       coeff_a, old_coeff_a - quotient * coeff_a
      __A ,						__A				 =       coeff_b, old_coeff_b - quotient * coeff_b

   # sign correction for negative numbers
   if a < 0:
      __A				 =       -old_coeff_a
   if b < 0:
      __A				 =       -old_coeff_b

   return old_coeff_a, old_coeff_b
def 				__UpperCamelCase  (		)  ->						Tuple:


   '''simple docstring'''







   if len(sys.argv     ) < 3:
      print('''2 integer arguments required'''     )
      return 1
   __A				 =       int(sys.argv[1]     )
   __A				 =       int(sys.argv[2]     )
   print(extended_euclidean_algorithm(snake_case    ,	snake_case     )     )
   return 0


if __name__ == "__main__":
    raise SystemExit(main())

