diff --git a/animation/optimization.py b/animation/optimization.py index 1739a46..073fd08 100644 --- a/animation/optimization.py +++ b/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) diff --git a/animation/utils/save_utils.py b/animation/utils/save_utils.py index c46e4a4..483af10 100644 --- a/animation/utils/save_utils.py +++ b/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" diff --git a/precompute_depth.py b/precompute_depth.py new file mode 100644 index 0000000..275f6f3 --- /dev/null +++ b/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 /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()