diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index fb256d8..e1b79d9 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -313,6 +313,22 @@ class GaussianDiffusion(nn.Module): def sample(self, image_size, batch_size = 16): return self.p_sample_loop((16, 3, image_size, image_size)) + @torch.no_grad() + def interpolate(self, x1, x2, t = None, lam = 0.5): + b, *_, device = *x1.shape, x1.device + t = default(t, self.num_timesteps - 1) + + assert x1.shape == x2.shape + + t_batched = torch.stack([torch.tensor(t, device=device)] * b) + xt1, xt2 = map(lambda x: self.q_sample(x, t=t_batched), (x1, x2)) + + img = (1 - lam) * xt1 + lam * xt2 + for i in tqdm(reversed(range(0, t)), desc='interpolation sample time step', total=t): + img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long)) + + return img + def q_sample(self, x_start, t, noise=None): noise = default(noise, lambda: torch.randn_like(x_start)) diff --git a/setup.py b/setup.py index d16aa88..27fdeb0 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages setup( name = 'denoising-diffusion-pytorch', packages = find_packages(), - version = '0.1.2', + version = '0.1.3', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',