Fix stage-2 OOM: offload Lyra-2 net to CPU during DA3 GS recon

app.py

@@ -219,16 +219,33 @@ def _stage1_video(
 def _stage2_splat(
     run_dir: Path,
     video_path: Path,
+    res1: Stage1Resources,
     res2: Stage2Resources,
     progress: gr.Progress,
 ) -> Path:
     progress(0.75, desc="Stage 2/2 — reconstructing 3D Gaussian splat (~3 min)")
     t0 = time.monotonic()
-    ply_path = run_stage2_single(
-        res2,
-        video_path=video_path,
-        output_dir=run_dir / "recon",
-    )
+
+    # Offload Lyra-2 diffusion net to CPU so DA3 GS recon has VRAM headroom.
+    # First-run measurement on A100-80GB: with stage-1 resident, DA3's DINOv2
+    # backbone OOMs trying to allocate 1.47 GiB (78.98/79.25 GiB already used).
+    # DA3 stays on GPU since it's reused across both stages.
+    res1.model.net.cpu()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    try:
+        ply_path = run_stage2_single(
+            res2,
+            video_path=video_path,
+            output_dir=run_dir / "recon",
+        )
+    finally:
+        # Restore the diffusion net to GPU for the next stage-1 request.
+        res1.model.net.to(device=res1.desired_device, dtype=res1.desired_dtype)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
     log.info("[stage2] done in %.1fs -> %s", time.monotonic() - t0, ply_path)
     if not ply_path.exists():
         raise gr.Error(f"Stage 2 did not produce {ply_path}.")
@@ -277,7 +294,7 @@ def generate(
 
     try:
         video_path = _stage1_video(run_dir, input_image, prompt.strip(), preset_params, res1, progress)
-        ply_path = _stage2_splat(run_dir, video_path, res2, progress)
+        ply_path = _stage2_splat(run_dir, video_path, res1, res2, progress)
     except gr.Error:
         _CONSECUTIVE_FAILURES += 1
         if torch.cuda.is_available():

entrypoint.sh

@@ -4,6 +4,11 @@
 # 2. Launch the Gradio app.
 set -euo pipefail
 
+# expandable_segments reduces fragmentation when the resident Lyra-2 diffusion
+# net is swapped between GPU/CPU to make room for DA3's stage-2 GS recon pass.
+# Without this, we can still OOM even though aggregate free memory is enough.
+export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
+
 cd /home/user/app
 
 # The checkpoint downloader is idempotent — on warm starts with a populated