feat: add timings

app.py

@@ -1,13 +1,13 @@
+import csv
 import os
-import uuid
 import tempfile
-import csv
+import time
+import uuid
 from pathlib import Path
+import torch
 import numpy as np
 import gradio as gr
-import asyncio
 import spaces
-import torch
 
 # Set matplotlib to use non-interactive backend before importing pyplot
 # This is required for Gradio which runs on worker threads
@@ -43,7 +43,6 @@ if HF_TOKEN is None:
 PLOT_TARGET_POINTS = int(os.environ.get("PLOT_TARGET_POINTS", "1500"))
 SEARCH_MAX_RESULTS = int(os.environ.get("SEARCH_MAX_RESULTS", "50"))
 
-
 # -----------------------------
 # Load pipeline (reloadable)
 # -----------------------------
@@ -70,6 +69,24 @@ load_pipeline(MODEL_ID, DEFAULT_SPECIES)
 # -----------------------------
 # Helpers
 # -----------------------------
+
+_t0 = None
+_tlast = None
+
+def tprint(msg: str):
+    "Function to print timing information"
+    global _t0, _tlast
+    if _t0 is None:
+        _t0 = _tlast = time.perf_counter()
+
+    # CUDA ops are async → synchronize to get real timings
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+    now = time.perf_counter()
+    print(f"[timing] {msg}: {now - _tlast:.3f}s (total {now - _t0:.3f}s)")
+    _tlast = now
+
 def _global_stride(L: int, target: int) -> int:
     if target <= 0 or L <= target:
         return 1
@@ -396,6 +413,8 @@ def predict(
     bigwig_selected: list[str],
     bed_elements: list[str],
 ):
+    tprint("start")
+
     # Debug: verify species is being passed
     if not species:
         raise gr.Error("Species parameter is missing. Please select a species.")
@@ -428,18 +447,24 @@ def predict(
     if "species" not in inputs:
         raise gr.Error(f"Internal error: species not found in inputs dict. Inputs: {list(inputs.keys())}")
 
+    tprint("inputs prepared")
+
     # move to GPU only once the ZeroGPU context is active
     device = "cuda" if torch.cuda.is_available() else "cpu"
     pipe.model.to(device)
     pipe.model.eval()
     print(f"Running on {next(pipe.model.parameters()).device}")
 
+    tprint(f"pipe.model.to({device})")
+
     # run inference
     out = pipe(inputs)
 
+    tprint("inference completed")
+
     # optional: move back to CPU so you don’t rely on any persistent CUDA context
-    if device == "cuda":
-        pipe.model.to("cpu")
+    # if device == "cuda":
+    #     pipe.model.to("cpu")
 
     bw_names = out.bigwig_track_names or []
     bw = out.bigwig_tracks_logits
@@ -516,7 +541,9 @@ def predict(
             eidx = bed_names.index(ename)
             series.append((ename, probs[:, eidx, 1][::stride].astype(float)))
 
+    tprint("figure data processed created")
     fig = _make_tracks_figure(x, series)
+    tprint("figure created")
 
     region = f"{out.chrom}:{out.pred_start}-{out.pred_end}" if out.chrom else f"{x0}-{x1}"
     if out.assembly:
@@ -895,7 +922,7 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
 
 <div class="intro-tip">
   <span class="intro-tip-icon">💡</span>
-  <span><strong>Tip:</strong> The demo includes default settings that you can use to get started, taking ~ 10 seconds to run.</span>
+  <span><strong>Tip:</strong> The demo includes default settings that you can use to get started, taking ~ 15 seconds to run for the example on human.</span>
 </div>
 
 <div style="margin-top: 16px; padding: 12px; background: rgba(0,0,0,0.03); border-radius: 12px; font-size: 0.95rem;">