fix

app.py

@@ -87,11 +87,13 @@ def create_user_temp_dir():
 # Global model initialization for Spatial Tracker
 print("🚀 Initializing tracking models...")
 
-vggt4track_model = VGGT4Track.from_pretrained(
-    "Yuxihenry/SpatialTrackerV2_Front")
+vggt4track_model = VGGT4Track.from_pretrained("Yuxihenry/SpatialTrackerV2_Front")
 vggt4track_model.eval()
 vggt4track_model = vggt4track_model.to("cuda")
 
+if not hasattr(vggt4track_model, 'infer'):
+    vggt4track_model.infer = vggt4track_model.forward
+
 tracker_model = Predictor.from_pretrained("Yuxihenry/SpatialTrackerV2-Offline")
 tracker_model.eval()
 
@@ -208,39 +210,78 @@ def render_from_pointcloud(rgb_frames, depth_frames, intrinsics, original_extrin
     return {'rendered': output_path, 'motion_signal': motion_signal_path, 'mask': mask_path}
 
 
+
 @spaces.GPU
-def run_spatial_tracker(video_tensor):
-    if not hasattr(vggt4track_model, "infer"):
-        vggt4track_model.infer = lambda x: vggt4track_model(x)
-    if tracker_model.spatrack.base_model is None:
-        tracker_model.spatrack.base_model = vggt4track_model
+def run_spatial_tracker(video_tensor: torch.Tensor):
+    """
+    GPU-intensive spatial tracking function.
 
+    Args:
+        video_tensor: Preprocessed video tensor (T, C, H, W)
+
+    Returns:
+        Dictionary containing tracking results
+    """
+    # Run VGGT to get depth and camera poses
     video_input = preprocess_image(video_tensor)[None].cuda()
+
     with torch.no_grad():
-        with torch.amp.autocast('cuda', dtype=torch.bfloat16):
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
             predictions = vggt4track_model(video_input / 255)
-            extrinsic, intrinsic = predictions["poses_pred"], predictions["intrs"]
-            depth_map, depth_conf = predictions["points_map"][...,
-                                                              2], predictions["unc_metric"]
-
-    depth_tensor, extrs, intrs = depth_map.squeeze().cpu().numpy(
-    ), extrinsic.squeeze().cpu().numpy(), intrinsic.squeeze().cpu().numpy()
+            extrinsic = predictions["poses_pred"]
+            intrinsic = predictions["intrs"]
+            depth_map = predictions["points_map"][..., 2]
+            depth_conf = predictions["unc_metric"]
+
+    depth_tensor = depth_map.squeeze().cpu().numpy()
+    extrs = extrinsic.squeeze().cpu().numpy()
+    intrs = intrinsic.squeeze().cpu().numpy()
+    video_tensor_gpu = video_input.squeeze()
     unc_metric = depth_conf.squeeze().cpu().numpy() > 0.5
+
+    # Setup tracker
     tracker_model.spatrack.track_num = 512
     tracker_model.to("cuda")
-    grid_pts = get_points_on_a_grid(
-        30, (video_input.shape[3], video_input.shape[4]), device="cpu")
-    query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[
-        0].numpy()
 
+    # Get grid points for tracking
+    frame_H, frame_W = video_tensor_gpu.shape[2:]
+    grid_pts = get_points_on_a_grid(30, (frame_H, frame_W), device="cpu")
+    query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[0].numpy()
+
+    # Run tracker
     with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-        c2w_traj, intrs_out, point_map, conf_depth, _, _, _, _, video_out = tracker_model.forward(
-            video_input.squeeze(), depth=depth_tensor, intrs=intrs, extrs=extrs, queries=query_xyt,
-            fps=1, unc_metric=unc_metric, support_frame=len(video_input.squeeze())-1
+        (
+            c2w_traj, intrs_out, point_map, conf_depth,
+            track3d_pred, track2d_pred, vis_pred, conf_pred, video_out
+        ) = tracker_model.forward(
+            video_tensor_gpu, depth=depth_tensor,
+            intrs=intrs, extrs=extrs,
+            queries=query_xyt,
+            fps=1, full_point=False, iters_track=4,
+            query_no_BA=True, fixed_cam=False, stage=1,
+            unc_metric=unc_metric,
+            support_frame=len(video_tensor_gpu)-1, replace_ratio=0.2
         )
-    return {'video_out': video_out.cpu(), 'point_map': point_map.cpu(), 'conf_depth': conf_depth.cpu(), 'intrs_out': intrs_out.cpu(), 'c2w_traj': c2w_traj.cpu()}
 
-# --- TTM WAN INFERENCE FUNCTION ---
+    # Resize outputs for rendering
+    max_size = 384
+    h, w = video_out.shape[2:]
+    scale = min(max_size / h, max_size / w)
+    if scale < 1:
+        new_h, new_w = int(h * scale), int(w * scale)
+        video_out = T.Resize((new_h, new_w))(video_out)
+        point_map = T.Resize((new_h, new_w))(point_map)
+        conf_depth = T.Resize((new_h, new_w))(conf_depth)
+        intrs_out[:, :2, :] = intrs_out[:, :2, :] * scale
+
+    # Move results to CPU and return
+    return {
+        'video_out': video_out.cpu(),
+        'point_map': point_map.cpu(),
+        'conf_depth': conf_depth.cpu(),
+        'intrs_out': intrs_out.cpu(),
+        'c2w_traj': c2w_traj.cpu(),
+    }
 
 
 @spaces.GPU