Upload app.py with huggingface_hub

app.py

@@ -93,21 +93,29 @@ def _apply_schedule(h_cycles: int, l_cycles: int) -> None:
     LAST_SCHEDULE["h"], LAST_SCHEDULE["l"] = h_cycles, l_cycles
 
 
-# Preset schedules from the paper's Table 4 sweep: (H, L) -> unrolled depth
-# 16 * H * (L + 1). H2L3 (the trained schedule) is the default.
+# Coherent schedules from the paper's Table 4 sweep: (label, H, L) -> unrolled
+# depth 16 * H * (L+1). H2L3 is the trained schedule and clearly the best; the
+# H1/H4 rows in the paper collapse (MMStar 0.5/24), so they are not offered.
 SCHEDULES = [
-    ("H1L1 (32 layers, fastest)", 1, 1),
-    ("H1L3 (64 layers)", 1, 3),
-    ("H2L1 (64 layers)", 2, 1),
-    ("H2L3 (128 layers, default)", 2, 3),
-    ("H2L5 (192 layers)", 2, 5),
-    ("H4L3 (256 layers)", 4, 3),
+    ("H2L3 · 128 layers (trained default, best)", 2, 3),
+    ("H2L2 · 96 layers (lighter, still solid)", 2, 2),
+    ("H3L3 · 192 layers (heavier, slightly lower)", 3, 3),
+    ("H3L2 · 144 layers", 3, 2),
+    ("H2L1 · 64 layers (noticeably degraded)", 2, 1),
+    ("H3L1 · 96 layers (noticeably degraded)", 3, 1),
 ]
 _SCHEDULE_LOOKUP = {name: (h, l) for name, h, l in SCHEDULES}
 
 
 def _estimate_duration(*args, **kwargs) -> int:
-    """Callable GPU duration: base time grows with unrolled depth and token budget."""
+    """Callable GPU duration, calibrated from measured calls on this Space.
+
+    Measured on ZeroGPU: ~0.081 s/token at 128 unrolled layers, roughly
+    proportional to depth; ~2.5 s fixed plus up to ~9 s one-time init on a
+    visitor's first call. The 20 s base covers init + prefill + margin.
+
+    Duration = 20 + budget x (0.12 s/token scaled by depth/128), capped.
+    """
     schedule = kwargs.get("schedule") or (args[2] if len(args) > 2 else None)
     budget = kwargs.get("max_tokens") or (args[3] if len(args) > 3 else 32)
     try:
@@ -115,14 +123,14 @@ def _estimate_duration(*args, **kwargs) -> int:
     except Exception:
         h, l = 2, 3
     depth = 16 * h * (l + 1)
-    return min(300, 20 + int(depth * 0.09) + int(budget * 0.35))
+    return min(120, 20 + int(float(budget) * 0.12 * depth / 128))
 
 
 @spaces.GPU(duration=_estimate_duration)
 def answer(
     image: Image.Image,
     prompt: str,
-    schedule: str = "H2L3 (128 layers, default)",
+    schedule: str = "H2L3 · 128 layers (trained default, best)",
     max_tokens: int = 32,
     use_trigger: bool = False,
     progress=gr.Progress(track_tqdm=True),
@@ -208,9 +216,11 @@ with gr.Blocks() as demo:
                 )
                 schedule_in = gr.Dropdown(
                     choices=[name for name, _, _ in SCHEDULES],
-                    value="H2L3 (128 layers, default)",
+                    value="H2L3 · 128 layers (trained default, best)",
                     label="Recurrent schedule (compute budget)",
-                    info="Unrolled depth = 16 × H × (L+1) effective layers per forward pass",
+                    info="Unrolled depth = 16 × H × (L+1) effective layers per forward pass. "
+                         "H2L3 is the trained schedule; other coherent schedules from the "
+                         "paper's Table 4 sweep are offered for comparison.",
                 )
                 run_btn = gr.Button("Answer", variant="primary")
                 with gr.Accordion("Advanced", open=False):