fix batch bug

threestudio-project · Jan 5, 2024 · ef9c9b9 · ef9c9b9
1 parent aa5eb83
commit ef9c9b9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/threestudio/models/guidance/stable_diffusion_guidance.py b/threestudio/models/guidance/stable_diffusion_guidance.py
@@ -264,8 +264,8 @@ def compute_grad_sds(
                 f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
             )
 
-        alpha = self.alphas[t] ** 0.5
-        sigma = (1 - self.alphas[t]) ** 0.5
+        alpha = (self.alphas[t] ** 0.5).view(-1, 1, 1, 1)
+        sigma = ((1 - self.alphas[t]) ** 0.5).view(-1, 1, 1, 1)
         latents_denoised = (latents_noisy - sigma * noise_pred) / alpha
         image_denoised = self.decode_latents(latents_denoised)
 

diff --git a/threestudio/models/guidance/stable_diffusion_vsd_guidance.py b/threestudio/models/guidance/stable_diffusion_vsd_guidance.py
@@ -541,8 +541,8 @@ def compute_grad_vsd(
 
         grad = w * (noise_pred_pretrain - noise_pred_est)
 
-        alpha = self.alphas[t] ** 0.5
-        sigma = (1 - self.alphas[t]) ** 0.5
+        alpha = (self.alphas[t] ** 0.5).view(-1, 1, 1, 1)
+        sigma = ((1 - self.alphas[t]) ** 0.5).view(-1, 1, 1, 1)
         # image-space SDS proposed in HiFA: https://hifa-team.github.io/HiFA-site/
         if self.cfg.use_img_loss:
             latents_denoised_pretrain = (