From 92d38cbed4280ceee305a81d36d53750d2b33da7 Mon Sep 17 00:00:00 2001
From: Lars Bogner <upyxo@student.kit.edu>
Date: Mon, 22 Jun 2026 08:23:58 +0200
Subject: [PATCH] Add --batch-size auto to predict, matching train

Estimates batch size from free GPU memory using the checkpoint's
hidden_dim/n_blocks, same as the train command.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 giant/cli.py | 43 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/giant/cli.py b/giant/cli.py
index ab71be6..cf846ec 100644
--- a/giant/cli.py
+++ b/giant/cli.py
@@ -255,8 +255,14 @@ def predict(
         ),
     ] = None,
     batch_size: Annotated[
-        int, typer.Option("--batch-size", "-b", help="Inference batch size")
-    ] = 4096,
+        str,
+        typer.Option(
+            "--batch-size",
+            "-b",
+            help="Inference batch size, or 'auto' to estimate from free GPU "
+            "memory (cuda devices only)",
+        ),
+    ] = "4096",
     steps: Annotated[
         int, typer.Option("--steps", "-s", help="Flow matching ODE steps")
     ] = 10,
@@ -266,6 +272,21 @@ def predict(
     ] = None,
 ) -> None:
     """Run trained model on a parquet file and save predictions."""
+    batch_size_auto = False
+    batch_size_value: Optional[int] = None
+    if batch_size.strip().lower() == "auto":
+        batch_size_auto = True
+    else:
+        try:
+            batch_size_value = int(batch_size)
+        except ValueError:
+            typer.echo(
+                f"error: --batch-size must be an integer or 'auto', "
+                f"got {batch_size!r}",
+                err=True,
+            )
+            raise typer.Exit(1)
+
     _device = torch.device(device) if device else gconfig.auto_device()
     typer.echo(f"device: {_device}")
 
@@ -279,6 +300,20 @@ def predict(
         raise typer.Exit(1)
 
     model_cfg = ckpt["model_config"]
+
+    if batch_size_auto:
+        try:
+            batch_size_value = gconfig.estimate_batch_size(
+                model_cfg["hidden_dim"], model_cfg["n_blocks"], _device
+            )
+        except ValueError as exc:
+            typer.echo(f"error: {exc}", err=True)
+            raise typer.Exit(1)
+        typer.echo(
+            f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
+        )
+    assert batch_size_value is not None
+    bs = batch_size_value
     pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
     mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
     cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
@@ -323,8 +358,8 @@ def predict(
 
             # Inference in batch_size slices
             pred_parts = []
-            for start in range(0, N, batch_size):
-                end = min(start + batch_size, N)
+            for start in range(0, N, bs):
+                end = min(start + bs, N)
                 cc = torch.from_numpy(cond_cont[start:end]).float().to(_device)
                 ck = torch.from_numpy(cond_cat[start:end]).long().to(_device)
                 pred_parts.append(sample_flow(model, cc, ck, steps=steps).cpu().numpy())