scheduling_sde_ve.py 6.46 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
# Copyright 2022 Google Brain and The HuggingFace Team. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
14
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Patrick von Platen's avatar
Patrick von Platen committed
15
16
17
# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

# TODO(Patrick, Anton, Suraj) - make scheduler framework indepedent and clean-up a bit
Nathan Lambert's avatar
Nathan Lambert committed
18
import pdb
19
20
21
22
23
24
25
26

import numpy as np
import torch

from ..configuration_utils import ConfigMixin
from .scheduling_utils import SchedulerMixin


Patrick von Platen's avatar
Patrick von Platen committed
27
class ScoreSdeVeScheduler(SchedulerMixin, ConfigMixin):
Nathan Lambert's avatar
Nathan Lambert committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    """
    The variance exploding stochastic differential equation (SDE) scheduler.

    :param snr: coefficient weighting the step from the score sample (from the network) to the random noise. :param
    sigma_min: initial noise scale for sigma sequence in sampling procedure. The minimum sigma should mirror the
            distribution of the data.
    :param sigma_max: :param sampling_eps: the end value of sampling, where timesteps decrease progessively from 1 to
    epsilon. :param correct_steps: number of correction steps performed on a produced sample. :param tensor_format:
    "np" or "pt" for the expected format of samples passed to the Scheduler.
    """

    def __init__(
        self,
        num_train_timesteps=2000,
        snr=0.15,
        sigma_min=0.01,
        sigma_max=1348,
        sampling_eps=1e-5,
        correct_steps=1,
        tensor_format="pt",
    ):
49
50
        super().__init__()
        self.register_to_config(
Nathan Lambert's avatar
Nathan Lambert committed
51
            num_train_timesteps=num_train_timesteps,
52
53
54
55
            snr=snr,
            sigma_min=sigma_min,
            sigma_max=sigma_max,
            sampling_eps=sampling_eps,
Nathan Lambert's avatar
Nathan Lambert committed
56
            correct_steps=correct_steps,
57
58
        )

Patrick von Platen's avatar
Patrick von Platen committed
59
60
61
62
        self.sigmas = None
        self.discrete_sigmas = None
        self.timesteps = None

Nathan Lambert's avatar
Nathan Lambert committed
63
64
65
        # TODO - update step to be torch-independant
        self.set_format(tensor_format=tensor_format)

Patrick von Platen's avatar
Patrick von Platen committed
66
    def set_timesteps(self, num_inference_steps):
Nathan Lambert's avatar
Nathan Lambert committed
67
68
69
70
71
72
73
        tensor_format = getattr(self, "tensor_format", "pt")
        if tensor_format == "np":
            self.timesteps = np.linspace(1, self.config.sampling_eps, num_inference_steps)
        elif tensor_format == "pt":
            self.timesteps = torch.linspace(1, self.config.sampling_eps, num_inference_steps)
        else:
            raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
Patrick von Platen's avatar
Patrick von Platen committed
74
75
76
77
78

    def set_sigmas(self, num_inference_steps):
        if self.timesteps is None:
            self.set_timesteps(num_inference_steps)

Nathan Lambert's avatar
Nathan Lambert committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
        tensor_format = getattr(self, "tensor_format", "pt")
        if tensor_format == "np":
            self.discrete_sigmas = np.exp(
                np.linspace(np.log(self.config.sigma_min), np.log(self.config.sigma_max), num_inference_steps)
            )
            self.sigmas = np.array(
                [self.config.sigma_min * (self.config.sigma_max / self.sigma_min) ** t for t in self.timesteps]
            )
        elif tensor_format == "pt":
            self.discrete_sigmas = torch.exp(
                torch.linspace(np.log(self.config.sigma_min), np.log(self.config.sigma_max), num_inference_steps)
            )
            self.sigmas = torch.tensor(
                [self.config.sigma_min * (self.config.sigma_max / self.sigma_min) ** t for t in self.timesteps]
            )
        else:
            raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")

    def get_adjacent_sigma(self, timesteps, t):
        tensor_format = getattr(self, "tensor_format", "pt")
        if tensor_format == "np":
            return np.where(timesteps == 0, np.zeros_like(t), self.discrete_sigmas[timesteps - 1])
        elif tensor_format == "pt":
            return torch.where(
                timesteps == 0, torch.zeros_like(t), self.discrete_sigmas[timesteps - 1].to(timesteps.device)
            )

        raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")

    def step_pred(self, score, x, t):
        """
        Predict the sample at the previous timestep by reversing the SDE.
        """
Patrick von Platen's avatar
Patrick von Platen committed
112
        # TODO(Patrick) better comments + non-PyTorch
Nathan Lambert's avatar
Nathan Lambert committed
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
        t = self.repeat_scalar(t, x.shape[0])
        timesteps = self.long((t * (len(self.timesteps) - 1)))

        sigma = self.discrete_sigmas[timesteps]
        adjacent_sigma = self.get_adjacent_sigma(timesteps, t)
        drift = self.zeros_like(x)
        diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5

        # equation 6 in the paper: the score modeled by the network is grad_x log pt(x)
        # also equation 47 shows the analog from SDE models to ancestral sampling methods
        drift = drift - diffusion[:, None, None, None] ** 2 * score

        #  equation 6: sample noise for the diffusion term of
        noise = self.randn_like(x)
        x_mean = x - drift  # subtract because `dt` is a small negative timestep
        # TODO is the variable diffusion the correct scaling term for the noise?
        x = x_mean + diffusion[:, None, None, None] * noise  # add impact of diffusion field g
        return x, x_mean
131

Nathan Lambert's avatar
Nathan Lambert committed
132
133
134
135
136
137
    def step_correct(self, score, x):
        """
        Correct the predicted sample based on the output score of the network. This is often run repeatedly after
        making the prediction for the previous timestep.
        """
        # TODO(Patrick) non-PyTorch
138

Nathan Lambert's avatar
Nathan Lambert committed
139
140
141
        # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
        # sample noise for correction
        noise = self.randn_like(x)
142

Nathan Lambert's avatar
Nathan Lambert committed
143
144
145
        # compute step size from the score, the noise, and the snr
        grad_norm = self.norm(score)
        noise_norm = self.norm(noise)
Patrick von Platen's avatar
Patrick von Platen committed
146
        step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
Nathan Lambert's avatar
Nathan Lambert committed
147
        step_size = self.repeat_scalar(step_size, x.shape[0])  # * self.ones(x.shape[0], device=x.device)
148

Nathan Lambert's avatar
Nathan Lambert committed
149
150
151
        # compute corrected sample: score term and noise term
        x_mean = x + step_size[:, None, None, None] * score
        x = x_mean + ((step_size * 2) ** 0.5)[:, None, None, None] * noise
152
153

        return x
Nathan Lambert's avatar
Nathan Lambert committed
154
155
156

    def __len__(self):
        return self.config.num_train_timesteps