minor edits

app.py

@@ -1,10 +1,9 @@
 import warnings
 
-# (optional) reduce noisy logs
 warnings.filterwarnings("ignore", message="Can't initialize NVML")
 warnings.filterwarnings("ignore", category=FutureWarning)
 
-# IMPORTANT: import spaces BEFORE torch/cuda gets touched anywhere
+# Must happen before torch/cuda is touched anywhere in imports
 try:
     import spaces  # noqa: F401
 except Exception:

app_v1.py

@@ -4,12 +4,20 @@ import zipfile
 import tempfile
 from functools import lru_cache
 import warnings
+
 warnings.filterwarnings("ignore", message="Can't initialize NVML")
 warnings.filterwarnings("ignore", category=FutureWarning)
+
 try:
-    import spaces  
+    import spaces  # must be imported before torch/cuda usage on ZeroGPU
 except Exception:
-    pass
+    class spaces:  # type: ignore
+        @staticmethod
+        def GPU(*args, **kwargs):
+            def deco(fn):
+                return fn
+            return deco
+
 
 import numpy as np
 import torch
@@ -22,16 +30,15 @@ from einops import rearrange
 
 from models.geometric_deeponet.geometric_deeponet import GeometricDeepONetTime
 
-# ---------------- Config ----------------
+
 REPO_ID = "BGLab/DeepONet-FlowBench-FPO"
 CKPTS = {
-    "1":  "checkpoints/time-dependent-deeponet_1in.ckpt",
-    "4":  "checkpoints/time-dependent-deeponet_4in.ckpt",
-    "8":  "checkpoints/time-dependent-deeponet_8in.ckpt",
+    "1": "checkpoints/time-dependent-deeponet_1in.ckpt",
+    "4": "checkpoints/time-dependent-deeponet_4in.ckpt",
+    "8": "checkpoints/time-dependent-deeponet_8in.ckpt",
     "16": "checkpoints/time-dependent-deeponet_16in.ckpt",
 }
 SAMPLES_DIR = Path("sample_cases")
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 TMP = Path(tempfile.gettempdir())
 
 RANGES = {
@@ -40,8 +47,11 @@ RANGES = {
 }
 
 
+def _device_str() -> str:
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 def _tag() -> str:
-    # unique per request (avoids filename collisions across sessions)
     return next(tempfile._get_candidate_names())
 
 
@@ -51,13 +61,11 @@ def _tmp(tag: str, name: str) -> str:
     return str(out_dir / name)
 
 
-# ---------------- Samples ----------------
 def list_samples():
     if not SAMPLES_DIR.is_dir():
         return []
     ids = []
     for p in SAMPLES_DIR.glob("sample_*_input.npy"):
-        # sample_{id}_input.npy
         sid = p.stem.split("_")[1]
         if sid.isdigit():
             ids.append(sid)
@@ -65,62 +73,64 @@ def list_samples():
 
 
 def load_sample(sample_id: str):
-    sdf = np.load(SAMPLES_DIR / f"sample_{sample_id}_input.npy").astype(np.float32)   # [1,H,W]
-    y = np.load(SAMPLES_DIR / f"sample_{sample_id}_output.npy").astype(np.float32)    # [T,2,H,W]
+    sdf = np.load(SAMPLES_DIR / f"sample_{sample_id}_input.npy").astype(np.float32)  # [1,H,W]
+    y = np.load(SAMPLES_DIR / f"sample_{sample_id}_output.npy").astype(np.float32)  # [T,2,H,W]
     return sdf, y
 
 
-# ---------------- Model ----------------
-@lru_cache(maxsize=4)
-def load_model(history_s: int) -> GeometricDeepONetTime:
+@lru_cache(maxsize=8)
+def load_model(history_s: int, device_str: str) -> GeometricDeepONetTime:
+    device = torch.device(device_str)
     ckpt_path = hf_hub_download(REPO_ID, CKPTS[str(history_s)])
-    model = GeometricDeepONetTime.load_from_checkpoint(ckpt_path, map_location=DEVICE)
-    return model.eval().to(DEVICE)
+    model = GeometricDeepONetTime.load_from_checkpoint(ckpt_path, map_location=device)
+    return model.eval().to(device)
 
 
-def static_tensors(hparams, sdf_np: np.ndarray):
+def static_tensors(hparams, sdf_np: np.ndarray, device: torch.device):
     _, H, W = sdf_np.shape
     x = np.linspace(0.0, float(hparams.domain_length_x), W, dtype=np.float32)
     y = np.linspace(0.0, float(hparams.domain_length_y), H, dtype=np.float32)
     yv, xv = np.meshgrid(y, x, indexing="ij")
     coords = np.stack([xv, yv], axis=0)[None]  # [1,2,H,W]
 
-    sdf_t = torch.from_numpy(sdf_np)[None].to(DEVICE)     # [1,1,H,W]
-    coords_t = torch.from_numpy(coords).to(DEVICE)        # [1,2,H,W]
-    re_t = torch.zeros_like(sdf_t)                        # [1,1,H,W]
+    sdf_t = torch.from_numpy(sdf_np)[None].to(device)  # [1,1,H,W]
+    coords_t = torch.from_numpy(coords).to(device)     # [1,2,H,W]
+    re_t = torch.zeros_like(sdf_t)                     # [1,1,H,W]
     return sdf_t, coords_t, re_t, H, W
 
 
-# ---------------- Rollout + metrics ----------------
 def rollout(sample_id: str, history_s: str):
-    s = int(history_s)
-    model = load_model(s)
+    s_req = int(history_s)
+
+    dev_str = _device_str()
+    device = torch.device(dev_str)
+    model = load_model(s_req, dev_str)
 
     sdf, y_true = load_sample(sample_id)
     T, C, H, W = y_true.shape
     if C != 2:
         raise ValueError(f"Expected 2 channels (u,v), got {C}")
 
-    s = min(s, T - 1)  # ensure s < T
-    sdf_t, coords_t, re_t, _, _ = static_tensors(model.hparams, sdf)
+    s = min(s_req, T - 1)
+    sdf_t, coords_t, re_t, _, _ = static_tensors(model.hparams, sdf, device)
 
     y_pred = np.zeros_like(y_true)
     y_pred[:s] = y_true[:s]
-    history = y_true[:s].copy()  # [s,2,H,W]
+    history = y_true[:s].copy()
 
     for t in range(s, T):
         branch = rearrange(history, "nb c h w -> (nb c) h w")[None]  # [1,s*2,H,W]
-        branch_t = torch.from_numpy(branch).to(DEVICE)
+        branch_t = torch.from_numpy(branch).to(device)
 
         with torch.no_grad():
             y_hat = model((branch_t, re_t, coords_t, sdf_t))  # [1,1,p,2]
 
-        frame = y_hat[0, 0].view(H, W, 2).permute(2, 0, 1).cpu().numpy()  # [2,H,W]
+        frame = y_hat[0, 0].view(H, W, 2).permute(2, 0, 1).cpu().numpy().astype(np.float32)  # [2,H,W]
         y_pred[t] = frame
 
         history = frame[None] if s == 1 else np.concatenate([history[1:], frame[None]], axis=0)
 
-    return y_true, y_pred, s
+    return y_true, y_pred, s, dev_str
 
 
 def rollout_errors(y_true: np.ndarray, y_pred: np.ndarray, s: int):
@@ -132,7 +142,9 @@ def rollout_errors(y_true: np.ndarray, y_pred: np.ndarray, s: int):
     def rel(comp: int):
         d = diff[:, comp].reshape(len(ts), -1)
         t = yt[:, comp].reshape(len(ts), -1)
-        return np.linalg.norm(d, axis=1) / np.linalg.norm(t, axis=1)
+        denom = np.linalg.norm(t, axis=1)
+        denom = np.where(denom == 0.0, 1.0, denom)
+        return np.linalg.norm(d, axis=1) / denom
 
     err_u = rel(0)
     err_v = rel(1)
@@ -140,7 +152,7 @@ def rollout_errors(y_true: np.ndarray, y_pred: np.ndarray, s: int):
 
 
 def pair_png(gt2d: np.ndarray, pred2d: np.ndarray, label: str, t: int) -> bytes:
-    vmin, vmax = RANGES.get(label, (-1.0, 1.0))  # fallback if label changes
+    vmin, vmax = RANGES.get(label, (-1.0, 1.0))
 
     fig, ax = plt.subplots(1, 2, figsize=(6.5, 2.6))
 
@@ -152,7 +164,6 @@ def pair_png(gt2d: np.ndarray, pred2d: np.ndarray, label: str, t: int) -> bytes:
     ax[1].set_title(f"{label} Pred – t={t}")
     ax[1].axis("off")
 
-    # Colorbar height == ax[1] image height
     divider = make_axes_locatable(ax[1])
     cax = divider.append_axes("right", size="5%", pad=0.05)
     fig.colorbar(im2, cax=cax)
@@ -176,7 +187,10 @@ def write_zip(tag: str, y_true: np.ndarray, y_pred: np.ndarray, comp: int, label
     path = _tmp(tag, f"{label}_frames.zip")
     with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
         for t in range(y_true.shape[0]):
-            zf.writestr(f"{label}_frame_{t:03d}.png", pair_png(y_true[t, comp], y_pred[t, comp], label, t))
+            zf.writestr(
+                f"{label}_frame_{t:03d}.png",
+                pair_png(y_true[t, comp], y_pred[t, comp], label, t),
+            )
     return path
 
 
@@ -206,11 +220,10 @@ def write_error_assets(tag: str, ts: np.ndarray, err_u: np.ndarray, err_v: np.nd
     return png, csv
 
 
-# ---------------- Gradio callback ----------------
 def predict_rollout(sample_id: str, history_s: str):
     tag = _tag()
 
-    y_true, y_pred, s = rollout(sample_id, history_s)
+    y_true, y_pred, s, dev_str = rollout(sample_id, history_s)
     ts, err_u, err_v, avg_u, avg_v = rollout_errors(y_true, y_pred, s)
 
     u_gif = write_gif(tag, y_true, y_pred, 0, "u")
@@ -220,6 +233,7 @@ def predict_rollout(sample_id: str, history_s: str):
     err_png, csv = write_error_assets(tag, ts, err_u, err_v)
 
     metrics = (
+        f"Device: {dev_str}\n"
         f"Rollout relative L2 error (averaged over t ≥ {s}):\n"
         f"  u: {avg_u:.3e}\n"
         f"  v: {avg_v:.3e}"
@@ -228,12 +242,16 @@ def predict_rollout(sample_id: str, history_s: str):
     return (u_gif, u_gif, u_zip, v_gif, v_gif, v_zip, err_png, csv, metrics)
 
 
-# ---------------- UI builder ----------------
+@spaces.GPU(duration=180)
+def predict_rollout_gpu(sample_id: str, history_s: str):
+    return predict_rollout(sample_id, history_s)
+
+
 def build_demo():
     sample_choices = list_samples() or ["0"]
 
     return gr.Interface(
-        fn=predict_rollout,
+        fn=predict_rollout_gpu,
         inputs=[
             gr.Radio(sample_choices, value=sample_choices[0], label="Sample ID"),
             gr.Radio(["1", "4", "8", "16"], value="16", label="History length s"),
@@ -249,7 +267,7 @@ def build_demo():
             gr.File(label="Download L2 vs time (CSV)"),
             gr.Textbox(label="Summary metrics"),
         ],
-        title="Time-Dependent DeepONet –  FPO Rollout Demo",
+        title="Time-Dependent DeepONet – FPO Rollout Demo",
         description=(
             "Auto-regressive 60-step rollout of u and v fields for a selected sample. "
             "Choose history length s (1, 4, 8, 16). Download videos/frames and relative error vs time (CSV)."

app_v2.py

@@ -4,12 +4,20 @@ import zipfile
 import tempfile
 from functools import lru_cache
 import warnings
+
 warnings.filterwarnings("ignore", message="Can't initialize NVML")
 warnings.filterwarnings("ignore", category=FutureWarning)
+
 try:
-    import spaces  
+    import spaces  # must be imported before torch/cuda usage on ZeroGPU
 except Exception:
-    pass
+    class spaces:  # type: ignore
+        @staticmethod
+        def GPU(*args, **kwargs):
+            def deco(fn):
+                return fn
+            return deco
+
 
 import numpy as np
 import torch
@@ -22,19 +30,16 @@ from einops import rearrange
 
 from models.geometric_deeponet.geometric_deeponet import GeometricDeepONetTime
 
-# ---------------- Config ----------------
+
 REPO_ID = "BGLab/DeepONet-FlowBench-FPO"
 CKPTS = {
-    "1":  "checkpoints/time-dependent-deeponet_1in.ckpt",
-    "4":  "checkpoints/time-dependent-deeponet_4in.ckpt",
-    "8":  "checkpoints/time-dependent-deeponet_8in.ckpt",
+    "1": "checkpoints/time-dependent-deeponet_1in.ckpt",
+    "4": "checkpoints/time-dependent-deeponet_4in.ckpt",
+    "8": "checkpoints/time-dependent-deeponet_8in.ckpt",
     "16": "checkpoints/time-dependent-deeponet_16in.ckpt",
 }
 
-# v2 samples live here (only 16 GT timesteps per sample)
 SAMPLES_DIR = Path("sample_cases") / "few_timesteps"
-
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 TMP = Path(tempfile.gettempdir())
 
 RANGES = {
@@ -42,15 +47,21 @@ RANGES = {
     "v": (-1.0, 1.0),
 }
 
+
+def _device_str() -> str:
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 def _tag() -> str:
     return next(tempfile._get_candidate_names())
 
+
 def _tmp(tag: str, name: str) -> str:
     out_dir = TMP / f"deeponet_fpo_{tag}"
     out_dir.mkdir(parents=True, exist_ok=True)
     return str(out_dir / name)
 
-# ---------------- Samples ----------------
+
 def list_samples():
     if not SAMPLES_DIR.is_dir():
         return []
@@ -61,32 +72,34 @@ def list_samples():
             ids.append(sid)
     return sorted(set(ids), key=int)
 
+
 def load_sample(sample_id: str):
     sdf = np.load(SAMPLES_DIR / f"sample_{sample_id}_input.npy").astype(np.float32)   # [1,H,W]
     y16 = np.load(SAMPLES_DIR / f"sample_{sample_id}_output.npy").astype(np.float32)  # [16,2,H,W]
     return sdf, y16
 
-# ---------------- Model ----------------
-@lru_cache(maxsize=4)
-def load_model(history_s: int) -> GeometricDeepONetTime:
+
+@lru_cache(maxsize=8)
+def load_model(history_s: int, device_str: str) -> GeometricDeepONetTime:
+    device = torch.device(device_str)
     ckpt_path = hf_hub_download(REPO_ID, CKPTS[str(history_s)])
-    model = GeometricDeepONetTime.load_from_checkpoint(ckpt_path, map_location=DEVICE)
-    return model.eval().to(DEVICE)
+    model = GeometricDeepONetTime.load_from_checkpoint(ckpt_path, map_location=device)
+    return model.eval().to(device)
 
-def static_tensors(hparams, sdf_np: np.ndarray):
-    _, H, W = sdf_np.shape
 
+def static_tensors(hparams, sdf_np: np.ndarray, device: torch.device):
+    _, H, W = sdf_np.shape
     x = np.linspace(0.0, float(hparams.domain_length_x), W, dtype=np.float32)
     y = np.linspace(0.0, float(hparams.domain_length_y), H, dtype=np.float32)
     yv, xv = np.meshgrid(y, x, indexing="ij")
     coords = np.stack([xv, yv], axis=0)[None]  # [1,2,H,W]
 
-    sdf_t = torch.from_numpy(sdf_np)[None].to(DEVICE)  # [1,1,H,W]
-    coords_t = torch.from_numpy(coords).to(DEVICE)     # [1,2,H,W]
+    sdf_t = torch.from_numpy(sdf_np)[None].to(device)  # [1,1,H,W]
+    coords_t = torch.from_numpy(coords).to(device)     # [1,2,H,W]
     re_t = torch.zeros_like(sdf_t)                     # [1,1,H,W]
     return sdf_t, coords_t, re_t, H, W
 
-# ---------------- Rollout ----------------
+
 def rollout_pred(sample_id: str, history_s: str, n_steps: int):
     s = int(history_s)
     n_steps = int(n_steps)
@@ -94,39 +107,48 @@ def rollout_pred(sample_id: str, history_s: str, n_steps: int):
     if n_steps <= 0:
         raise ValueError("Number of rollout steps must be a positive integer.")
     if n_steps < s:
-        n_steps = s  # must have at least s frames to seed
+        n_steps = s
+
+    dev_str = _device_str()
+    device = torch.device(dev_str)
+    model = load_model(s, dev_str)
 
-    model = load_model(s)
     sdf, y16 = load_sample(sample_id)
 
-    # Expect [16,2,H,W] (or more), but we ONLY use first s to seed the model.
     if y16.ndim != 4 or y16.shape[1] != 2:
         raise ValueError(f"Expected y shape [T,2,H,W], got {y16.shape}")
     if y16.shape[0] < s:
         raise ValueError(f"Sample only has {y16.shape[0]} timesteps, but checkpoint needs s={s}.")
 
     _, _, H, W = y16.shape
-    sdf_t, coords_t, re_t, _, _ = static_tensors(model.hparams, sdf)
+    sdf_t, coords_t, re_t, _, _ = static_tensors(model.hparams, sdf, device)
 
-    seed = y16[:s].copy()  # [s,2,H,W] (GT seed only)
+    seed = y16[:s].copy()
     y_out = np.zeros((n_steps, 2, H, W), dtype=np.float32)
     y_out[:s] = seed
 
     history = seed.copy()
     for t in range(s, n_steps):
         branch = rearrange(history, "nb c h w -> (nb c) h w")[None]  # [1,s*2,H,W]
-        branch_t = torch.from_numpy(branch).to(DEVICE)
+        branch_t = torch.from_numpy(branch).to(device)
 
         with torch.no_grad():
             y_hat = model((branch_t, re_t, coords_t, sdf_t))  # [1,1,p,2]
 
-        frame = y_hat[0, 0].view(H, W, 2).permute(2, 0, 1).cpu().numpy().astype(np.float32)  # [2,H,W]
+        frame = (
+            y_hat[0, 0]
+            .view(H, W, 2)
+            .permute(2, 0, 1)
+            .cpu()
+            .numpy()
+            .astype(np.float32)
+        )  # [2,H,W]
         y_out[t] = frame
         history = frame[None] if s == 1 else np.concatenate([history[1:], frame[None]], axis=0)
 
-    return y_out, s
+    return y_out, s, dev_str
+
 
-# ---------------- Rendering (prediction-only) ----------------
 def single_png(field2d: np.ndarray, label: str, t: int) -> bytes:
     vmin, vmax = RANGES.get(label, (-1.0, 1.0))
 
@@ -144,6 +166,7 @@ def single_png(field2d: np.ndarray, label: str, t: int) -> bytes:
     plt.close(fig)
     return buf.getvalue()
 
+
 def write_gif(tag: str, y: np.ndarray, comp: int, label: str) -> str:
     path = _tmp(tag, f"{label}_rollout.gif")
     with imageio.get_writer(path, mode="I", duration=0.1, loop=0) as w:
@@ -152,6 +175,7 @@ def write_gif(tag: str, y: np.ndarray, comp: int, label: str) -> str:
             w.append_data(imageio.imread(io.BytesIO(png)))
     return path
 
+
 def write_zip(tag: str, y: np.ndarray, comp: int, label: str) -> str:
     path = _tmp(tag, f"{label}_frames.zip")
     with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
@@ -159,10 +183,10 @@ def write_zip(tag: str, y: np.ndarray, comp: int, label: str) -> str:
             zf.writestr(f"{label}_frame_{t:03d}.png", single_png(y[t, comp], label, t))
     return path
 
-# ---------------- Gradio callback ----------------
+
 def run_v2(sample_id: str, history_s: str, n_steps: int):
     tag = _tag()
-    y, s = rollout_pred(sample_id, history_s, n_steps)
+    y, s, dev_str = rollout_pred(sample_id, history_s, n_steps)
 
     u_gif = write_gif(tag, y, comp=0, label="u")
     v_gif = write_gif(tag, y, comp=1, label="v")
@@ -170,8 +194,9 @@ def run_v2(sample_id: str, history_s: str, n_steps: int):
     v_zip = write_zip(tag, y, comp=1, label="v")
 
     summary = (
+        f"Device: {dev_str}\n"
         f"Seeded with s={s} timesteps from {SAMPLES_DIR}.\n"
-        f"Generated rollout length N={y.shape[0]} (frames labeled seed for t<s, pred for t≥s)."
+        f"Generated rollout length N={y.shape[0]} (seed frames t<s, predicted frames t≥s)."
     )
 
     return (
@@ -180,13 +205,18 @@ def run_v2(sample_id: str, history_s: str, n_steps: int):
         summary,
     )
 
-# ---------------- UI builder ----------------
+
+@spaces.GPU(duration=180)
+def run_v2_gpu(sample_id: str, history_s: str, n_steps: int):
+    return run_v2(sample_id, history_s, n_steps)
+
+
 def build_demo():
     sample_choices = list_samples() or ["0"]
     history_choices = ["1", "4", "8", "16"]
 
     return gr.Interface(
-        fn=run_v2,
+        fn=run_v2_gpu,
         inputs=[
             gr.Radio(sample_choices, value=sample_choices[0], label="Sample ID"),
             gr.Radio(history_choices, value="16", label="History length s (checkpoint)"),
@@ -201,10 +231,10 @@ def build_demo():
             gr.File(label="Download all v frames (ZIP)"),
             gr.Textbox(label="Run summary"),
         ],
-        title="Time-Dependent DeepONet –  FPO Rollout Demo",
+        title="Time-Dependent DeepONet – FPO Rollout Demo",
         description=(
-            "Auto-regressive rollout of u and v fields for a selected sample. "
-            "Choose history length s (1, 4, 8, 16). Download videos/frames."
+            "Prediction-only auto-regressive rollout of u and v fields for a selected sample. "
+            "Choose history length s (1, 4, 8, 16) and rollout length N."
         ),
     )