fix trainer max_steps bug

threestudio-project · Jan 3, 2024 · 585232b · 585232b
1 parent 3546097
commit 585232b
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 6 deletions.
diff --git a/launch.py b/launch.py
@@ -167,7 +167,6 @@ def main(args, extras) -> None:
 
     dm = threestudio.find(cfg.data_type)(cfg.data)
     # guidance needs to know max_steps for annealing
-    cfg.system.guidance.trainer_max_steps = cfg.trainer.max_steps
     system: BaseSystem = threestudio.find(cfg.system_type)(
         cfg.system, resumed=cfg.resume is not None
     )

diff --git a/threestudio/models/guidance/stable_diffusion_guidance.py b/threestudio/models/guidance/stable_diffusion_guidance.py
@@ -446,7 +446,6 @@ def __call__(
         # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
         loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
 
-
         guidance_out = {
             "loss_sds": loss_sds,
             "grad_norm": grad.norm(),
@@ -463,7 +462,7 @@ def __call__(
                 0.5 * F.mse_loss(rgb_BCHW_512, target_img, reduction="sum") / batch_size
             )
             guidance_out["loss_sds_img"] = loss_sds_img
-            
+
         if guidance_eval:
             guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
             texts = []

diff --git a/threestudio/models/guidance/stable_diffusion_unified_guidance.py b/threestudio/models/guidance/stable_diffusion_unified_guidance.py
@@ -671,14 +671,19 @@ def forward(
         # image-space loss proposed in HiFA: https://hifa-team.github.io/HiFA-site
         if self.cfg.use_img_loss:
             if self.cfg.guidance_type == "vsd":
-                latents_denoised_est = (latents_noisy - self.sigmas[t] * eps_phi) / self.alphas[t]
+                latents_denoised_est = (
+                    latents_noisy - self.sigmas[t] * eps_phi
+                ) / self.alphas[t]
                 image_denoised_est = self.vae_decode(
                     self.pipe.vae, latents_denoised_est
                 )
             else:
                 image_denoised_est = rgb_BCHW_512
             grad_img = (
-                w * (image_denoised_est - image_denoised_pretrain) * self.alphas[t] / self.sigmas[t]
+                w
+                * (image_denoised_est - image_denoised_pretrain)
+                * self.alphas[t]
+                / self.sigmas[t]
             )
             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/systems/dreamfusion.py b/threestudio/systems/dreamfusion.py
@@ -75,7 +75,7 @@ def training_step(self, batch, batch_idx):
             loss_z_variance = out["z_variance"][out["opacity"] > 0.5].mean()
             self.log("train/loss_z_variance", loss_z_variance)
             loss += loss_z_variance * self.C(self.cfg.loss.lambda_z_variance)
-            
+
         for name, value in self.cfg.loss.items():
             self.log(f"train_params/{name}", self.C(value))