Add CPU-first inference with automatic SDPA fallback

.gitignore

@@ -16,4 +16,4 @@ weights/
 .DS_Store
 **/.DS_Store
 docs/
-
+checkpoints/

README.md

@@ -123,13 +123,27 @@ pip install --index-url https://pypi.org/simple flashinfer-python
 > `--index-url https://pypi.org/simple` is only needed if your default pip index is an internal mirror that doesn't have `flashinfer-python`.
 > (Optional) For faster first-use, you can additionally install a CUDA-specific JIT cache: `pip install flashinfer-jit-cache -f https://flashinfer.ai/whl/cu128/flashinfer-jit-cache/`.
 > See [FlashInfer installation](https://docs.flashinfer.ai/installation.html) for details. If FlashInfer is not installed, the model falls back to SDPA (PyTorch native attention) via `--use_sdpa`.
+> On CPU (`--device cpu`, the default), SDPA is selected automatically — FlashInfer is not required.
 
 **5. Visualization dependencies (optional)**
 
 ```bash
 pip install -e ".[vis]"
 ```
 
+**6. CPU / low-spec machines (optional)**
+
+If you do not have a capable GPU, cannot install FlashInfer, or have limited VRAM (the checkpoint is ~4.4 GB), install CPU PyTorch and run with `--device cpu`:
+
+```bash
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+pip install -e .
+python demo.py --model_path /path/to/checkpoint.pt \
+    --image_folder /path/to/images/ --device cpu
+```
+
+`--device auto` uses CUDA only when free VRAM looks sufficient (~8 GiB+); otherwise it falls back to CPU. Laptop GPUs with ≤4–6 GB VRAM typically need `--device cpu`.
+
 ## 📦 Model Download
 
 | Model Name | Huggingface Repository | ModelScope Repository | Description |
@@ -149,6 +163,8 @@ python demo.py --model_path /path/to/lingbot-map-long.pt \
     --image_folder example/courthouse --mask_sky
 ```
 
+> **CPU / low-spec:** `--device` defaults to `cpu` (SDPA, no FlashInfer required). Use `--device cuda` or `--device auto` when you have sufficient GPU VRAM.
+
 This launches an interactive [viser](https://github.com/nerfstudio-project/viser) viewer at `http://localhost:8080`. See [Interactive Demo](#-interactive-demo-demopy) below for the full set of scenes and flags, or jump to [Offline Rendering Pipeline](#-offline-rendering-pipeline-demo_renderbatch_demopy) for long-sequence batch rendering.
 
 ## 🎬 Interactive Demo (`demo.py`)
@@ -305,10 +321,21 @@ python demo.py --model_path /path/to/checkpoint.pt \
     --image_folder /path/to/images/ --use_sdpa
 ```
 
+SDPA is also selected automatically on CPU (`--device cpu`, the default) or when FlashInfer is not installed — `--use_sdpa` is only needed to force SDPA on CUDA.
+
+#### CPU inference (low-spec / no GPU)
+
+```bash
+python demo.py --model_path /path/to/checkpoint.pt \
+    --image_folder /path/to/images/ --device cpu \
+    --keyframe_interval 4 --num_scale_frames 4
+```
+
 #### Running on Limited GPU Memory
 
 If you run into out-of-memory issues, try one (or both) of the following:
 
+- **`--device cpu`** or **`--device auto`** — prefer host RAM when VRAM cannot hold the ~4.4 GB checkpoint (typical on ≤4–6 GB laptop GPUs).
 - **`--offload_to_cpu`** — offload per-frame predictions to CPU during inference (on by default; use `--no-offload_to_cpu` only if you have memory to spare).
 - **`--num_scale_frames 2`** — reduce the number of bidirectional scale frames from the default 8 down to 2, which shrinks the activation peak of the initial scale phase.
 

benchmark/configs/methods/lingbot_map.yaml

@@ -1,10 +1,15 @@
 model: lingbot_map
 env: lingbot-map
 _checkpoint: /path/to/lingbot-map.pt
-_device: cuda
+# Device: cpu (portable default), cuda, or auto (CUDA only if free VRAM ≥ ~8 GiB).
+# Laptop / low-VRAM GPUs should keep cpu — the checkpoint alone is ~4.4 GB.
+_device: cpu
 _mode: streaming
+# AMP is ignored / disabled on CPU by the method wrapper.
 _use_amp: true
-_use_sdpa: false  # Force FlashInfer backend (paged KV cache attention)
+# true = force SDPA. false = try FlashInfer on CUDA when installed; SDPA otherwise.
+# FlashInfer is optional and not required for correct results.
+_use_sdpa: true
 _image_size: 518
 _patch_size: 14
 _enable_3d_rope: true

benchmark/configs/methods/lingbot_map_v1.yaml

@@ -1,10 +1,12 @@
 model: lingbot_map
 env: lingbot-map
 _checkpoint: /path/to/lingbot-map.pt
-_device: cuda
+# Device: cpu (portable default), cuda, or auto (CUDA only if free VRAM ≥ ~8 GiB).
+_device: cpu
 _mode: streaming
 _use_amp: true
-_use_sdpa: false  # Force FlashInfer backend (paged KV cache attention)
+# true = force SDPA. false = try FlashInfer on CUDA when installed; SDPA otherwise.
+_use_sdpa: true
 _image_size: 518
 _patch_size: 14
 _enable_3d_rope: true

benchmark/methods/lingbot_map.py

@@ -13,6 +13,12 @@
 
 from benchmark.method.base import BaseMethod
 from benchmark.core.loader import BSSLoader
+from lingbot_map.utils.device import (
+    autocast_context,
+    resolve_device,
+    resolve_dtype,
+    select_attention_backend,
+)
 
 
 # Mirrors lingbot-map/demo.py:413-421 — above this frame count, the KV cache
@@ -48,7 +54,7 @@ class LingbotMapMethod(BaseMethod):
     def __init__(
         self,
         checkpoint: str = None,
-        device: str = 'cuda',
+        device: str = 'cpu',
         mode: str = 'streaming',
         use_amp: bool = True,
         use_sdpa: bool = False,
@@ -77,10 +83,13 @@ def __init__(
         )
 
         self.checkpoint = checkpoint
-        self.device = device
+        # Resolve device early so AMP / SDPA defaults are correct on CPU.
+        self.device = str(resolve_device(device))
         self.mode = mode
-        self.use_amp = use_amp
-        self.use_sdpa = use_sdpa
+        self.use_amp = use_amp and self.device.startswith("cuda")
+        # Auto-select SDPA on CPU or when FlashInfer is missing; honor explicit True.
+        backend = select_attention_backend(self.device, force_sdpa=bool(use_sdpa))
+        self.use_sdpa = backend == "sdpa"
         self.image_size = image_size
         self.patch_size = patch_size
         self.enable_3d_rope = enable_3d_rope
@@ -143,20 +152,22 @@ def _prepare_images(self, rgb_list):
         from torchvision import transforms as TF
 
         to_tensor = TF.ToTensor()
+        # Keep batch on CPU; inference_* moves frames to the model device.
         images = torch.stack([to_tensor(rgb) for rgb in rgb_list])
-        return images.to(self.device)
+        return images
 
     def _run_inference(self, images):
         """Run LingbotMap inference and return raw predictions dict."""
-        if self.use_amp:
-            dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+        device = torch.device(self.device)
+        if self.use_amp and device.type == "cuda":
+            dtype = resolve_dtype(device)
         else:
             dtype = torch.float32
 
-        print(f"  → Running {self.mode} inference (dtype: {dtype})")
+        print(f"  → Running {self.mode} inference (dtype: {dtype}, device: {device}, sdpa: {self.use_sdpa})")
 
         num_frames = images.shape[0]
-        with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+        with torch.no_grad(), autocast_context(device, dtype):
             if self.mode == 'streaming':
                 keyframe_interval = _resolve_keyframe_interval(
                     self.keyframe_interval, num_frames, self.auto_keyframe_threshold

demo.py

@@ -47,6 +47,13 @@
 from lingbot_map.utils.pose_enc import pose_encoding_to_extri_intri
 from lingbot_map.utils.geometry import closed_form_inverse_se3_general
 from lingbot_map.utils.load_fn import load_and_preprocess_images
+from lingbot_map.utils.device import (
+    autocast_context,
+    flashinfer_usable,
+    resolve_device,
+    resolve_dtype,
+    select_attention_backend,
+)
 
 
 # =============================================================================
@@ -211,15 +218,17 @@ def _warm_streaming(model, images, scale_frames, warm_stream_n, dtype,
     warm_stream_n = max(1, min(int(warm_stream_n), num_avail - scale_frames))
     kf_int = max(int(keyframe_interval), 1)
 
-    # images: [S, 3, H, W] on device already; slice + add batch dim, no copy of
-    # spatial dims so warmup shape == real inference shape (H, W).
-    warm_scale = images[:scale_frames].unsqueeze(0).to(dtype)
-    warm_stream = images[scale_frames:scale_frames + warm_stream_n].unsqueeze(0).to(dtype)
+    # images stay on CPU host; move warmup slices to the model device below.
+    model_device = next(model.parameters()).device
+    warm_scale = images[:scale_frames].unsqueeze(0).to(device=model_device, dtype=dtype)
+    warm_stream = images[scale_frames:scale_frames + warm_stream_n].unsqueeze(0).to(
+        device=model_device, dtype=dtype
+    )
 
     for _ in range(passes):
         model.clean_kv_cache()
         torch.compiler.cudagraph_mark_step_begin()
-        with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+        with torch.no_grad(), autocast_context(model_device, dtype):
             model.forward(
                 warm_scale,
                 num_frame_for_scale=scale_frames,
@@ -231,7 +240,7 @@ def _warm_streaming(model, images, scale_frames, warm_stream_n, dtype,
             if not is_keyframe:
                 model._set_skip_append(True)
             torch.compiler.cudagraph_mark_step_begin()
-            with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+            with torch.no_grad(), autocast_context(model_device, dtype):
                 model.forward(
                     warm_stream[:, i:i + 1],
                     num_frame_for_scale=scale_frames,
@@ -379,15 +388,25 @@ def main():
     parser.add_argument("--camera_num_iterations", type=int, default=4,
                         help="Camera head iterative-refinement steps. Default 4; set 1 for faster inference "
                             "(skips 3 refinement passes at a small accuracy cost).")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        choices=["cpu", "cuda", "auto"],
+        help="Inference device. Default: cpu (portable). "
+             "'cuda' forces GPU when available; "
+             "'auto' uses CUDA only if free VRAM looks sufficient (~8 GiB+).",
+    )
     parser.add_argument("--use_sdpa", action="store_true", default=False,
-                        help="Use SDPA backend (no flashinfer needed). Default: FlashInfer")
+                        help="Force SDPA attention backend (no FlashInfer). "
+                             "SDPA is selected automatically on CPU or when FlashInfer is missing.")
     parser.add_argument("--compile", action="store_true", default=False,
                         help="torch.compile hot modules (reduce-overhead) with a CUDA-graph warmup. "
-                            "Streaming mode only; ~5 FPS faster at 518x378. Adds ~30-60 s warmup time.")
+                            "Streaming mode + CUDA only; ~5 FPS faster at 518x378. Adds ~30-60 s warmup time.")
     parser.add_argument(
         "--offload_to_cpu",
         action=argparse.BooleanOptionalAction,
-        default=False,
+        default=True,
         help="Offload per-frame predictions to CPU during inference to cut GPU peak memory "
             "(on by default).  Use --no-offload_to_cpu to keep outputs on GPU.",
     )
@@ -418,7 +437,21 @@ def main():
     assert args.image_folder or args.video_path, \
         "Provide --image_folder or --video_path"
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = resolve_device(args.device)
+    backend = select_attention_backend(device, force_sdpa=args.use_sdpa)
+    args.use_sdpa = backend == "sdpa"
+    fi_ok, fi_reason = flashinfer_usable()
+    dtype = resolve_dtype(device)
+
+    print(f"Device: {device} (requested={args.device})")
+    print(f"Attention backend: {backend}"
+          + ("" if fi_ok else f" (FlashInfer skipped: {fi_reason})"))
+    print(f"Dtype: {dtype}")
+    if device.type == "cpu":
+        print(
+            "CPU mode: expected to be much slower than GPU; "
+            "FlashInfer is not required. Tip: export OMP_NUM_THREADS to match cores."
+        )
 
     # ── Load images & model ──────────────────────────────────────────────────
     t0 = time.time()
@@ -444,26 +477,21 @@ def main():
     model = load_model(args, device)
     print(f"Total load time: {time.time() - t0:.1f}s")
 
-    # Pick inference dtype; autocast still runs for the ops that need fp32 (e.g. LayerNorm).
-    if torch.cuda.is_available():
-        dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
-    else:
-        dtype = torch.float32
-
     # Cast the aggregator (DINOv2-style trunk) to the inference dtype to remove the
     # redundant fp32 master weight copy + autocast bf16 weight cache (~2-3 GB saved,
     # no measurable quality change). gct_base._predict_* upcasts inputs to fp32 and
     # runs each head under `autocast(enabled=False)`, so camera/depth/point heads
-    # keep fp32 weights automatically.
+    # keep fp32 weights automatically. Skipped on CPU (already fp32).
     if dtype != torch.float32 and getattr(model, "aggregator", None) is not None:
         print(f"Casting aggregator to {dtype} (heads kept in fp32)")
         model.aggregator = model.aggregator.to(dtype=dtype)
 
-    images = images.to(device)
+    # Keep the full image batch on CPU; inference_* moves frames to the model device.
+    images = images.cpu()
     num_frames = images.shape[0]
-    print(f"Input: {num_frames} frames, shape {tuple(images.shape)}")
+    print(f"Input: {num_frames} frames, shape {tuple(images.shape)} (host/CPU storage)")
     print(f"Mode: {args.mode}")
-    if torch.cuda.is_available():
+    if device.type == "cuda":
         torch.cuda.empty_cache()
         print(
             f"GPU mem after load: "
@@ -498,9 +526,11 @@ def main():
                 f"(window_size={args.window_size} keyframes, scale={args.num_scale_frames})."
             )
 
-    # ── Optional: torch.compile + CUDA-graph warmup (streaming only) ────────
+    # ── Optional: torch.compile + CUDA-graph warmup (streaming + CUDA only) ─
     if args.compile:
-        if args.mode != "streaming":
+        if device.type != "cuda":
+            print("--compile requires CUDA; skipping on CPU.")
+        elif args.mode != "streaming":
             print(
                 f"--compile only applies to --mode streaming (got {args.mode!r}); "
                 "skipping compile."
@@ -537,12 +567,12 @@ def main():
             print(f"  compiled warmup: {time.time() - t_warm:.1f}s")
 
     # ── Inference ────────────────────────────────────────────────────────────
-    print(f"Running {args.mode} inference (dtype={dtype})...")
+    print(f"Running {args.mode} inference (dtype={dtype}, device={device}, backend={backend})...")
     t0 = time.time()
 
     output_device = torch.device("cpu") if args.offload_to_cpu else None
 
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+    with torch.no_grad(), autocast_context(device, dtype):
         if args.mode == "streaming":
             predictions = model.inference_streaming(
                 images,
@@ -562,7 +592,7 @@ def main():
             )
 
     print(f"Inference done in {time.time() - t0:.1f}s")
-    if torch.cuda.is_available():
+    if device.type == "cuda":
         print(
             f"GPU peak during inference: "
             f"{torch.cuda.max_memory_allocated()/1e9:.2f} GB "

gct_profile.py

@@ -252,6 +252,12 @@ def main():
     if args.keyframe_interval < 1:
         parser.error(f"--keyframe_interval must be >= 1 (got {args.keyframe_interval})")
 
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "gct_profile.py requires CUDA. For portable inference use "
+            "`python demo.py --device cpu` (SDPA, no FlashInfer)."
+        )
+
     dtype_map = {'bf16': torch.bfloat16, 'fp32': torch.float32}
     backends = ['sdpa', 'flashinfer'] if args.backend == 'both' else [args.backend]
     dtypes = ['bf16', 'fp32'] if args.dtype == 'both' else [args.dtype]

lingbot_map/aggregator/stream.py

@@ -88,9 +88,26 @@ def __init__(
         kwargs.pop('use_flexflash', None)
         use_sdpa = kwargs.pop('use_sdpa', False)
 
-        # Backend selection: SDPA (no extra deps) or FlashInfer (paged KV cache)
-        self.use_sdpa = use_sdpa
-        self.use_flashinfer = not use_sdpa  # FlashInfer is default unless SDPA requested
+        # Backend selection: SDPA (portable) or FlashInfer (optional CUDA).
+        # Soft-fallback when FlashInfer was requested but is not importable — avoid
+        # crashing at first forward / FlashInferBlock construction.
+        if use_sdpa or not use_flashinfer:
+            self.use_sdpa = True
+            self.use_flashinfer = False
+        else:
+            from lingbot_map.utils.device import flashinfer_usable
+
+            usable, reason = flashinfer_usable()
+            if usable:
+                self.use_sdpa = False
+                self.use_flashinfer = True
+            else:
+                logger.warning(
+                    "FlashInfer not usable (%s); falling back to SDPA backend.",
+                    reason,
+                )
+                self.use_sdpa = True
+                self.use_flashinfer = False
 
         # Call parent __init__
         super().__init__(**kwargs)

lingbot_map/layers/attention.py

@@ -377,7 +377,10 @@ def __init__(
         kv_cache_camera_only: bool = False,
     ) -> None:
         if not FLASHINFER_AVAILABLE:
-            raise RuntimeError("FlashInfer is not available. Please install flashinfer.")
+            raise RuntimeError(
+                "FlashInfer is not available. Install flashinfer-python for the "
+                "CUDA fast path, or use the SDPA backend (--use_sdpa / use_sdpa=True)."
+            )
 
         super().__init__(
             dim=dim,