wassname/denoising-diffusion-pytorch
1
0
Fork 0
mirror of https://github.com/wassname/denoising-diffusion-pytorch.git synced 2026-09-11 12:11:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

get working version of gaussian diffusion with continuous time (only beta linear schedule for now, but will eventually contain alpha cosine schedule as well as parameterized, learned monotonic MLP)

Browse Source
This commit is contained in:
Phil Wang
2022-06-07 15:59:29 -07:00
parent 6b93fa48f6
commit 3bbb6ebf16
5 changed files with 203 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+2 -7
View File
@@ -324,11 +324,6 @@ def extract(a, t, x_shape):
out = a.gather(-1, t)
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
def noise_like(shape, device, repeat=False):
repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
noise = lambda: torch.randn(shape, device=device)
return repeat_noise() if repeat else noise()
def linear_beta_schedule(timesteps):
scale = 1000 / timesteps
beta_start = scale * 0.0001
@@ -444,10 +439,10 @@ class GaussianDiffusion(nn.Module):
return model_mean, posterior_variance, posterior_log_variance
@torch.no_grad()
def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
def p_sample(self, x, t, clip_denoised=True):
b, *_, device = *x.shape, x.device
model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
noise = noise_like(x.shape, device, repeat_noise)
noise = torch.randn_like(x)
# no noise when t == 0
nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise