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

Revert "Revert "fix ddim sampling""

Browse Source 
This reverts commit cd8329cdd7.
This commit is contained in:
lukovnikov
2022-08-31 15:28:35 +02:00
parent cd8329cdd7
commit e4a4e4acaa
1 changed files with 15 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+15 -14
View File
@@ -568,35 +568,36 @@ class GaussianDiffusion(nn.Module):
def ddim_sample(self, shape, clip_denoised = True):
batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
times = torch.linspace(0., total_timesteps, steps = sampling_timesteps + 2)[:-1]
times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1) # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
times = list(reversed(times.int().tolist()))
time_pairs = list(filter(lambda a: a[0] > a[1], zip(times[:-1], times[1:])))
time_pairs = list(zip(times[:-1], times[1:])) # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]
img = torch.randn(shape, device = device)
x_start = None
for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
alpha = self.alphas_cumprod[time]
alpha_next = self.alphas_cumprod[time_next]
time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
self_cond = x_start if self.self_condition else None
pred_noise, x_start, *_ = self.model_predictions(img, time_cond, self_cond)
if clip_denoised:
x_start.clamp_(-1., 1.)
sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
c = ((1 - alpha_next) - sigma ** 2).sqrt()
if time_next > -1:
alpha = self.alphas_cumprod[time]
alpha_next = self.alphas_cumprod[time_next]
noise = torch.randn_like(img) if time_next > 0 else 0.
sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
c = (1 - alpha_next - sigma ** 2).sqrt()
img = x_start * alpha_next.sqrt() + \
c * pred_noise + \
sigma * noise
noise = torch.randn_like(img)
img = x_start * alpha_next.sqrt() + \
c * pred_noise + \
sigma * noise
else:
img = x_start
img = unnormalize_to_zero_to_one(img)
return img