CoCalc -- test_scheduler

GitHub Repository: shivamshrirao/diffusers
Path: blob/main/tests/schedulers/test_scheduler_heun.py
¹⁴⁴⁸ views
1
import torch
2

3
from diffusers import HeunDiscreteScheduler
4
from diffusers.utils import torch_device
5

6
from .test_schedulers import SchedulerCommonTest
7

8

9
class HeunDiscreteSchedulerTest(SchedulerCommonTest):
10
    scheduler_classes = (HeunDiscreteScheduler,)
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_prediction_type(self):
37
        for prediction_type in ["epsilon", "v_prediction"]:
38
            self.check_over_configs(prediction_type=prediction_type)
39

40
    def test_full_loop_no_noise(self):
41
        scheduler_class = self.scheduler_classes[0]
42
        scheduler_config = self.get_scheduler_config()
43
        scheduler = scheduler_class(**scheduler_config)
44

45
        scheduler.set_timesteps(self.num_inference_steps)
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)
57
            sample = output.prev_sample
58

59
        result_sum = torch.sum(torch.abs(sample))
60
        result_mean = torch.mean(torch.abs(sample))
61

62
        if torch_device in ["cpu", "mps"]:
63
            assert abs(result_sum.item() - 0.1233) < 1e-2
64
            assert abs(result_mean.item() - 0.0002) < 1e-3
65
        else:
66
            # CUDA
67
            assert abs(result_sum.item() - 0.1233) < 1e-2
68
            assert abs(result_mean.item() - 0.0002) < 1e-3
69

70
    def test_full_loop_with_v_prediction(self):
71
        scheduler_class = self.scheduler_classes[0]
72
        scheduler_config = self.get_scheduler_config(prediction_type="v_prediction")
73
        scheduler = scheduler_class(**scheduler_config)
74

75
        scheduler.set_timesteps(self.num_inference_steps)
76

77
        model = self.dummy_model()
78
        sample = self.dummy_sample_deter * scheduler.init_noise_sigma
79
        sample = sample.to(torch_device)
80

81
        for i, t in enumerate(scheduler.timesteps):
82
            sample = scheduler.scale_model_input(sample, t)
83

84
            model_output = model(sample, t)
85

86
            output = scheduler.step(model_output, t, sample)
87
            sample = output.prev_sample
88

89
        result_sum = torch.sum(torch.abs(sample))
90
        result_mean = torch.mean(torch.abs(sample))
91

92
        if torch_device in ["cpu", "mps"]:
93
            assert abs(result_sum.item() - 4.6934e-07) < 1e-2
94
            assert abs(result_mean.item() - 6.1112e-10) < 1e-3
95
        else:
96
            # CUDA
97
            assert abs(result_sum.item() - 4.693428650170972e-07) < 1e-2
98
            assert abs(result_mean.item() - 0.0002) < 1e-3
99

100
    def test_full_loop_device(self):
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

107
        model = self.dummy_model()
108
        sample = self.dummy_sample_deter.to(torch_device) * scheduler.init_noise_sigma
109

110
        for t in scheduler.timesteps:
111
            sample = scheduler.scale_model_input(sample, t)
112

113
            model_output = model(sample, t)
114

115
            output = scheduler.step(model_output, t, sample)
116
            sample = output.prev_sample
117

118
        result_sum = torch.sum(torch.abs(sample))
119
        result_mean = torch.mean(torch.abs(sample))
120

121
        if str(torch_device).startswith("cpu"):
122
            # The following sum varies between 148 and 156 on mps. Why?
123
            assert abs(result_sum.item() - 0.1233) < 1e-2
124
            assert abs(result_mean.item() - 0.0002) < 1e-3
125
        elif str(torch_device).startswith("mps"):
126
            # Larger tolerance on mps
127
            assert abs(result_mean.item() - 0.0002) < 1e-2
128
        else:
129
            # CUDA
130
            assert abs(result_sum.item() - 0.1233) < 1e-2
131
            assert abs(result_mean.item() - 0.0002) < 1e-3
132

133
Product

Resources

Company
