Upload train_jit.py with huggingface_hub

train_jit.py

@@ -185,16 +185,13 @@ def build_jit(cfg: JiTTrainConfig) -> JiT:
 
 class CFMFlowWrapper(nn.Module):
     """
-    torchdyn NeuralODE expects f(t, x) with dx/dt returned.
-    - velocity mode: JiT predicts v directly → return model(x, t).
-    - x_pred mode (JiT denoiser): JiT predicts x1 (clean); v = (x_pred - x) / (1-t).
+    torchdyn NeuralODE expects f(t, x) with same batch as x.
+    JiT is forward(x, t) with t shape (N,).
     """
 
-    def __init__(self, model: JiT, prediction_mode: str = "velocity", t_eps: float = 1e-5):
+    def __init__(self, model: JiT):
         super().__init__()
         self.model = model
-        self.prediction_mode = "x_pred"
-        self.t_eps = t_eps
 
     def forward(self, t: torch.Tensor, x: torch.Tensor, y=None, *args, **kwargs) -> torch.Tensor:
         batch = x.shape[0]
@@ -203,12 +200,6 @@ class CFMFlowWrapper(nn.Module):
             t_flat = t_flat.expand(batch)
         elif t_flat.shape[0] != batch:
             t_flat = t_flat[:batch]
-
-        if self.prediction_mode == "x_pred":
-            x_pred = self.model(x, t_flat)
-            one_minus_t = (1.0 - t_flat).clamp(min=self.t_eps)
-            t_bc = one_minus_t.reshape(-1, *([1] * (x.dim() - 1)))
-            return (x_pred - x) / t_bc
         return self.model(x, t_flat)
 
 
@@ -277,14 +268,8 @@ def main() -> None:
             x0 = torch.randn_like(x1)
             t, xt, ut = fm.sample_location_and_conditional_flow(x0, x1)
             t_b = t.reshape(-1).float()
-            # vt = net_model(xt, t_b)
-            # loss = torch.mean((vt - ut) ** 2)
-            x_pred = net_model(xt, t_b)
-            one_minus_t = (1.0 - t_b).clamp(min=1.0e-5)
-            t_bc = one_minus_t.reshape(-1, *([1] * (xt.dim() - 1)))
-            v_pred = (x_pred - xt) / t_bc
-            v_target = (x1 - xt) / t_bc
-            loss = torch.mean((v_target - v_pred) ** 2)
+            vt = net_model(xt, t_b)
+            loss = torch.mean((vt - ut) ** 2)
 
             optim.zero_grad(set_to_none=True)
             loss.backward()