quick fix for interpolate in cfg example

Author: lucidrains

denoising_diffusion_pytorch/classifier_free_guidance.py

@@ -699,7 +699,7 @@ def sample(self, classes, cond_scale = 6., rescaled_phi = 0.7):
         return sample_fn(classes, (batch_size, channels, image_size, image_size), cond_scale, rescaled_phi)
 
     @torch.no_grad()
-    def interpolate(self, x1, x2, t = None, lam = 0.5):
+    def interpolate(self, x1, x2, classes, t = None, lam = 0.5):
         b, *_, device = *x1.shape, x1.device
         t = default(t, self.num_timesteps - 1)
 
@@ -709,8 +709,9 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
         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))
+            img, _ = self.p_sample(img, i, classes)
 
         return img
 
@@ -795,3 +796,12 @@ def forward(self, img, *args, **kwargs):
     )
 
     sampled_images.shape # (8, 3, 128, 128)
+
+    # interpolation
+
+    interpolate_out = diffusion.interpolate(
+        training_images[:1],
+        training_images[:1],
+        image_classes[:1]
+    )
+

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.14'
+__version__ = '1.8.15'