Open Flawless1202 opened 1 year ago
The code is here:
class LangevinCorrector(Corrector): def __init__(self, sde, score_fn, snr, n_steps): super().__init__(sde, score_fn, snr, n_steps) if not isinstance(sde, sde_lib.VPSDE) \ and not isinstance(sde, sde_lib.VESDE) \ and not isinstance(sde, sde_lib.subVPSDE): raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.") def update_fn(self, x, t): sde = self.sde score_fn = self.score_fn n_steps = self.n_steps target_snr = self.snr if isinstance(sde, sde_lib.VPSDE) or isinstance(sde, sde_lib.subVPSDE): timestep = (t * (sde.N - 1) / sde.T).long() alpha = sde.alphas.to(t.device)[timestep] else: alpha = torch.ones_like(t) for i in range(n_steps): grad = score_fn(x, t) noise = torch.randn_like(x) grad_norm = torch.norm(grad.reshape(grad.shape[0], -1), dim=-1).mean() noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean() step_size = (target_snr * noise_norm / grad_norm) ** 2 * 2 * alpha x_mean = x + step_size[:, None, None, None] * grad x = x_mean + torch.sqrt(step_size * 2)[:, None, None, None] * noise return x, x_mean
Could you please give me some explanation or references? Thx a lot! @yang-song @patrickvonplaten
snr is the variable r in Algorithm 4 and Algorithm 5 in the original paper
The code is here:
Could you please give me some explanation or references? Thx a lot! @yang-song @patrickvonplaten