| 58 | |
| 59 | |
| 60 | class LinearSchedule: |
| 61 | def alpha_t(self, t): |
| 62 | return t |
| 63 | |
| 64 | def alpha_dt_t(self, t): |
| 65 | return 1 |
| 66 | |
| 67 | def sigma_t(self, t): |
| 68 | return 1 - t |
| 69 | |
| 70 | def sigma_dt_t(self, t): |
| 71 | return -1 |
| 72 | |
| 73 | """ Legacy functions to work with SiT Sampler """ |
| 74 | |
| 75 | def compute_alpha_t(self, t): |
| 76 | return self.alpha_t(t), self.alpha_dt_t(t) |
| 77 | |
| 78 | def compute_sigma_t(self, t): |
| 79 | """Compute the noise coefficient along the path""" |
| 80 | return self.sigma_t(t), self.sigma_dt_t(t) |
| 81 | |
| 82 | def compute_d_alpha_alpha_ratio_t(self, t): |
| 83 | """Compute the ratio between d_alpha and alpha""" |
| 84 | return 1 / t |
| 85 | |
| 86 | def compute_drift(self, x, t): |
| 87 | """We always output sde according to score parametrization; """ |
| 88 | t = pad_v_like_x(t, x) |
| 89 | alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t) |
| 90 | sigma_t, d_sigma_t = self.compute_sigma_t(t) |
| 91 | drift = alpha_ratio * x |
| 92 | diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t |
| 93 | |
| 94 | return -drift, diffusion |
| 95 | |
| 96 | def compute_diffusion(self, x, t, form="constant", norm=1.0): |
| 97 | """Compute the diffusion term of the SDE |
| 98 | Args: |
| 99 | x: [batch_dim, ...], data point |
| 100 | t: [batch_dim,], time vector |
| 101 | form: str, form of the diffusion term |
| 102 | norm: float, norm of the diffusion term |
| 103 | """ |
| 104 | t = pad_v_like_x(t, x) |
| 105 | choices = { |
| 106 | "constant": norm, |
| 107 | "SBDM": norm * self.compute_drift(x, t)[1], |
| 108 | "sigma": norm * self.compute_sigma_t(t)[0], |
| 109 | "linear": norm * (1 - t), |
| 110 | "decreasing": 0.25 * (norm * torch.cos(np.pi * t) + 1) ** 2, |
| 111 | "increasing-decreasing": norm * torch.sin(np.pi * t) ** 2, |
| 112 | } |
| 113 | |
| 114 | try: diffusion = choices[form] |
| 115 | except KeyError: raise NotImplementedError(f"Diffusion form {form} not implemented") |
| 116 | |
| 117 | return diffusion |