fix: DDIM restoration when batch_size > 1

jolibrain · Aug 22, 2023 · ccb445b · ccb445b
1 parent bd16f1e
commit ccb445b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/modules/diffusion_generator.py b/models/modules/diffusion_generator.py
@@ -378,12 +378,12 @@ def ddim_p_mean_variance(
             y_0_hat.clamp_(-1.0, 1.0)
 
         gamma_t = self.extract(
-            getattr(self.denoise_fn.model, "gammas_" + phase), t, x_shape=(1, 1)
+            getattr(self.denoise_fn.model, "gammas_" + phase), t, x_shape=(1, 1, 1, 1)
         ).to(y_t.device)
         gamma_prevt = self.extract(
             getattr(self.denoise_fn.model, "gammas_prev_" + phase),
             prevt + 1,
-            x_shape=(1, 1),
+            x_shape=(1, 1, 1, 1),
         ).to(y_t.device)
 
         ## denoising formula for model_mean witih DDIM