fix hifa batch bug (#400)

threestudio-project · Jan 5, 2024 · 3dc4f1f · 3dc4f1f
1 parent d85e1f8
commit 3dc4f1f
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 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_unified_guidance.py b/threestudio/models/guidance/stable_diffusion_unified_guidance.py
@@ -673,7 +673,7 @@ def forward(
             if self.cfg.guidance_type == "vsd":
                 latents_denoised_est = (
                     latents_noisy - self.sigmas[t] * eps_phi
-                ) / self.alphas[t]
+                ) / self.alphas[t].view(-1, 1, 1, 1)
                 image_denoised_est = self.vae_decode(
                     self.pipe.vae, latents_denoised_est
                 )
@@ -682,8 +682,8 @@ def forward(
             grad_img = (
                 w
                 * (image_denoised_est - image_denoised_pretrain)
-                * self.alphas[t]
-                / self.sigmas[t]
+                * self.alphas[t].view(-1, 1, 1, 1)
+                / self.sigmas[t].view(-1, 1, 1, 1)
             )
             if self.grad_clip_val is not None:
                 grad_img = grad_img.clamp(-self.grad_clip_val, self.grad_clip_val)

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 = (