





'''simple docstring'''






import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest
class 					__snake_case  (					__SCREAMING_SNAKE_CASE							):






   '''simple docstring'''





   lowerCamelCase__    							=      (DPMSolverSinglestepScheduler,)
   lowerCamelCase__    							=      (('''num_inference_steps''', 25),)

   def 		__UpperCamelCase						(							self	,		**__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Union[str, Any]    =  {
        """num_train_timesteps""": 1_0_0_0,
        """beta_start""": 0.0001,
        """beta_end""": 0.02,
        """beta_schedule""": """linear""",
        """solver_order""": 2,
        """prediction_type""": """epsilon""",
        """thresholding""": False,
        """sample_max_value""": 1.0,
        """algorithm_type""": """dpmsolver++""",
        """solver_type""": """midpoint""",
        """lambda_min_clipped""": -float("""inf"""       ),
        """variance_type""": None,
    }

    config.update(**__SCREAMING_SNAKE_CASE       )
    return config

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE=0	,		**__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					str    =  dict(self.forward_default_kwargs       )
    snake_case__     :					Optional[int]    =  kwargs.pop("""num_inference_steps"""	,		__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Optional[Any]    =  self.dummy_sample
    snake_case__     :					Dict    =  0.1 * sample
    snake_case__     :					Any    =  [residual + 0.2, residual + 0.15, residual + 0.10]

    for scheduler_class in self.scheduler_classes:
     snake_case__     :					Dict    =  self.get_scheduler_config(**__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Tuple    =  scheduler_class(**__SCREAMING_SNAKE_CASE       )
     scheduler.set_timesteps(__SCREAMING_SNAKE_CASE       )
     # copy over dummy past residuals
     snake_case__     :					Union[str, Any]    =  dummy_past_residuals[: scheduler.config.solver_order]

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

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

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

   def 		__UpperCamelCase						(							self       ):
    pass

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE=0	,		**__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					List[str]    =  dict(self.forward_default_kwargs       )
    snake_case__     :					Tuple    =  kwargs.pop("""num_inference_steps"""	,		__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Union[str, Any]    =  self.dummy_sample
    snake_case__     :					Tuple    =  0.1 * sample
    snake_case__     :					List[str]    =  [residual + 0.2, residual + 0.15, residual + 0.10]

    for scheduler_class in self.scheduler_classes:
     snake_case__     :					List[str]    =  self.get_scheduler_config()
     snake_case__     :					Union[str, Any]    =  scheduler_class(**__SCREAMING_SNAKE_CASE       )
     scheduler.set_timesteps(__SCREAMING_SNAKE_CASE       )

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

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

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

     snake_case__     :					str    =  scheduler.step(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		**__SCREAMING_SNAKE_CASE       ).prev_sample
     snake_case__     :					Dict    =  new_scheduler.step(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		**__SCREAMING_SNAKE_CASE       ).prev_sample

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

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE=None	,		**__SCREAMING_SNAKE_CASE       ):
    if scheduler is None:
     snake_case__     :					Any    =  self.scheduler_classes[0]
     snake_case__     :					List[str]    =  self.get_scheduler_config(**__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Union[str, Any]    =  scheduler_class(**__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Any    =  self.scheduler_classes[0]
    snake_case__     :					Optional[Any]    =  self.get_scheduler_config(**__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[Any]    =  scheduler_class(**__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Tuple    =  1_0
    snake_case__     :					Union[str, Any]    =  self.dummy_model()
    snake_case__     :					str    =  self.dummy_sample_deter
    scheduler.set_timesteps(__SCREAMING_SNAKE_CASE       )

    for i, t in enumerate(scheduler.timesteps       ):
     snake_case__     :					int    =  model(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     snake_case__     :					str    =  scheduler.step(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ).prev_sample

    return sample

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Union[str, Any]    =  DPMSolverSinglestepScheduler(**self.get_scheduler_config()       )
    snake_case__     :					str    =  5_0
    snake_case__     :					str    =  self.dummy_model()
    snake_case__     :					Tuple    =  self.dummy_sample_deter
    scheduler.set_timesteps(__SCREAMING_SNAKE_CASE       )

    # make sure that the first t is uneven
    for i, t in enumerate(scheduler.timesteps[3:]       ):
     snake_case__     :					Any    =  model(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Optional[Any]    =  scheduler.step(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ).prev_sample

    snake_case__     :					Union[str, Any]    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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

   def 		__UpperCamelCase						(							self       ):
    for timesteps in [2_5, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
     self.check_over_configs(num_train_timesteps=__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    # make sure that iterating over schedulers with same config names gives same results
    # for defaults
    snake_case__     :					Optional[Any]    =  DPMSolverSinglestepScheduler(**self.get_scheduler_config()       )
    snake_case__     :					List[Any]    =  self.full_loop(scheduler=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Union[str, Any]    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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

    snake_case__     :					Optional[int]    =  DEISMultistepScheduler.from_config(scheduler.config       )
    snake_case__     :					List[str]    =  DPMSolverMultistepScheduler.from_config(scheduler.config       )
    snake_case__     :					List[str]    =  UniPCMultistepScheduler.from_config(scheduler.config       )
    snake_case__     :					Union[str, Any]    =  DPMSolverSinglestepScheduler.from_config(scheduler.config       )

    snake_case__     :					List[str]    =  self.full_loop(scheduler=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Dict    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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

   def 		__UpperCamelCase						(							self       ):
    self.check_over_configs(thresholding=__SCREAMING_SNAKE_CASE       )
    for order in [1, 2, 3]:
     for solver_type in ["midpoint", "heun"]:
      for threshold in [0.5, 1.0, 2.0]:
       for prediction_type in ["epsilon", "sample"]:
        self.check_over_configs(
            thresholding=__SCREAMING_SNAKE_CASE	,		prediction_type=__SCREAMING_SNAKE_CASE	,		sample_max_value=__SCREAMING_SNAKE_CASE	,		algorithm_type="""dpmsolver++"""	,		solver_order=__SCREAMING_SNAKE_CASE	,		solver_type=__SCREAMING_SNAKE_CASE	,		)

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

   def 		__UpperCamelCase						(							self       ):
    for algorithm_type in ["dpmsolver", "dpmsolver++"]:
     for solver_type in ["midpoint", "heun"]:
      for order in [1, 2, 3]:
       for prediction_type in ["epsilon", "sample"]:
        self.check_over_configs(
            solver_order=__SCREAMING_SNAKE_CASE	,		solver_type=__SCREAMING_SNAKE_CASE	,		prediction_type=__SCREAMING_SNAKE_CASE	,		algorithm_type=__SCREAMING_SNAKE_CASE	,		)
        snake_case__     :					Tuple    =  self.full_loop(
            solver_order=__SCREAMING_SNAKE_CASE	,		solver_type=__SCREAMING_SNAKE_CASE	,		prediction_type=__SCREAMING_SNAKE_CASE	,		algorithm_type=__SCREAMING_SNAKE_CASE	,		)
        assert not torch.isnan(__SCREAMING_SNAKE_CASE       ).any(), "Samples have nan numbers"

   def 		__UpperCamelCase						(							self       ):
    self.check_over_configs(lower_order_final=__SCREAMING_SNAKE_CASE       )
    self.check_over_configs(lower_order_final=__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    self.check_over_configs(lambda_min_clipped=-float("""inf"""       )       )
    self.check_over_configs(lambda_min_clipped=-5.1       )

   def 		__UpperCamelCase						(							self       ):
    self.check_over_configs(variance_type=__SCREAMING_SNAKE_CASE       )
    self.check_over_configs(variance_type="""learned_range"""       )

   def 		__UpperCamelCase						(							self       ):
    for num_inference_steps in [1, 2, 3, 5, 1_0, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
     self.check_over_forward(num_inference_steps=__SCREAMING_SNAKE_CASE	,		time_step=0       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[int]    =  self.full_loop()
    snake_case__     :					str    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[Any]    =  self.full_loop(use_karras_sigmas=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Any    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[str]    =  self.full_loop(prediction_type="""v_prediction"""       )
    snake_case__     :					str    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					int    =  self.full_loop(prediction_type="""v_prediction"""	,		use_karras_sigmas=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Tuple    =  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE       )       )

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


   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[Any]    =  self.scheduler_classes[0]
    snake_case__     :					str    =  self.get_scheduler_config(thresholding=__SCREAMING_SNAKE_CASE	,		dynamic_thresholding_ratio=0       )
    snake_case__     :					List[str]    =  scheduler_class(**__SCREAMING_SNAKE_CASE       )

    snake_case__     :					List[Any]    =  1_0
    snake_case__     :					Optional[int]    =  self.dummy_model()
    snake_case__     :					Any    =  self.dummy_sample_deter.half()
    scheduler.set_timesteps(__SCREAMING_SNAKE_CASE       )

    for i, t in enumerate(scheduler.timesteps       ):
     snake_case__     :					Tuple    =  model(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Optional[int]    =  scheduler.step(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ).prev_sample

    assert sample.dtype == torch.floataa






"""simple docstring"""



import math
import unittest


def 			_UpperCamelCase						(    _A   )     ->	bool:





    """simple docstring"""

    assert isinstance(_A       ,   _A   ) and (
        number >= 0
    ), "'number' must been an int and positive"

    if 1 < number < 4:
        # 2 and 3 are primes
        return True
    elif number < 2 or number % 2 == 0 or number % 3 == 0:
        # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
        return False

    # All primes number are in format of 6k +/- 1
    for i in range(5       ,   int(math.sqrt(_A   ) + 1   )       ,   6   ):
        if number % i == 0 or number % (i + 2) == 0:
            return False
    return True




class 		a_    (			unittest.TestCase						):







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






        '''simple docstring'''






        self.assertTrue(is_prime(2  )  )
        self.assertTrue(is_prime(3  )  )
        self.assertTrue(is_prime(5  )  )
        self.assertTrue(is_prime(7  )  )
        self.assertTrue(is_prime(11  )  )
        self.assertTrue(is_prime(13  )  )
        self.assertTrue(is_prime(17  )  )
        self.assertTrue(is_prime(19  )  )
        self.assertTrue(is_prime(23  )  )
        self.assertTrue(is_prime(29  )  )







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






        '''simple docstring'''






        with self.assertRaises(__UpperCamelCase  ):
            is_prime(-19  )
        self.assertFalse(
            is_prime(0  )  ,       """Zero doesn't have any positive factors, primes must have exactly two."""  ,       )
        self.assertFalse(
            is_prime(1  )  ,       """One only has 1 positive factor, primes must have exactly two."""  ,       )
        self.assertFalse(is_prime(2 * 2  )  )
        self.assertFalse(is_prime(2 * 3  )  )
        self.assertFalse(is_prime(3 * 3  )  )
        self.assertFalse(is_prime(3 * 5  )  )
        self.assertFalse(is_prime(3 * 5 * 7  )  )


if __name__ == "__main__":
 unittest.main()