DFE cleanup

toolkit/models/diffusion_feature_extraction.py

@@ -1000,7 +1000,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)
+        # velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2)
+        velocity_equiv_weight = 1.0
 
         with torch.no_grad():
             target = self.get_pred(target_0_1)
@@ -1094,7 +1095,6 @@ class DiffusionFeatureExtractor9(nn.Module):
         self.model = Sapiens2(arch="sapiens2_1b", img_size=(1024, 768), patch_size=16).eval().cuda()  # img_size is (H, W)
         self.model.load_state_dict(load_file(ckpt_path))
         self.model.to(device, dtype=dtype)
-        self.model.enable_gradient_checkpointing()
 
         self.losses = {}
         self.log_every = 100

toolkit/models/tipsv2.py

@@ -369,7 +369,8 @@ class VisionTransformer(nn.Module):
         self, x: torch.Tensor, collect_indices: Optional[Sequence[int]] = None
     ):
         collected = [] if collect_indices is not None else None
-        use_ckpt = self.gradient_checkpointing and self.training
+        # gate on is_grad_enabled, not self.training: DFE keeps this in eval
+        use_ckpt = self.gradient_checkpointing and torch.is_grad_enabled()
         for i, blk in enumerate(self.blocks):
             if use_ckpt:
                 x = torch.utils.checkpoint.checkpoint(blk, x, use_reentrant=False)