From 4bfe944792d73c0d2c3026ad02135bd0c9f4e716 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 13 May 2026 09:06:59 -0600 Subject: [PATCH] Scale dfe 7 with velocity_equiv_weight --- .../models/diffusion_feature_extraction.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py index 32a1d0f..a035b26 100644 --- a/toolkit/models/diffusion_feature_extraction.py +++ b/toolkit/models/diffusion_feature_extraction.py @@ -949,6 +949,8 @@ class DiffusionFeatureExtractor7(nn.Module): device = self.model.device dtype = self.model.dtype + + velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2) with torch.no_grad(): target = self.get_pred(target_0_1) @@ -957,20 +959,24 @@ class DiffusionFeatureExtractor7(nn.Module): pred = self.get_pred(pred_images) head_loss = torch.nn.functional.mse_loss( - pred.head.float(), target.head.float() - ) + pred.head.float(), target.head.float(), reduction='none' + ) * velocity_equiv_weight + head_loss = head_loss.mean() depth_loss = torch.nn.functional.l1_loss( - pred.depth.float(), target.depth.float() - ) + pred.depth.float(), target.depth.float(), reduction='none' + ) * velocity_equiv_weight + depth_loss = depth_loss.mean() normals_loss = torch.nn.functional.l1_loss( - pred.normals.float(), target.normals.float() - ) + pred.normals.float(), target.normals.float(), reduction='none' + ) * velocity_equiv_weight + normals_loss = normals_loss.mean() segmentation_loss = torch.nn.functional.l1_loss( - pred.segmentation.float(), target.segmentation.float() - ) + pred.segmentation.float(), target.segmentation.float(), reduction='none' + ) * velocity_equiv_weight + segmentation_loss = segmentation_loss.mean() total_loss = (head_loss + depth_loss + normals_loss + segmentation_loss) / 4.0