Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
shivamshrirao
GitHub Repository: shivamshrirao/diffusers
 Path: blob/main/tests/schedulers/test_scheduler_kdpm2_ancestral.py
1448 views
1
import torch

2


3
from diffusers import KDPM2AncestralDiscreteScheduler

4
from diffusers.utils import torch_device

5


6
from .test_schedulers import SchedulerCommonTest

7


8


9
class KDPM2AncestralDiscreteSchedulerTest(SchedulerCommonTest):

10
    scheduler_classes = (KDPM2AncestralDiscreteScheduler,)

11
    num_inference_steps = 10

12


13
    def get_scheduler_config(self, **kwargs):

14
        config = {

15
            "num_train_timesteps": 1100,

16
            "beta_start": 0.0001,

17
            "beta_end": 0.02,

18
            "beta_schedule": "linear",

19
        }

20


21
        config.update(**kwargs)

22
        return config

23


24
    def test_timesteps(self):

25
        for timesteps in [10, 50, 100, 1000]:

26
            self.check_over_configs(num_train_timesteps=timesteps)

27


28
    def test_betas(self):

29
        for beta_start, beta_end in zip([0.00001, 0.0001, 0.001], [0.0002, 0.002, 0.02]):

30
            self.check_over_configs(beta_start=beta_start, beta_end=beta_end)

31


32
    def test_schedules(self):

33
        for schedule in ["linear", "scaled_linear"]:

34
            self.check_over_configs(beta_schedule=schedule)

35


36
    def test_full_loop_no_noise(self):

37
        if torch_device == "mps":

38
            return

39
        scheduler_class = self.scheduler_classes[0]

40
        scheduler_config = self.get_scheduler_config()

41
        scheduler = scheduler_class(**scheduler_config)

42


43
        scheduler.set_timesteps(self.num_inference_steps)

44


45
        generator = torch.manual_seed(0)

46


47
        model = self.dummy_model()

48
        sample = self.dummy_sample_deter * scheduler.init_noise_sigma

49
        sample = sample.to(torch_device)

50


51
        for i, t in enumerate(scheduler.timesteps):

52
            sample = scheduler.scale_model_input(sample, t)

53


54
            model_output = model(sample, t)

55


56
            output = scheduler.step(model_output, t, sample, generator=generator)

57
            sample = output.prev_sample

58


59
        result_sum = torch.sum(torch.abs(sample))

60
        result_mean = torch.mean(torch.abs(sample))

61


62
        assert abs(result_sum.item() - 13849.3877) < 1e-2

63
        assert abs(result_mean.item() - 18.0331) < 5e-3

64


65
    def test_prediction_type(self):

66
        for prediction_type in ["epsilon", "v_prediction"]:

67
            self.check_over_configs(prediction_type=prediction_type)

68


69
    def test_full_loop_with_v_prediction(self):

70
        if torch_device == "mps":

71
            return

72
        scheduler_class = self.scheduler_classes[0]

73
        scheduler_config = self.get_scheduler_config(prediction_type="v_prediction")

74
        scheduler = scheduler_class(**scheduler_config)

75


76
        scheduler.set_timesteps(self.num_inference_steps)

77


78
        model = self.dummy_model()

79
        sample = self.dummy_sample_deter * scheduler.init_noise_sigma

80
        sample = sample.to(torch_device)

81


82
        generator = torch.manual_seed(0)

83


84
        for i, t in enumerate(scheduler.timesteps):

85
            sample = scheduler.scale_model_input(sample, t)

86


87
            model_output = model(sample, t)

88


89
            output = scheduler.step(model_output, t, sample, generator=generator)

90
            sample = output.prev_sample

91


92
        result_sum = torch.sum(torch.abs(sample))

93
        result_mean = torch.mean(torch.abs(sample))

94


95
        assert abs(result_sum.item() - 328.9970) < 1e-2

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

97


98
    def test_full_loop_device(self):

99
        if torch_device == "mps":

100
            return

101
        scheduler_class = self.scheduler_classes[0]

102
        scheduler_config = self.get_scheduler_config()

103
        scheduler = scheduler_class(**scheduler_config)

104


105
        scheduler.set_timesteps(self.num_inference_steps, device=torch_device)

106
        generator = torch.manual_seed(0)

107


108
        model = self.dummy_model()

109
        sample = self.dummy_sample_deter.to(torch_device) * scheduler.init_noise_sigma

110


111
        for t in scheduler.timesteps:

112
            sample = scheduler.scale_model_input(sample, t)

113


114
            model_output = model(sample, t)

115


116
            output = scheduler.step(model_output, t, sample, generator=generator)

117
            sample = output.prev_sample

118


119
        result_sum = torch.sum(torch.abs(sample))

120
        result_mean = torch.mean(torch.abs(sample))

121


122
        assert abs(result_sum.item() - 13849.3818) < 1e-1

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

124


125