Support <=12GB GPUs: lazy depth model, frame resize, lower tracker caps

animation/optimization.py

@@ -296,13 +296,21 @@ def optimization(
             quats_to_optimize, global_quats, global_trans, n_iters
         )
 
-        # Initialize depth module and flow weights
-        depth_module = DepthModule(
-            encoder='vitl',
-            device=self.device,
-            input_size=images_batch.shape[1],
-            fp32=False
-        )
+        # Initialize depth module lazily: if depth cache already exists
+        # (e.g. written by precompute_depth.py), skip loading the ~1.5GB
+        # Video-Depth-Anything model so <=12GB GPUs don't OOM.
+        depth_cache = os.path.join(flow_dirs.replace('flow', 'depth'),
+                                   'depth_gt_raw.pt')
+        if os.path.exists(depth_cache):
+            print("Depth cache found — skipping VideoDepthAnything load")
+            depth_module = None
+        else:
+            depth_module = DepthModule(
+                encoder='vitl',
+                device=self.device,
+                input_size=images_batch.shape[1],
+                fp32=False
+            )
 
         # Prepare masks
         real_rgb = images_batch[..., :3]
@@ -467,6 +475,8 @@ def load_and_prepare_data(args):
     images = []
     for f in img_files:
         img = Image.open(f).convert("RGBA")
+        if img.size != (args.img_size, args.img_size):
+            img = img.resize((args.img_size, args.img_size), Image.LANCZOS)
         arr = np.array(img, dtype=np.float32) / 255.0
         t = torch.from_numpy(arr).to(args.device)
         images.append(t)

animation/utils/save_utils.py

@@ -238,8 +238,10 @@ def save_track_points(point_vis_mask, renderer, model, img_path, out_dir, args):
     os.makedirs(track_2d_point_path, exist_ok=True)
 
     num_visible = len(visible_indices)
-    MAX_VISIBLE_POINTS = 15000
-    MAX_SAMPLE_POINTS = 4000
+    # Lowered from 15000/4000 for <=12GB GPUs: dense meshes (200k+ verts)
+    # OOM CoTracker's corr-volume einsum in "full" mode.
+    MAX_VISIBLE_POINTS = 6000
+    MAX_SAMPLE_POINTS = 2000
 
     # Determine tracking strategy
     tracking_mode = "full" if num_visible <= MAX_VISIBLE_POINTS else "sampled"

precompute_depth.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""
+precompute_depth.py — Run Video-Depth-Anything standalone so the heavy model
+is freed before Puppeteer's optimization loop starts (fixes OOM on <=12GB
+GPUs). Writes <seq>/depth/depth_gt_raw.pt in exactly the format
+utils/data_loader.prepare_depth caches.
+
+Usage (from animation/ dir, venv active):
+  PYTHONPATH=. python ../precompute_depth.py --input_path ../examples --seq_name deer --img_size 512
+"""
+import argparse
+import glob
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+
+from utils.loss_utils import DepthModule
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--input_path", required=True)
+    p.add_argument("--seq_name", required=True)
+    p.add_argument("--img_size", type=int, default=512)
+    args = p.parse_args()
+
+    base = os.path.join(args.input_path, args.seq_name)
+    img_paths = sorted(glob.glob(os.path.join(base, "imgs", "frame_*.png")))
+    assert img_paths, f"no frames under {base}/imgs"
+
+    frames = []
+    for pth in img_paths:
+        img = Image.open(pth).convert("RGB").resize((args.img_size, args.img_size))
+        frames.append(np.asarray(img, dtype=np.float32) / 255.0)
+    batch = torch.from_numpy(np.stack(frames))  # [T,H,W,3] in 0..1
+
+    depth = DepthModule(encoder="vitl", device="cuda")
+    with torch.no_grad():
+        depth_gt_raw = depth.get_depth_maps(batch)
+    os.makedirs(os.path.join(base, "depth"), exist_ok=True)
+    out = os.path.join(base, "depth", "depth_gt_raw.pt")
+    torch.save(depth_gt_raw.cpu(), out)
+    print("wrote", out, tuple(depth_gt_raw.shape))
+
+
+if __name__ == "__main__":
+    main()