fix trainer max_steps bug

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
  )

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 = []

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)

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))