From e4a4e4acaa67fab9b89fb286eac344489127f973 Mon Sep 17 00:00:00 2001 From: lukovnikov Date: Wed, 31 Aug 2022 15:28:35 +0200 Subject: [PATCH] Revert "Revert "fix ddim sampling"" This reverts commit cd8329cdd73a42b9252c9ef996ec554bc75d3396. --- .../denoising_diffusion_pytorch.py | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 884f81f..b3854c0 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -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