SIMIT demo: BAGEL, Lance, Qwen3.8-27B-FP8 + FLUX.2-klein on ZeroGPU

.gitattributes

@@ -33,3 +33,24 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/bagel_amazon/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_amazon/demo_1.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_amazon/query.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_infovqa_val_449/query.jpg filter=lfs diff=lfs merge=lfs -text
+examples/bagel_ok_vqa_val2014_492_always/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_ok_vqa_val2014_492_always/demo_1.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_345_always/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_345_always/demo_1.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_345_always/query.jpg filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_47_always/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_47_always/demo_1.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_47_always/query.jpg filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_87_always/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_87_always/demo_1.png filter=lfs diff=lfs merge=lfs -text
+examples/bagel_vizwiz_vqa_val_87_always/query.jpg filter=lfs diff=lfs merge=lfs -text
+examples/lance_ok_vqa_val2014_253/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/lance_ok_vqa_val2014_253/query.jpg filter=lfs diff=lfs merge=lfs -text
+examples/lance_ok_vqa_val2014_8_always/demo_0.png filter=lfs diff=lfs merge=lfs -text
+examples/lance_ok_vqa_val2014_8_always/demo_1.png filter=lfs diff=lfs merge=lfs -text
+examples/lance_ok_vqa_val2014_8_always/demo_2.png filter=lfs diff=lfs merge=lfs -text
+examples/lance_ok_vqa_val2014_8_always/demo_3.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,77 @@
 ---
-title: Simit
-emoji: 💻
-colorFrom: pink
-colorTo: green
+title: SIMIT - models that imagine their own practice examples
+emoji: 🥯
+colorFrom: yellow
+colorTo: red
 sdk: gradio
 sdk_version: 6.29.1
-python_version: '3.12'
+python_version: "3.12"
 app_file: app.py
 pinned: false
+license: apache-2.0
+startup_duration_timeout: 1h
+short_description: A VLM imagines practice examples for your question
+models:
+  - ByteDance-Seed/BAGEL-7B-MoT
+  - bytedance-research/Lance
+  - Qwen/Qwen3.8-27B-FP8
+  - black-forest-labs/FLUX.2-klein-4B
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SIMIT demo
+
+Upload an image, ask a question, and compare the model's plain answer with its
+SIMIT-ICL answer: before answering again, the model imagines a few practice
+examples (question, answer decided first, an image made to fit), checks them,
+and keeps them in context.
+
+Built on the [`simit`](https://github.com/monurcan/simit) library. See `tools/` for how the per-budget
+hyperparameters were chosen and how the cached examples were produced.
+
+## Hardware notes (ZeroGPU)
+
+* Each request runs in a fresh `@spaces.GPU` process and uses the visitor's own
+  ZeroGPU quota for the selected time budget (20-180 s).
+* BAGEL-7B-MoT and Lance run on the default 48 GB size; Qwen3.8-27B-FP8 +
+  FLUX.2-klein needs the 96 GB size (2x quota).
+* All weights stay on CPU in the main process (~91 GB of RAM, ~92 GB of disk for
+  the model files). If the Space's disk is too small, attach persistent storage
+  and set `HF_HOME=/data/.huggingface`, or limit the models with
+  `SIMIT_DEMO_MODELS=bagel,lance`.
+* The first request after start-up compiles and autotunes GPU kernels (Qwen:
+  ~25 s); later requests reuse the on-disk cache.
+
+## How the time budgets map to settings
+
+Each budget (20 / 40 / 60 / 90 / 180 s of GPU time per request) uses its own settings per model
+(`presets.json`, built by `tools/build_presets.py`):
+
+1. **Accuracy per K_max** (`tools/tune_presets.py`): one synthesis cache with 4 candidates per query on a
+   mixed 64-question general-QA validation set: 8 each from VizWiz, VQAv2, OK-VQA, TextVQA, ChartQA,
+   DocVQA, MMBench and AI2D (LMMs-Eval-Lite rows 300-307), each scored with its own metric. Then 800
+   Optuna trials over the ABA/DF thresholds and K_max ∈ {1,2,3,4}. Per K_max, the pick is regularized:
+   within one validation query of the best score, the widest difficulty-filter band.
+2. **Latency** (`tools/latency.py`): per-request timelines on the demo's own code path in fresh forked
+   processes (as on ZeroGPU, weight transfer included). Measured on an H100 and scaled by spaces'
+   own factor for the ZeroGPU GPU sizes (1.5× on `large`, 1× on `xlarge`).
+3. **Choice**: the tuned (best-quality) speed setting whenever some K fits the budget for ≥80% of
+   requests; among those, the K within one validation query of the best score with the widest band.
+   More budget also buys more attempts and verification rounds. If no K beats zero-shot on
+   validation, SIMIT keeps the base answer by default; "Imagine even when confident" still runs it.
+
+Validation scores (64 general questions):
+
+| model | zero-shot | K=1 | K=2 | K=3 | K=4 | used by default |
+|---|---|---|---|---|---|---|
+| BAGEL-7B-MoT | 0.703 | 0.708 | **0.750** | 0.760 | 0.755 | K=2 from 60 s (20-40 s: zero-shot) |
+| Lance | 0.432 | 0.442 | 0.442 | 0.458 | **0.468** | K=1 at 20 s, K=4 from 40 s |
+| Qwen3.8-27B-FP8 + FLUX | **0.854** | 0.839 | 0.823 | 0.839 | 0.844 | zero-shot (no gain found) |
+
+Lance's answers are sampled (greedy decoding degenerates on this checkpoint), so its outputs vary
+between runs and SIMIT can also hurt a correct answer. Qwen3.8-27B is strong enough that imagined
+examples did not improve it on these questions.
+
+The cached examples are real runs of this app (BAGEL-7B-MoT and Lance) on queries from the paper's
+figures, plus one product-page screenshot (`examples/bagel_amazon`). Several of them used "Imagine even when confident" (shown in the example), because the base
+model was confidently wrong (e.g. "Unanswerable"). Qwen3.8-27B already answered those queries
+correctly, so it has no cached example.

app.py

@@ -0,0 +1,297 @@
+"""SIMIT demo: a vision-language model imagines its own practice examples for
+your question, then answers again with them in context.
+
+Runs on Hugging Face ZeroGPU (each request uses the visitor's own GPU quota)
+and on any machine with a CUDA GPU (``python app.py``).
+"""
+
+import os
+
+# ZeroGPU starts a fresh process per request: keep Triton's autotuning results on disk
+# (otherwise every request re-tunes the FP8 / linear-attention kernels, ~20 s).
+os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")
+os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")
+
+import spaces  # before anything CUDA-related: ZeroGPU patches torch  # noqa: E402  # isort: skip
+
+import base64  # noqa: E402
+import html  # noqa: E402
+import json  # noqa: E402
+from pathlib import Path  # noqa: E402
+
+import gradio as gr  # noqa: E402
+from PIL import Image  # noqa: E402
+
+from demo_models import BUDGETS, ENABLED, MODELS, code_snippet, run  # noqa: E402
+
+HERE = Path(__file__).parent
+ASSETS = HERE / "assets"
+EXAMPLES_DIR = HERE / "examples"
+
+
+def _data_uri(path: Path) -> str:
+    return "data:image/webp;base64," + base64.b64encode(path.read_bytes()).decode()
+
+
+MASCOT = {p.stem: _data_uri(p) for p in ASSETS.glob("*.webp")}
+STAGES = {  # stage -> (mascot, message)
+    "idle": ("initial", "Ask me anything about an image"),
+    "loading": ("initial", "Loading the model weights (first request only)"),
+    "warmup": ("initial", "Waking up the model on the GPU"),
+    "solving": ("initial", "Solving it on its own first"),
+    "synthesizing": ("synthesizer", "Writing practice questions it already knows the answers to"),
+    "routing": ("router", "Choosing how to draw each practice image"),
+    "drawing": ("artist", "Painting a practice image"),
+    "coding": ("architect", "Writing code for a chart, diagram or document"),
+    "checking": ("critic", "Checking that each picture matches its question"),
+    "improving": ("improved", "Answering again, with its own practice examples in context"),
+}
+LABELS = {MODELS[k].label: k for k in ENABLED}
+BUDGET_LABELS = [f"{b} s" for b in BUDGETS]
+GPU_OVERHEAD = 20  # seconds requested on top of the budget (worker start-up, final answer); weight transfer counts inside it
+
+
+def _seconds(label: str) -> int:
+    return int(str(label).split()[0])
+
+
+# --------------------------------------------------------------------- HTML
+def status_html(stage: str, elapsed: float = 0.0, limit: float = 0.0, n_demos: int = 0, note: str = "") -> str:
+    mascot, message = STAGES.get(stage, STAGES["solving"])
+    dots = "" if stage == "idle" else '<span class="simit-dots"></span>'
+    bar = ""
+    if limit:
+        pct = max(2.0, min(100.0, 100.0 * elapsed / limit))
+        bar = (f'<div class="simit-bar"><div style="width:{pct:.0f}%"></div></div>'
+               f'<div class="simit-sub">{elapsed:.0f} s of {limit:.0f} s'
+               + (f' · {n_demos} practice example{"s" if n_demos != 1 else ""} ready' if n_demos else "")
+               + "</div>")
+    return (f'<div class="simit-status"><img class="simit-mascot" src="{MASCOT[mascot]}" alt="">'
+            f'<div><div class="simit-msg">{html.escape(message)}{dots}</div>'
+            f'{bar}{f"<div class=simit-sub>{html.escape(note)}</div>" if note else ""}</div></div>')
+
+
+def done_html(note: str) -> str:
+    return (f'<div class="simit-status done"><img class="simit-mascot still" src="{MASCOT["improved"]}" alt="">'
+            f'<div><div class="simit-msg">Done</div><div class="simit-sub">{html.escape(note)}</div></div></div>')
+
+
+def error_html(note: str) -> str:
+    return (f'<div class="simit-status error"><img class="simit-mascot still" src="{MASCOT["critic"]}" alt="">'
+            f'<div><div class="simit-msg">Something went wrong</div><div class="simit-sub">{html.escape(note)}'
+            '</div></div></div>')
+
+
+def answer_card(kind: str, answer: str = "", sub: str = "", mark: str = "") -> str:
+    title = "Base model (greedy)" if kind == "base" else "SIMIT-ICL (self-improved)"
+    body = html.escape(answer) if answer else '<span class="simit-wait">…</span>'
+    size = " long" if len(answer) > 60 else ""
+    badge = {"right": '<span class="simit-ok">✓</span>', "wrong": '<span class="simit-no">✗</span>'}.get(mark, "")
+    return (f'<div class="simit-card {kind}"><div class="simit-card-title">{title}{badge}</div>'
+            f'<div class="simit-answer{size}">{body}</div><div class="simit-card-sub">{html.escape(sub)}</div></div>')
+
+
+def demo_caption(d) -> str:
+    bits = [f"Q: {d.question}", f"A: {d.answer}"]
+    meta = [d.skill] if d.skill else []
+    if d.verify_score is not None:
+        meta.append(f"critic {d.verify_score}/100")
+    if d.confidence is not None:
+        meta.append(f"confidence {d.confidence:.2f}")
+    return "\n".join(bits) + (f"\n({', '.join(meta)})" if meta else "")
+
+
+# ------------------------------------------------------------- GPU requests
+def _duration(image, question, key, budget_label, always, *args, **kwargs):
+    return _seconds(budget_label) + GPU_OVERHEAD
+
+
+def _session(image, question, key, budget_label, always):
+    """Runs inside the GPU call: turns pipeline events into UI updates."""
+    spec, budget = MODELS[key], _seconds(budget_label)
+    try:
+        yield from _events(spec, image, question, budget, always)
+    except Exception as e:  # show it in the page instead of a bare error toast
+        print(f"[demo] request failed: {type(e).__name__}: {e}", flush=True)
+        yield (error_html(f"{type(e).__name__}: {str(e)[:300]}"), answer_card("base"), answer_card("simit"), [], "")
+
+
+def _events(spec, image, question, budget, always):
+    greedy, k_budget, demos, gallery = None, 0, [], []
+    base, improved = answer_card("base"), answer_card("simit")
+    status = status_html("warmup")
+    for event in run(spec, image, question, budget, always):
+        kind = event[0]
+        if kind == "stage":
+            _, stage, info = event
+            status = status_html(stage, info.get("elapsed", 0), info.get("limit", 0), len(demos))
+        elif kind == "greedy":
+            _, zs, k_budget = event
+            greedy = zs
+            sub = f"confidence p0 = {zs.confidence:.2f}" if zs.confidence is not None else ""
+            base = answer_card("base", zs.answer, sub)
+        elif kind == "demo":
+            demos.append(event[1])
+            gallery = [(d.image, demo_caption(d)) for d in demos]
+        else:  # final
+            _, answer, info = event
+            reason = info.get("skipped")
+            if reason == "confident":
+                sub = "The model was already confident, so SIMIT kept its answer (adaptive budget: 0 examples)."
+            elif reason == "policy":
+                sub = ("On general questions, our tuning found no gain from imagination for this model at this "
+                       "budget, so SIMIT keeps the base answer. Tick 'Imagine even when confident' to try it anyway.")
+            elif reason == "no_time":
+                sub = "Not enough time budget left to imagine examples; kept the base answer."
+            elif reason == "none_passed":
+                sub = "No imagined example passed the checks in time; kept the base answer."
+            else:
+                changed = "changed" if answer.strip() != greedy.answer.strip() else "kept"
+                sub = f"{changed} the answer after {info['demos']} imagined example{'s' if info['demos'] > 1 else ''}"
+            improved = answer_card("simit", answer, sub)
+            status = done_html(f"{info['elapsed']:.0f} s on the GPU")
+        yield status, base, improved, gallery, _details(greedy, k_budget, demos, spec, budget)
+
+
+@spaces.GPU(duration=_duration, size="large")
+def gpu_large(image, question, key, budget_label, always):
+    yield from _session(image, question, key, budget_label, always)
+
+
+@spaces.GPU(duration=_duration, size="xlarge")
+def gpu_xlarge(image, question, key, budget_label, always):
+    yield from _session(image, question, key, budget_label, always)
+
+
+def _details(greedy, k_budget, demos, spec, budget) -> str:
+    if greedy is None:
+        return ""
+    p0 = f"{greedy.confidence:.2f}" if greedy.confidence is not None else "n/a"
+    return (f"**{spec.label}**, {budget} s budget. Zero-shot confidence p0 = {p0}, so the adaptive budget asked "
+            f"for **{k_budget}** example{'s' if k_budget != 1 else ''}; **{len(demos)}** passed the checks. "
+            "Each example is a question the model wrote, whose answer it decided first, with an image made "
+            "to fit that answer.")
+
+
+def submit(image, question, model_label, budget_label, always):
+    if image is None:
+        raise gr.Error("Please upload an image.")
+    if not question or not question.strip():
+        raise gr.Error("Please type a question about the image.")
+    key = LABELS[model_label]
+    spec = MODELS[key]
+    image = image.convert("RGB")
+    if spec._weights is None:  # main process, CPU: no GPU time is spent here
+        yield (status_html("loading"), answer_card("base"), answer_card("simit"), [], "")
+        spec.load()
+    fn = gpu_xlarge if spec.gpu_size == "xlarge" else gpu_large
+    yield from fn(image, question.strip(), key, budget_label, always)
+
+
+# ------------------------------------------------------------------ cached
+def _load_examples():
+    index = EXAMPLES_DIR / "index.json"
+    return json.loads(index.read_text()) if index.exists() else []
+
+
+EXAMPLES = [e for e in _load_examples() if e["model"] in ENABLED]
+
+
+def show_cached(image, question, model_label, budget_label, always=False):
+    """A cached example: the stored outputs of a real run, shown without using the GPU."""
+    entry = next((e for e in EXAMPLES if e["question"] == question and MODELS[e["model"]].label == model_label
+                  and f"{e['budget']} s" == budget_label and bool(e.get("always")) == bool(always)), None)
+    if entry is None:
+        return status_html("idle", note="Press Submit to run this example."), answer_card("base"), \
+            answer_card("simit"), [], ""
+    d = EXAMPLES_DIR / entry["id"]
+    mark = lambda ok: "right" if ok else "wrong"  # noqa: E731
+    gt = f"reference answer: {entry['reference']}"
+    base = answer_card("base", entry["greedy"], f"confidence p0 = {entry['p0']:.2f} · {gt}", mark(entry["greedy_ok"]))
+    n = len(entry["demos"])
+    improved = answer_card("simit", entry["simit"], f"after {n} imagined example{'s' if n != 1 else ''} · {gt}",
+                           mark(entry["simit_ok"]))
+    gallery = [(str(d / x["image"]), "\n".join([f"Q: {x['question']}", f"A: {x['answer']}",
+                                                f"({x['skill']}" + (f", critic {x['verify_score']}/100" if
+                                                                    x.get("verify_score") is not None else "") + ")"]))
+               for x in entry["demos"]]
+    note = (f"Cached result from a real run ({entry['elapsed']:.0f} s on an H100, {entry['budget']} s budget"
+            + (", imagining even when confident" if entry.get("always") else "") + f"). Source: {entry['source']}.")
+    details = (f"**{MODELS[entry['model']].label}**, {entry['budget']} s budget. Zero-shot confidence "
+               f"p0 = {entry['p0']:.2f}; {n} imagined example{'s' if n != 1 else ''} passed the checks.")
+    return done_html(note), base, improved, gallery, details
+
+
+# ---------------------------------------------------------------------- UI
+CSS = (ASSETS / "style.css").read_text()
+INTRO = """
+<div class="simit-hero">
+  <h1>SIMIT: models that imagine their own practice examples</h1>
+  <p>Ask any question about an image. The model first answers on its own. Then, without any labels, it writes
+  similar practice questions whose answers it decides first, makes an image for each one, keeps only those it
+  can verify, and answers your question again with them in context.</p>
+</div>
+"""
+
+with gr.Blocks(title="SIMIT demo") as demo:
+    gr.HTML(INTRO)
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=5):
+            image = gr.Image(type="pil", label="Image", height=360)
+            question = gr.Textbox(label="Question", placeholder="e.g. In what country would you find this hat?",
+                                  lines=2)
+            model = gr.Dropdown(list(LABELS), value=next(iter(LABELS)), label="Model")
+            budget = gr.Radio(BUDGET_LABELS, value="60 s", label="Max GPU time per request",
+                              info="More time lets the model imagine and check more practice examples.")
+            with gr.Accordion("Advanced", open=False):
+                always = gr.Checkbox(False, label="Imagine even when the model is already confident",
+                                     info="By default SIMIT skips imagination for confident answers.")
+            run_btn = gr.Button("Submit", variant="primary")
+        with gr.Column(scale=6):
+            status = gr.HTML(status_html("idle", note="Upload an image and type a question, or pick an example below."))
+            with gr.Row():
+                base_out = gr.HTML(answer_card("base"), min_width=260)
+                simit_out = gr.HTML(answer_card("simit"), min_width=260)
+            with gr.Accordion("See the imagined practice examples", open=False):
+                details = gr.Markdown()
+                gallery = gr.Gallery(columns=4, height=300, object_fit="contain", show_label=False)
+            with gr.Accordion("Use SIMIT in your own code", open=False):
+                code = gr.Code(code_snippet(MODELS[next(iter(LABELS.values()))], 60, ""), language="python",
+                               interactive=False)
+
+    outputs = [status, base_out, simit_out, gallery, details]
+    # On ZeroGPU each request gets its own GPU process; run locally, requests share this process and GPU.
+    from spaces.config import Config as _SpacesConfig
+    run_btn.click(submit, [image, question, model, budget, always], outputs,
+                  concurrency_limit=4 if _SpacesConfig.zero_gpu else 1)
+
+    def _code(model_label, budget_label, q, force):
+        return code_snippet(MODELS[LABELS[model_label]], _seconds(budget_label), q or "", force)
+    for comp in (model, budget, question, always):
+        comp.change(_code, [model, budget, question, always], code, queue=False)
+
+    if EXAMPLES:
+        gr.Markdown("### Examples\nCached results from real runs of this demo (click one; press Submit to "
+                    "run it again live).")
+        examples = gr.Gallery(
+            [(str(EXAMPLES_DIR / e["id"] / e["query_file"]),
+              f"{e['question'].splitlines()[0]}  ({MODELS[e['model']].label})") for e in EXAMPLES],
+            columns=4, height=300 * ((len(EXAMPLES) + 3) // 4) + 20, object_fit="cover", allow_preview=False,
+            show_label=False)
+
+        def pick(evt: gr.SelectData):
+            e = EXAMPLES[evt.index]
+            inputs = [Image.open(EXAMPLES_DIR / e["id"] / e["query_file"]), e["question"],
+                      MODELS[e["model"]].label, f"{e['budget']} s", bool(e.get("always"))]
+            return (*inputs, *show_cached(*inputs))
+        examples.select(pick, None, [image, question, model, budget, always] + outputs)
+    gr.Markdown("Models: " + ", ".join(f"[{MODELS[k].label}](https://huggingface.co/{MODELS[k].repo})"
+                                      for k in ENABLED)
+                + ". Each request uses your own ZeroGPU quota for the selected time budget.")
+
+if os.environ.get("SIMIT_DEMO_PRELOAD", "1") == "1":
+    for key in ENABLED:  # load every model on CPU before serving (main process, no CUDA)
+        MODELS[key].load()
+
+if __name__ == "__main__":
+    demo.queue(max_size=32).launch(css=CSS, theme=gr.themes.Soft(primary_hue="orange", secondary_hue="amber"))

assets/style.css

@@ -0,0 +1,32 @@
+.simit-hero h1 { font-size: 1.7rem; margin: 0.2rem 0 0.4rem; }
+.simit-hero p { margin: 0 0 0.6rem; opacity: 0.85; max-width: 60rem; }
+
+.simit-status { display: flex; align-items: center; gap: 1rem; min-height: 132px; padding: 0.6rem 0.8rem;
+  border-radius: 14px; background: var(--block-background-fill); border: 1px solid var(--border-color-primary); }
+.simit-mascot { height: 112px; width: auto; animation: simit-bob 1.6s ease-in-out infinite; transform-origin: 50% 90%; }
+.simit-mascot.still { animation: simit-pop 0.5s ease-out 1; }
+.simit-msg { font-size: 1.05rem; font-weight: 600; }
+.simit-sub { font-size: 0.85rem; opacity: 0.75; margin-top: 0.3rem; }
+.simit-dots::after { content: ""; animation: simit-dots 1.4s steps(4, end) infinite; }
+.simit-bar { height: 8px; width: 16rem; max-width: 100%; border-radius: 4px; margin-top: 0.5rem;
+  background: var(--border-color-primary); overflow: hidden; }
+.simit-bar > div { height: 100%; border-radius: 4px; transition: width 0.4s linear;
+  background: linear-gradient(90deg, #f2a65a, #e07a3f); }
+
+.simit-card { border-radius: 14px; padding: 0.8rem 1rem; min-height: 120px;
+  border: 1px solid var(--border-color-primary); background: var(--block-background-fill); }
+.simit-card.base { border-top: 5px solid #c9c4bd; }
+.simit-card.simit { border-top: 5px solid #e07a3f; }
+.simit-card-title { font-size: 0.85rem; text-transform: uppercase; letter-spacing: 0.04em; opacity: 0.7;
+  display: flex; justify-content: space-between; }
+.simit-answer { font-size: 1.35rem; font-weight: 650; margin: 0.4rem 0; white-space: pre-wrap; }
+.simit-card-sub { font-size: 0.8rem; opacity: 0.7; }
+.simit-wait { opacity: 0.4; }
+.simit-ok { color: #2e9b5a; font-size: 1.1rem; }
+.simit-no { color: #c94a3c; font-size: 1.1rem; }
+
+@keyframes simit-bob { 0%, 100% { transform: translateY(0) rotate(-2deg); } 50% { transform: translateY(-8px) rotate(2deg); } }
+@keyframes simit-pop { 0% { transform: scale(0.8); } 70% { transform: scale(1.06); } 100% { transform: scale(1); } }
+@keyframes simit-dots { 0% { content: ""; } 25% { content: "."; } 50% { content: ".."; } 75% { content: "..."; } }
+@media (prefers-reduced-motion: reduce) { .simit-mascot, .simit-dots::after { animation: none; } }
+.simit-answer.long { font-size: 1.0rem; font-weight: 550; line-height: 1.4; }

demo_models.py

@@ -0,0 +1,251 @@
+"""Models of the demo and how one request runs on ZeroGPU.
+
+ZeroGPU forks the Space process for every ``@spaces.GPU`` call and kills the
+fork afterwards. So the weights are loaded once, on CPU, in the main process
+(no CUDA there), and each call moves only the chosen model to the GPU and
+builds a fresh SIMIT engine around it.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from simit import SIMIT, SIMITConfig
+
+HERE = Path(__file__).parent
+PRESETS = json.loads((HERE / "presets.json").read_text())
+BUDGETS = [20, 40, 60, 90, 180]           # seconds of GPU time per request
+ANSWER_TOKENS = 128
+ANSWER_RESERVE = 5.0                      # seconds kept for the final SIMIT-ICL answer
+
+
+@dataclass
+class ModelSpec:
+    key: str
+    label: str
+    repo: str
+    kind: str                               # "bagel" | "lance" | "hf"
+    gpu_size: str                           # ZeroGPU size: "large" (48 GB) or "xlarge" (96 GB)
+    image_generator: Optional[str] = None
+    blurb: str = ""
+    _weights: object = None
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    # ----------------------------------------------------- main process (CPU)
+    def load(self):
+        """Load the weights on CPU (main process; nothing touches CUDA)."""
+        with self._lock:
+            if self._weights is None:
+                t = time.time()
+                self._weights = _LOADERS[self.kind](self)
+                print(f"[demo] loaded {self.label} on CPU in {time.time() - t:.0f}s", flush=True)
+        return self._weights
+
+    # ------------------------------------------------- inside @spaces.GPU call
+    def build(self, preset: dict) -> SIMIT:
+        """Move the weights to the GPU and wrap them in a SIMIT engine."""
+        w = self.load()
+        # Without ZeroGPU every request runs in this process: keep one model on the GPU at a time.
+        # (On ZeroGPU each request is a fresh fork in which nothing is on the GPU yet.)
+        if _ON_GPU[0] not in (None, self.key):
+            MODELS[_ON_GPU[0]]._to("cpu")
+            torch.cuda.empty_cache()
+        _ON_GPU[0] = self.key
+        cfg = SIMITConfig(**{k: v for k, v in preset.items() if k in SIMITConfig.__dataclass_fields__})
+        if self.kind == "bagel":
+            from simit.backends.bagel import BagelBackend
+            backend = BagelBackend(w, device="cuda", image_steps=preset.get("image_steps"), use_cuda_graphs=False)
+            return SIMIT(backend, config=cfg, model_id=self.repo)
+        if self.kind == "lance":
+            from simit.backends.lance import LanceBackend
+            backend = LanceBackend(w, device="cuda", image_steps=preset.get("image_steps"), use_cuda_graphs=False)
+            return SIMIT(backend, config=cfg, model_id=self.repo)
+        from simit.backends.hf import HFBackend
+        from simit.backends.imagegen import DiffusersImageGenerator
+        model, processor, pipe = w
+        model.to("cuda")
+        pipe.to("cuda")
+        backend = HFBackend(model, processor=processor, max_batch_size=8, max_batch_tokens=32768)
+        return SIMIT(backend, image_generator=DiffusersImageGenerator(pipe), config=cfg, model_id=self.repo)
+
+    def _to(self, device: str):
+        if self.kind in ("bagel", "lance"):
+            self._weights.to(device)
+        else:
+            model, _, pipe = self._weights
+            model.to(device)
+            pipe.to(device)
+
+    def preset(self, budget: int) -> dict:
+        return dict(PRESETS[self.key][str(budget)])
+
+
+_ON_GPU = [None]   # key of the model whose weights are on the GPU in this process
+
+
+@contextlib.contextmanager
+def _fp8_without_gpu():
+    """transformers dequantizes FP8 checkpoints to bf16 when it sees no GPU at load
+    time. FP8 weights only need a GPU to *run*, so let the load keep them in FP8
+    on CPU; they are moved to the GPU inside each request."""
+    saved = torch.cuda.is_available, torch.cuda.get_device_capability
+    torch.cuda.is_available = lambda: True
+    torch.cuda.get_device_capability = lambda *a, **k: (9, 0)
+    try:
+        yield
+    finally:
+        torch.cuda.is_available, torch.cuda.get_device_capability = saved
+
+
+def _load_bagel(spec):
+    from simit.backends.bagel import BagelModel
+    return BagelModel.load(spec.repo, device="cpu")
+
+
+def _load_lance(spec):
+    from simit.backends.lance import LanceModel
+    return LanceModel.load(spec.repo, device="cpu")
+
+
+def _load_hf(spec):
+    from diffusers import DiffusionPipeline
+    from transformers import AutoModelForMultimodalLM, AutoProcessor
+
+    from simit.backends.hf import _dense_fp8_fix
+    os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")  # Triton FP8 kernels: no nvcc JIT at runtime
+    processor = AutoProcessor.from_pretrained(spec.repo)
+    kw = dict(device_map="cpu", dtype=torch.bfloat16)
+    _dense_fp8_fix(spec.repo, kw)  # transformers would skip FP8 conversion of gate_proj for this checkpoint
+    with _fp8_without_gpu():
+        model = AutoModelForMultimodalLM.from_pretrained(spec.repo, **kw)
+    model.eval()
+    pipe = DiffusionPipeline.from_pretrained(spec.image_generator, dtype=torch.bfloat16)
+    return model, processor, pipe
+
+
+_LOADERS = {"bagel": _load_bagel, "lance": _load_lance, "hf": _load_hf}
+
+MODELS = {m.key: m for m in [
+    ModelSpec("bagel", "BAGEL-7B-MoT", "ByteDance-Seed/BAGEL-7B-MoT", "bagel", "large",
+              blurb="Unified model: draws its own photos and writes code for charts and diagrams."),
+    ModelSpec("lance", "Lance", "bytedance-research/Lance", "lance", "large",
+              blurb="Small unified model: fast native image generation."),
+    ModelSpec("qwen", "Qwen3.8-27B (FP8) + FLUX.2-klein", "Qwen/Qwen3.8-27B-FP8", "hf", "xlarge",
+              image_generator="black-forest-labs/FLUX.2-klein-4B",
+              blurb="Standard VLM: FLUX.2-klein draws the photos; it writes code for charts and diagrams."),
+]}
+ENABLED = [k for k in os.environ.get("SIMIT_DEMO_MODELS", "bagel,lance,qwen").split(",") if k in MODELS]
+
+
+# ---------------------------------------------------------------- a request
+# Pipeline call tags (``on_progress``) -> what the loading screen shows.
+STAGE_OF_TAG = {"synthesis": "synthesizing", "synthesis_fresh": "synthesizing", "revise": "synthesizing",
+                "route": "routing", "skill_spec": "coding", "render": "coding", "image": "drawing",
+                "critic": "checking", "critic_score": "checking"}
+
+
+def run(spec: ModelSpec, image, question: str, budget: int, always: bool):
+    """Generator of progress events for one request (runs inside the GPU call):
+    ("stage", name, info) / ("greedy", ZeroShot, k) / ("demo", Demo) / ("final", answer, info)."""
+    t_start = time.time()
+    preset = spec.preset(budget)
+    yield ("stage", "warmup", {"elapsed": 0.0})
+    sim = spec.build(preset)
+    try:
+        yield ("stage", "solving", {"elapsed": time.time() - t_start})
+        zs = sim.zero_shot(image, question, max_new_tokens=ANSWER_TOKENS)   # the greedy answer + p0 (ABA input)
+        k_budget = sim.config.budget(zs.confidence)
+        yield ("greedy", zs, k_budget)
+
+        # what is left of the budget, keeping room for the final answer (measured: ~1-1.5 s after the last demo)
+        time_left = budget - (time.time() - t_start) - ANSWER_RESERVE
+        if (k_budget == 0 and not always) or time_left < 3:
+            if time_left < 3:
+                reason = "no_time"
+            else:
+                reason = "confident" if preset.get("imagine_by_default", True) else "policy"
+            yield ("final", zs.answer, {"demos": 0, "skipped": reason, "elapsed": time.time() - t_start})
+            return
+
+        events: "queue.Queue" = queue.Queue()
+        current = ["synthesis"]
+        result = {}
+
+        def work():
+            try:
+                result["det"] = sim.imagine(image, question, k=preset["k_max"] if always else None,
+                                            return_details=True, max_new_tokens=ANSWER_TOKENS, time_limit=time_left,
+                                            on_demo=events.put, on_progress=lambda tag: current.__setitem__(0, tag))
+            except Exception as e:  # surfaced below
+                result["error"] = e
+        th = threading.Thread(target=work, daemon=True)
+        th.start()
+        while th.is_alive() or not events.empty():
+            try:
+                yield ("demo", events.get(timeout=0.4))
+                continue
+            except queue.Empty:
+                pass
+            stage = STAGE_OF_TAG.get(current[0].split(":")[0], "synthesizing")
+            yield ("stage", stage, {"elapsed": time.time() - t_start, "limit": budget})
+        if "error" in result:
+            raise result["error"]
+        det = result["det"]
+        if not det.demos:
+            yield ("final", zs.answer, {"demos": 0, "skipped": "none_passed", "elapsed": time.time() - t_start,
+                                        "stats": det.stats})
+            return
+        yield ("stage", "improving", {"elapsed": time.time() - t_start})
+        answer = sim.answer(image, question, det.demos, max_new_tokens=ANSWER_TOKENS)
+        yield ("final", answer, {"demos": len(det.demos), "elapsed": time.time() - t_start, "stats": det.stats})
+    finally:
+        sim.close()
+
+
+def code_snippet(spec: ModelSpec, budget: int, question: str, always: bool = False) -> str:
+    """Copy-paste code reproducing this request with the library."""
+    p = spec.preset(budget)
+    force = always or not p.get("imagine_by_default", True)
+    keys = ["k_max", "epsilon", "A0", "A1", "B", "t_low", "t_high", "verify", "verify_think", "use_skills",
+            "attempts_per_slot", "verify_rounds"]
+    args = [f"{k}={round(p[k], 3) if isinstance(p[k], float) else p[k]!r}" for k in keys if k in p]
+    lines, line = [], ""
+    for a in args:  # wrap the config arguments, a few per line
+        if line and len(line) + len(a) > 62:
+            lines.append(line.rstrip())
+            line = ""
+        line += a + ", "
+    lines.append(line.rstrip(", "))
+    cfg = "\n        " + "\n        ".join(lines) + ",\n    "
+    gen = f',\n    image_generator="{spec.image_generator}",' if spec.image_generator else ","
+    q = question.strip().replace('"""', "'''") or "What is shown in this image?"
+    steps = f"\n    image_steps={p['image_steps']}," if p.get("image_steps") else ""
+    return f'''# pip install simit
+from simit import SIMIT, SIMITConfig
+
+model = SIMIT.from_pretrained(
+    "{spec.repo}"{gen}
+    config=SIMITConfig({cfg}),{steps}
+)
+
+image = "my_image.jpg"
+question = """{q}"""
+
+greedy = model.greedy(image, question)                          # base model, zero-shot
+demos = model.imagine(image, question{f", k={p['k_max']}" if force else ""}, time_limit={max(5, budget - 15)})   # imagined examples
+improved = model.answer(image, question, demos)                 # SIMIT-ICL
+print(greedy, "->", improved)
+
+from simit import save_demos
+save_demos(demos, "imagined/")                                  # PNGs + jsonl, e.g. to fine-tune on
+'''

examples/index.json

@@ -0,0 +1,274 @@
+[
+ {
+  "id": "bagel_ok_vqa_val2014_492_always",
+  "query_file": "query.jpg",
+  "model": "bagel",
+  "budget": 60,
+  "always": true,
+  "question": "In what country would you find this hat?\nWhen the provided information is insufficient, respond with 'Unanswerable'.\nAnswer the question using a single word or phrase.",
+  "reference": "vietnam",
+  "greedy": "Unanswerable",
+  "p0": 0.8883122759482616,
+  "simit": "Vietnam",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "How many wheels does the bicycle have?",
+    "answer": "Two",
+    "skill": "natural",
+    "verify_score": 100,
+    "confidence": null
+   },
+   {
+    "image": "demo_1.png",
+    "question": "What color is the hat the person is wearing?",
+    "answer": "Brown",
+    "skill": "natural",
+    "verify_score": 70,
+    "confidence": null
+   }
+  ],
+  "elapsed": 16.20292091369629,
+  "source": "OK-VQA row 492, paper Fig. (qualitative samples)"
+ },
+ {
+  "id": "bagel_infovqa_val_449",
+  "query_file": "query.jpg",
+  "model": "bagel",
+  "budget": 60,
+  "always": false,
+  "question": "How many points are listed under when to wash your hands?\nAnswer the question using a single word or phrase.",
+  "reference": "6",
+  "greedy": "7",
+  "p0": 0.6461464658834921,
+  "simit": "6",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "How many steps are there in the handwashing process?",
+    "answer": "11",
+    "skill": "figure",
+    "verify_score": 70,
+    "confidence": 0.38085715656328817
+   }
+  ],
+  "elapsed": 55.53464651107788,
+  "source": "InfographicVQA row 449, paper Fig. (qualitative samples)"
+ },
+ {
+  "id": "bagel_amazon",
+  "query_file": "query.png",
+  "model": "bagel",
+  "budget": 60,
+  "always": false,
+  "question": "Does this contain any trans fat?",
+  "reference": "No",
+  "greedy": "Yes, the bagel contains trans fat. According to the nutrition facts label on the packaging, the bagel has 0.5 grams of trans fat per serving.",
+  "p0": 0.7576221457584613,
+  "simit": "No",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "How many calories are in one serving of the Turkish Sesame Bagel?",
+    "answer": "290",
+    "skill": "natural",
+    "verify_score": 75,
+    "confidence": 0.3624771327791413
+   },
+   {
+    "image": "demo_1.png",
+    "question": "How many calories are listed for each serving of Turkish Sessame Bagel?",
+    "answer": "290",
+    "skill": "natural",
+    "verify_score": 100,
+    "confidence": 0.147100984163098
+   }
+  ],
+  "elapsed": 27.87025022506714,
+  "source": "Amazon product page (screenshot)"
+ },
+ {
+  "id": "bagel_vizwiz_vqa_val_345_always",
+  "query_file": "query.jpg",
+  "model": "bagel",
+  "budget": 60,
+  "always": true,
+  "question": "What color is this cat?\nWhen the provided information is insufficient, respond with 'Unanswerable'.\nAnswer the question using a single word or phrase.",
+  "reference": "brown",
+  "greedy": "Unanswerable",
+  "p0": 0.6977114977650781,
+  "simit": "Brown",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "How many cats are in the image?",
+    "answer": "One",
+    "skill": "natural",
+    "verify_score": 70,
+    "confidence": null
+   },
+   {
+    "image": "demo_1.png",
+    "question": "What type of furniture is visible in the background?",
+    "answer": "A chair",
+    "skill": "natural",
+    "verify_score": 70,
+    "confidence": null
+   }
+  ],
+  "elapsed": 10.399970293045044,
+  "source": "VizWiz row 345, paper teaser"
+ },
+ {
+  "id": "bagel_vizwiz_vqa_val_47_always",
+  "query_file": "query.jpg",
+  "model": "bagel",
+  "budget": 60,
+  "always": true,
+  "question": "What is this item?\nWhen the provided information is insufficient, respond with 'Unanswerable'.\nAnswer the question using a single word or phrase.",
+  "reference": "angel",
+  "greedy": "Unanswerable.",
+  "p0": 0.8626286828795229,
+  "simit": "Angel",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "How many cats are depicted in the image?",
+    "answer": "One",
+    "skill": "natural",
+    "verify_score": 80,
+    "confidence": null
+   },
+   {
+    "image": "demo_1.png",
+    "question": "How many cats are featured in the image?",
+    "answer": "One",
+    "skill": "natural",
+    "verify_score": 70,
+    "confidence": null
+   }
+  ],
+  "elapsed": 15.055564641952515,
+  "source": "VizWiz row 47, paper appendix"
+ },
+ {
+  "id": "bagel_vizwiz_vqa_val_87_always",
+  "query_file": "query.jpg",
+  "model": "bagel",
+  "budget": 60,
+  "always": true,
+  "question": "What is this a picture of?\nWhen the provided information is insufficient, respond with 'Unanswerable'.\nAnswer the question using a single word or phrase.",
+  "reference": "dog in lace dress",
+  "greedy": "Unanswerable.",
+  "p0": 0.8959448319074929,
+  "simit": "Dog",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "How many legs does the dog have?",
+    "answer": "Four",
+    "skill": "natural",
+    "verify_score": 100,
+    "confidence": null
+   },
+   {
+    "image": "demo_1.png",
+    "question": "What color is the dog's belly?",
+    "answer": "Pink",
+    "skill": "natural",
+    "verify_score": 100,
+    "confidence": null
+   }
+  ],
+  "elapsed": 9.160282850265503,
+  "source": "VizWiz row 87, paper appendix"
+ },
+ {
+  "id": "lance_ok_vqa_val2014_253",
+  "query_file": "query.jpg",
+  "model": "lance",
+  "budget": 60,
+  "always": false,
+  "question": "What time period is the action pictured here based off of?\nAnswer the question using a single word or phrase.",
+  "reference": "medieval",
+  "greedy": "Polo",
+  "p0": 0.6023501950034147,
+  "simit": "Medieval",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "What color does the umbrella look like?",
+    "answer": "Red",
+    "skill": "natural",
+    "verify_score": null,
+    "confidence": 0.36976529385784346
+   }
+  ],
+  "elapsed": 16.72637391090393,
+  "source": "OK-VQA row 253, paper appendix"
+ },
+ {
+  "id": "lance_ok_vqa_val2014_8_always",
+  "query_file": "query.jpg",
+  "model": "lance",
+  "budget": 60,
+  "always": true,
+  "question": "What kind of house is next to the stop sign?\nAnswer the question using a single word or phrase.",
+  "reference": "apartment",
+  "greedy": "Single story",
+  "p0": 0.5895590491401231,
+  "simit": "Apartment",
+  "greedy_ok": false,
+  "simit_ok": true,
+  "demos": [
+   {
+    "image": "demo_0.png",
+    "question": "What are the colors of the umbrellas in the painting?",
+    "answer": "No",
+    "skill": "natural",
+    "verify_score": null,
+    "confidence": null
+   },
+   {
+    "image": "demo_1.png",
+    "question": "What color does the door have?",
+    "answer": "Brown",
+    "skill": "natural",
+    "verify_score": null,
+    "confidence": null
+   },
+   {
+    "image": "demo_2.png",
+    "question": "What color is the umbrella?",
+    "answer": "No",
+    "skill": "natural",
+    "verify_score": null,
+    "confidence": null
+   },
+   {
+    "image": "demo_3.png",
+    "question": "What is the color of the stop sign?",
+    "answer": "Red",
+    "skill": "natural",
+    "verify_score": null,
+    "confidence": null
+   }
+  ],
+  "elapsed": 7.836957216262817,
+  "source": "OK-VQA row 8, paper appendix"
+ }
+]

packages.txt

@@ -0,0 +1,18 @@
+libnss3
+libnspr4
+libatk1.0-0
+libatk-bridge2.0-0
+libcups2
+libdrm2
+libxkbcommon0
+libxcomposite1
+libxdamage1
+libxfixes3
+libxrandr2
+libgbm1
+libasound2
+libpango-1.0-0
+libcairo2
+libxrender1
+libxext6
+fonts-dejavu-core

presets.json

@@ -0,0 +1,233 @@
+{
+  "_note": "built by tools/build_presets.py from tune_presets.py + latency.py runs",
+  "bagel": {
+    "20": {
+      "epsilon": 1.0,
+      "A0": 0.5631057319035666,
+      "A1": 0.6871778307441961,
+      "B": 0.3818684686169197,
+      "t_low": 0.01568469261109281,
+      "t_high": 0.9941888593223634,
+      "k_max": 1,
+      "attempts_per_slot": 2,
+      "repair_retries": 0,
+      "verify_rounds": 1,
+      "imagine_by_default": false,
+      "verify": false,
+      "image_steps": 25
+    },
+    "40": {
+      "epsilon": 1.0,
+      "A0": 0.5631057319035666,
+      "A1": 0.6871778307441961,
+      "B": 0.3818684686169197,
+      "t_low": 0.01568469261109281,
+      "t_high": 0.9941888593223634,
+      "k_max": 1,
+      "attempts_per_slot": 2,
+      "repair_retries": 1,
+      "verify_rounds": 1,
+      "imagine_by_default": false,
+      "verify": true,
+      "verify_think": true,
+      "image_steps": 50
+    },
+    "60": {
+      "epsilon": 0.13150801401314569,
+      "A0": 0.5697034469957714,
+      "A1": 0.6928285171478893,
+      "B": 0.4731007761788684,
+      "t_low": 0.012450289537920928,
+      "t_high": 0.702696037286609,
+      "k_max": 2,
+      "attempts_per_slot": 3,
+      "repair_retries": 1,
+      "verify_rounds": 2,
+      "imagine_by_default": true,
+      "verify": true,
+      "verify_think": true,
+      "image_steps": 50
+    },
+    "90": {
+      "epsilon": 0.13150801401314569,
+      "A0": 0.5697034469957714,
+      "A1": 0.6928285171478893,
+      "B": 0.4731007761788684,
+      "t_low": 0.012450289537920928,
+      "t_high": 0.702696037286609,
+      "k_max": 2,
+      "attempts_per_slot": 4,
+      "repair_retries": 2,
+      "verify_rounds": 3,
+      "imagine_by_default": true,
+      "verify": true,
+      "verify_think": true,
+      "image_steps": 50
+    },
+    "180": {
+      "epsilon": 0.13150801401314569,
+      "A0": 0.5697034469957714,
+      "A1": 0.6928285171478893,
+      "B": 0.4731007761788684,
+      "t_low": 0.012450289537920928,
+      "t_high": 0.702696037286609,
+      "k_max": 2,
+      "attempts_per_slot": 4,
+      "repair_retries": 2,
+      "verify_rounds": 5,
+      "imagine_by_default": true,
+      "verify": true,
+      "verify_think": true,
+      "image_steps": 50
+    }
+  },
+  "lance": {
+    "20": {
+      "epsilon": 0.10540046109011536,
+      "A0": 0.5284418351778891,
+      "A1": 0.7185259059153423,
+      "B": 0.20052864010313778,
+      "t_low": 0.24157862499461874,
+      "t_high": 0.8741991398657891,
+      "k_max": 1,
+      "attempts_per_slot": 2,
+      "repair_retries": 0,
+      "verify_rounds": 1,
+      "imagine_by_default": true,
+      "image_steps": 30
+    },
+    "40": {
+      "epsilon": 0.11989932190804062,
+      "A0": 0.35382385806255806,
+      "A1": 0.4638199725338076,
+      "B": 0.2584592356693373,
+      "t_low": 0.25642426481228886,
+      "t_high": 0.9450464322610527,
+      "k_max": 4,
+      "attempts_per_slot": 2,
+      "repair_retries": 1,
+      "verify_rounds": 1,
+      "imagine_by_default": true,
+      "image_steps": 30
+    },
+    "60": {
+      "epsilon": 0.11989932190804062,
+      "A0": 0.35382385806255806,
+      "A1": 0.4638199725338076,
+      "B": 0.2584592356693373,
+      "t_low": 0.25642426481228886,
+      "t_high": 0.9450464322610527,
+      "k_max": 4,
+      "attempts_per_slot": 3,
+      "repair_retries": 1,
+      "verify_rounds": 2,
+      "imagine_by_default": true,
+      "image_steps": 30
+    },
+    "90": {
+      "epsilon": 0.11989932190804062,
+      "A0": 0.35382385806255806,
+      "A1": 0.4638199725338076,
+      "B": 0.2584592356693373,
+      "t_low": 0.25642426481228886,
+      "t_high": 0.9450464322610527,
+      "k_max": 4,
+      "attempts_per_slot": 4,
+      "repair_retries": 2,
+      "verify_rounds": 3,
+      "imagine_by_default": true,
+      "image_steps": 30
+    },
+    "180": {
+      "epsilon": 0.11989932190804062,
+      "A0": 0.35382385806255806,
+      "A1": 0.4638199725338076,
+      "B": 0.2584592356693373,
+      "t_low": 0.25642426481228886,
+      "t_high": 0.9450464322610527,
+      "k_max": 4,
+      "attempts_per_slot": 4,
+      "repair_retries": 2,
+      "verify_rounds": 5,
+      "imagine_by_default": true,
+      "image_steps": 30
+    }
+  },
+  "qwen": {
+    "20": {
+      "epsilon": 1.0,
+      "A0": 0.29123083240427616,
+      "A1": 0.892922351072982,
+      "B": 0.051022405374014035,
+      "t_low": 0.10443837804741735,
+      "t_high": 0.24890099152089373,
+      "k_max": 1,
+      "attempts_per_slot": 2,
+      "repair_retries": 0,
+      "verify_rounds": 1,
+      "imagine_by_default": false,
+      "verify": false,
+      "use_skills": false
+    },
+    "40": {
+      "epsilon": 1.0,
+      "A0": 0.29123083240427616,
+      "A1": 0.892922351072982,
+      "B": 0.051022405374014035,
+      "t_low": 0.10443837804741735,
+      "t_high": 0.24890099152089373,
+      "k_max": 1,
+      "attempts_per_slot": 2,
+      "repair_retries": 1,
+      "verify_rounds": 1,
+      "imagine_by_default": false,
+      "verify": false,
+      "use_skills": false
+    },
+    "60": {
+      "epsilon": 1.0,
+      "A0": 0.4895372103191979,
+      "A1": 0.6789573167296538,
+      "B": 0.39026458814322773,
+      "t_low": 0.05913721293446661,
+      "t_high": 0.6612150885625033,
+      "k_max": 2,
+      "attempts_per_slot": 3,
+      "repair_retries": 1,
+      "verify_rounds": 2,
+      "imagine_by_default": false,
+      "verify": false,
+      "use_skills": true
+    },
+    "90": {
+      "epsilon": 1.0,
+      "A0": 0.4895372103191979,
+      "A1": 0.6789573167296538,
+      "B": 0.39026458814322773,
+      "t_low": 0.05913721293446661,
+      "t_high": 0.6612150885625033,
+      "k_max": 2,
+      "attempts_per_slot": 4,
+      "repair_retries": 2,
+      "verify_rounds": 3,
+      "imagine_by_default": false,
+      "verify": false,
+      "use_skills": true
+    },
+    "180": {
+      "epsilon": 1.0,
+      "A0": 0.29123083240427616,
+      "A1": 0.892922351072982,
+      "B": 0.051022405374014035,
+      "t_low": 0.10443837804741735,
+      "t_high": 0.24890099152089373,
+      "k_max": 1,
+      "attempts_per_slot": 4,
+      "repair_retries": 2,
+      "verify_rounds": 5,
+      "imagine_by_default": false,
+      "verify": true,
+      "use_skills": true
+    }
+  }
+}

requirements.txt

@@ -0,0 +1,4 @@
+# gradio, spaces and huggingface_hub are preinstalled and managed by the Spaces platform: don't list them.
+simit>=0.1.1
+# simit needs torch >= 2.12 (torch.nn.attention.varlen); 2.13.0 is on ZeroGPU's supported list
+torch==2.13.0

tools/build_presets.py

@@ -0,0 +1,91 @@
+"""Combine accuracy (tune_presets.py) and latency (latency.py) into presets.json:
+for every model and time budget, the largest K_max whose demos arrive in time
+for most requests, at the best-quality speed setting that achieves it, with the
+ABA/DF thresholds tuned for that K_max.
+
+python tools/build_presets.py RUN_DIR
+"""
+import json
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).parent.parent
+sys.path.insert(0, str(HERE))
+from demo_models import ANSWER_RESERVE, BUDGETS, MODELS  # noqa: E402
+
+run = Path(sys.argv[1])
+# ZeroGPU vs. our H100 (spaces' own duration_factor: 1.5 on the half GPU, none on the full one)
+SLOWDOWN = {"large": 1.5, "xlarge": 1.0}
+QUALITY = {  # speed settings, best quality first (as timed by latency.py)
+    "bagel": ["think50", "nothink50", "nothink25", "nocritic25"],
+    "lance": ["steps30", "steps20"],
+    "qwen": ["critic", "skills", "natural"],
+}
+SETTING = {
+    "think50": dict(verify=True, verify_think=True, image_steps=50),
+    "nothink50": dict(verify=True, verify_think=False, image_steps=50),
+    "nothink25": dict(verify=True, verify_think=False, image_steps=25),
+    "nocritic25": dict(verify=False, image_steps=25),
+    "steps30": dict(image_steps=30), "steps20": dict(image_steps=20),
+    "critic": dict(verify=True, use_skills=True), "skills": dict(verify=False, use_skills=True),
+    "natural": dict(verify=False, use_skills=False),
+}
+SHARE = 0.8            # fraction of requests that must get all K demos in time
+
+out_file = HERE / "presets.json"
+presets = json.loads(out_file.read_text()) if out_file.exists() else {}   # models without data keep theirs
+presets["_note"] = "built by tools/build_presets.py from tune_presets.py + latency.py runs"
+report = []
+for key in MODELS:
+    tune_file, lat_file = run / f"presets_{key}.json", run / f"latency_{key}.json"
+    if not (tune_file.exists() and lat_file.exists()):
+        print(f"skip {key}: missing {tune_file.name if not tune_file.exists() else lat_file.name}")
+        continue
+    tuning = json.loads(tune_file.read_text())
+    tuned, zero_shot = tuning["robust_per_k_max"], tuning["zero_shot"]
+    timelines = json.loads(lat_file.read_text())
+    f = SLOWDOWN[MODELS[key].gpu_size]
+    scores = {int(k): v["score"] for k, v in tuned.items()}
+    presets[key] = {}
+    n_val = tuning["n_val"]
+    band = {int(k): v["params"]["t_high"] - v["params"]["t_low"] for k, v in tuned.items()}
+
+    def fits(setting, k, deadline):
+        """Most requests have k demos before the deadline. Pooled over every run that asked for at least
+        k demos (when its k-th demo arrived): 6 requests per cell alone are too noisy."""
+        runs = [r for r in timelines if r["setting"] == setting and r["k"] >= k]
+        if not runs:
+            return False
+        return sum(len(r.get("demos", [])) >= k and r["demos"][k - 1] <= deadline for r in runs) >= SHARE * len(runs)
+
+    for budget in BUDGETS:
+        deadline = (budget - ANSWER_RESERVE) / f      # the demo's imagination deadline, in H100 seconds
+        choice = None
+        # The tuned accuracies hold for the setting the tuning cache was made with (the first, best-quality
+        # one), so use it whenever any K fits; faster settings only when it cannot fit at all.
+        for setting in QUALITY[key]:
+            ks = [k for k in sorted(scores) if fits(setting, k, deadline)]
+            if not ks:
+                continue
+            top = max(scores[k] for k in ks)
+            near = [k for k in ks if scores[k] >= top - 1.0 / n_val]   # within one validation query
+            k = min(near, key=lambda k: (-band[k], k))                  # widest band, then fewest demos
+            choice = (k, setting)
+            break
+        if choice is None:  # nothing fits: the fastest setting, one demo; the deadline decides
+            choice = (1, QUALITY[key][-1])
+        k, setting = choice
+        # more time: more attempts and checking rounds (the paper's 4 / 2 / 5 at the largest budget)
+        effort = {20: (2, 0, 1), 40: (2, 1, 1), 60: (3, 1, 2), 90: (4, 2, 3), 180: (4, 2, 5)}[budget]
+        params = {p: v for p, v in tuned[str(k)]["params"].items() if p != "k_max"}
+        helps = scores[k] > zero_shot + 0.5 / n_val
+        if not helps:  # imagination does not beat zero-shot here: ABA never imagines (the UI's checkbox still can)
+            params["epsilon"] = 1.0
+        presets[key][str(budget)] = dict(params, k_max=k, attempts_per_slot=effort[0], repair_retries=effort[1],
+                                         verify_rounds=effort[2], imagine_by_default=helps, **SETTING[setting])
+        report.append((key, budget, k, setting, round(scores[k], 3), round(zero_shot, 3), round(band[k], 2),
+                       "" if helps else "  -> no gain: zero-shot by default"))
+
+out_file.write_text(json.dumps(presets, indent=2))
+for row in report:
+    print("%-6s %4ss  K_max=%d  %-11s val %.3f (zero-shot %.3f)  DF band %.2f%s" % row)

tools/fork_test.py

@@ -0,0 +1,55 @@
+"""Simulate ZeroGPU locally: load weights on CPU in the parent (CUDA must stay
+uninitialized there), then run whole requests in forked children, which is
+what ``@spaces.GPU`` does on ZeroGPU.
+
+CUDA_VISIBLE_DEVICES=0 python tools/fork_test.py bagel [n_requests]
+"""
+import multiprocessing as mp
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+sys.path.insert(0, str(Path(__file__).parent))
+
+# ZeroGPU's own main-process emulation (fake CUDA queries, CUDA init forbidden); children undo it,
+# as spaces' worker does before attaching the GPU.
+from spaces.zero.torch import patching  # noqa: E402
+import torch  # noqa: E402
+
+patching.patch()
+from demo_models import MODELS, run  # noqa: E402
+from general_qa import general_qa  # noqa: E402
+
+key = sys.argv[1]
+n = int(sys.argv[2]) if len(sys.argv) > 2 else 2
+spec = MODELS[key]
+t = time.time()
+spec.load()
+print(f"parent: loaded on CPU in {time.time() - t:.0f}s; CUDA initialized in parent: {torch.cuda.is_initialized()}",
+      flush=True)
+assert not torch.cuda.is_initialized(), "the main process must not initialize CUDA (ZeroGPU forks it)"
+queries = general_qa(per_subset=1, offset=420)[:n]
+
+
+def child(i, q):
+    patching.unpatch()
+    t0 = time.time()
+    for ev in run(spec, q["image"], q["question"], 60, False):
+        if ev[0] == "greedy":
+            print(f"  [{i}] greedy={ev[1].answer!r} p0={ev[1].confidence:.2f} k={ev[2]} at {time.time() - t0:.1f}s",
+                  flush=True)
+        elif ev[0] == "demo":
+            print(f"  [{i}] demo [{ev[1].skill}] {ev[1].question!r} -> {ev[1].answer!r} at {time.time() - t0:.1f}s",
+                  flush=True)
+        elif ev[0] == "final":
+            print(f"  [{i}] final={ev[1]!r} {ev[2]} refs={q['answers'][:3]} total {time.time() - t0:.1f}s", flush=True)
+
+
+ctx = mp.get_context("fork")
+for i, q in enumerate(queries):
+    p = ctx.Process(target=child, args=(i, q))
+    t0 = time.time()
+    p.start()
+    p.join()
+    print(f"request {i} ({q['subset']}): exit={p.exitcode} in {time.time() - t0:.1f}s", flush=True)

tools/general_qa.py

@@ -0,0 +1,70 @@
+"""A mixed "general visual QA" set from LMMs-Eval-Lite, close to what demo
+visitors ask: everyday photos (VizWiz, VQAv2), knowledge (OK-VQA), reading text
+(TextVQA), charts and documents (ChartQA, DocVQA) and multiple-choice reasoning
+(MMBench, AI2D). Each example carries its benchmark's own metric.
+"""
+
+import ast
+
+from datasets import load_dataset
+
+VQA = "Answer the question using a single word or phrase."
+
+
+def _list(v):
+    if isinstance(v, list):
+        return v
+    try:
+        v = ast.literal_eval(v)
+    except Exception:
+        pass
+    return v if isinstance(v, list) else [v]
+
+
+def _options(row, keys):
+    opts = [(k, str(row[k])) for k in keys if str(row.get(k, "nan")) not in ("nan", "None", "")]
+    return "\n".join(f"{k}. {v}" for k, v in opts)
+
+
+def _mmbench(r):
+    hint = str(r.get("hint", "nan"))
+    q = (hint + "\n" if hint not in ("nan", "None", "") else "") + r["question"]
+    return (q + "\n" + _options(r, "ABCD") + "\nAnswer with the option's letter from the given choices directly.",
+            [str(r["answer"])], "multiple_choice")
+
+
+def _ai2d(r):
+    opts = _list(r["options"])
+    q = r["question"] + "\n" + "\n".join(f"{chr(65 + i)}. {o}" for i, o in enumerate(opts))
+    return q + "\nAnswer with the option's letter from the given choices directly.", [chr(65 + int(r["answer"]))], \
+        "multiple_choice"
+
+
+SUBSETS = {
+    "vizwiz_vqa_val": lambda r: (f"{r['question']}\nWhen the provided information is insufficient, respond with "
+                                 f"'Unanswerable'.\n{VQA}", [str(a) for a in _list(r["answers"])], "vqa_accuracy"),
+    "vqav2_val": lambda r: (f"{r['question']}\n{VQA}",
+                            [a["answer"] if isinstance(a, dict) else str(a) for a in _list(r["answers"])],
+                            "vqa_accuracy"),
+    "ok_vqa_val2014": lambda r: (f"{r['question']}\n{VQA}", [str(a) for a in _list(r["answers"])], "vqa_accuracy"),
+    "textvqa_val": lambda r: (f"{r['question']}\n{VQA}", [str(a) for a in _list(r["answers"])], "vqa_accuracy"),
+    "chartqa": lambda r: (f"{r['question']}\nAnswer the question with a single word.",
+                          [str(a) for a in _list(r["answer"])], "relaxed_accuracy"),
+    "docvqa_val": lambda r: (f"{r['question']}\n{VQA}", [str(a) for a in _list(r["answers"])], "anls"),
+    "mmbench_en_dev": _mmbench,
+    "ai2d": _ai2d,
+}
+
+
+def general_qa(per_subset: int = 8, offset: int = 300) -> list[dict]:
+    """``per_subset`` examples from each subset, starting at row ``offset``
+    (rows below 300 were used for the package benchmarks)."""
+    out = []
+    for sub, fn in SUBSETS.items():
+        ds = load_dataset("lmms-lab/LMMs-Eval-Lite", sub)["lite"]
+        for i in range(offset, offset + per_subset):
+            r = ds[i]
+            question, answers, metric = fn(r)
+            out.append({"image": r["image"].convert("RGB"), "question": question, "answers": answers,
+                        "metric": metric, "subset": sub})
+    return out

tools/latency.py

@@ -0,0 +1,74 @@
+"""Per-request timelines on the demo's own code path, as ZeroGPU runs it (a fresh
+forked process per request, weights moved to the GPU each time): when the greedy
+answer and each imagined demo arrive, for forced K and a given speed setting.
+
+CUDA_VISIBLE_DEVICES=1 python tools/latency.py lance OUT.json
+"""
+import json
+import multiprocessing as mp
+import os
+import sys
+import time
+from pathlib import Path
+
+os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")
+os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")
+sys.path.insert(0, str(Path(__file__).parent.parent))
+sys.path.insert(0, str(Path(__file__).parent))
+from spaces.zero.torch import patching  # noqa: E402
+import torch  # noqa: E402
+
+patching.patch()
+from demo_models import MODELS, run  # noqa: E402
+from general_qa import general_qa  # noqa: E402
+
+ABA = dict(epsilon=0.06, A0=0.5, A1=0.85, B=0.4, t_low=0.0, t_high=1.0)
+SETTINGS = {  # speed settings to time (accuracy comes from tune_presets.py)
+    "bagel": {"think50": dict(verify=True, verify_think=True, image_steps=50),
+              "nothink50": dict(verify=True, verify_think=False, image_steps=50),
+              "nothink25": dict(verify=True, verify_think=False, image_steps=25),
+              "nocritic25": dict(verify=False, image_steps=25)},
+    "lance": {"steps30": dict(image_steps=30), "steps20": dict(image_steps=20)},
+    "qwen": {"natural": dict(verify=False, use_skills=False), "skills": dict(verify=False, use_skills=True),
+             "critic": dict(verify=True, use_skills=True)},
+}
+KS = {"bagel": [1, 2, 4], "lance": [1, 2, 4], "qwen": [1, 2]}
+
+key, out = sys.argv[1], Path(sys.argv[2])
+n_q = int(sys.argv[3]) if len(sys.argv) > 3 else 6
+spec = MODELS[key]
+spec.load()
+assert not torch.cuda.is_initialized()
+queries = general_qa(per_subset=1, offset=440)[:n_q]
+
+
+def child(conn, q, preset):
+    patching.unpatch()
+    t0 = time.time()
+    rec = {"demos": []}
+    spec.preset = lambda budget: dict(preset)   # this process only
+    for ev in run(spec, q["image"], q["question"], 1000, True):
+        t = time.time() - t0
+        if ev[0] == "greedy":
+            rec["greedy"] = t
+        elif ev[0] == "demo":
+            rec["demos"].append(t)
+        elif ev[0] == "final":
+            rec["total"] = t
+    conn.send(rec)
+
+
+results = []
+for sname, setting in SETTINGS[key].items():
+    for k in KS[key]:
+        preset = dict(ABA, k_max=k, attempts_per_slot=4, repair_retries=1, verify_rounds=2, **setting)
+        for qi, q in enumerate(queries):
+            a, b = mp.get_context("fork").Pipe()
+            p = mp.get_context("fork").Process(target=child, args=(b, q, preset))
+            p.start()
+            rec = a.recv() if a.poll(900) else {"error": "timeout"}
+            p.join(timeout=30)
+            rec.update(setting=sname, k=k, query=qi, subset=q["subset"])
+            results.append(rec)
+            print(json.dumps(rec), flush=True)
+            out.write_text(json.dumps(results, indent=1))

tools/make_examples.py

@@ -0,0 +1,121 @@
+"""Cached demo examples: real runs of the demo pipeline on the paper's figure
+queries (LMMs-Eval-Lite rows), stored with both answers, the imagined demos and
+the reference answer. ``--keep improved`` keeps only runs where SIMIT-ICL fixed
+the base answer.
+
+CUDA_VISIBLE_DEVICES=0 python tools/make_examples.py bagel 60 [--keep improved]
+"""
+import argparse
+import json
+import os
+import shutil
+import sys
+import time
+from pathlib import Path
+
+os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")
+os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")
+HERE = Path(__file__).parent.parent
+sys.path.insert(0, str(HERE))
+sys.path.insert(0, str(HERE / "tools"))
+
+from datasets import load_dataset  # noqa: E402
+from PIL import Image  # noqa: E402
+
+from demo_models import MODELS, run  # noqa: E402
+from general_qa import _list  # noqa: E402
+from simit.metrics import get_metric  # noqa: E402
+
+VQA = "Answer the question using a single word or phrase."
+UNANSWERABLE = "When the provided information is insufficient, respond with 'Unanswerable'."
+# (subset, row, prompt suffix, metric, where it appears in the paper)
+CANDIDATES = [
+    ("ok_vqa_val2014", 492, f"{UNANSWERABLE}\n{VQA}", "vqa_accuracy", "paper Fig. (qualitative samples)"),
+    ("infovqa_val", 449, VQA, "anls", "paper Fig. (qualitative samples)"),
+    ("vizwiz_vqa_val", 345, f"{UNANSWERABLE}\n{VQA}", "vqa_accuracy", "paper teaser"),
+    ("vizwiz_vqa_val", 47, f"{UNANSWERABLE}\n{VQA}", "vqa_accuracy", "paper appendix"),
+    ("vizwiz_vqa_val", 87, f"{UNANSWERABLE}\n{VQA}", "vqa_accuracy", "paper appendix"),
+] + [("ok_vqa_val2014", i, VQA, "vqa_accuracy", "paper appendix") for i in (115, 151, 210, 22, 253, 272, 291, 381,
+                                                                           477, 8)]
+SUBSET_NAME = {"ok_vqa_val2014": "OK-VQA", "infovqa_val": "InfographicVQA", "vizwiz_vqa_val": "VizWiz"}
+
+ap = argparse.ArgumentParser()
+ap.add_argument("model")
+ap.add_argument("budget", type=int)
+ap.add_argument("--keep", choices=["all", "improved"], default="all")
+ap.add_argument("--always", action="store_true", help="imagine even when the model is confident (UI checkbox)")
+ap.add_argument("--out", default=str(HERE / "examples"))
+ap.add_argument("--only", nargs="*", default=None, help="subset:row items to run, e.g. ok_vqa_val2014:492")
+ap.add_argument("--image", help="a custom query image (instead of the dataset candidates)")
+ap.add_argument("--question", help="the question for --image")
+ap.add_argument("--reference", nargs="+", help="reference answer(s) for --image")
+ap.add_argument("--metric", default="contains", help="metric for --image (free-form answers: 'contains')")
+ap.add_argument("--id", help="example id (folder name) for --image")
+ap.add_argument("--source", default="custom example", help="source line shown in the UI for --image")
+args = ap.parse_args()
+out = Path(args.out)
+out.mkdir(parents=True, exist_ok=True)
+index_file = out / "index.json"
+index = json.loads(index_file.read_text()) if index_file.exists() else []
+spec = MODELS[args.model]
+spec.load()
+suffix_id = "_always" if args.always else ""
+
+
+def jobs():
+    """(example id, image, question, references, metric name, source, keep the folder's own files)"""
+    if args.image:
+        img = Image.open(args.image)
+        img.load()
+        yield (args.id + suffix_id, img, args.question, args.reference, args.metric, args.source, True)
+        return
+    for sub, row, suffix, metric_name, where in CANDIDATES:
+        if args.only is not None and f"{sub}:{row}" not in args.only:
+            continue
+        r = load_dataset("lmms-lab/LMMs-Eval-Lite", sub)["lite"][row]
+        yield (f"{args.model}_{sub}_{row}{suffix_id}", r["image"].convert("RGB"), f"{r['question']}\n{suffix}",
+               [str(a) for a in _list(r.get("answers", r.get("answer")))], metric_name,
+               f"{SUBSET_NAME.get(sub, sub)} row {row}, {where}", False)
+
+
+for ex_id, image, question, refs, metric_name, source, custom in jobs():
+    metric = get_metric(metric_name)
+    t0 = time.time()
+    greedy, demos, final, info = None, [], None, {}
+    for ev in run(spec, image, question, args.budget, args.always):
+        if ev[0] == "greedy":
+            greedy = ev[1]
+        elif ev[0] == "demo":
+            demos.append(ev[1])
+        elif ev[0] == "final":
+            final, info = ev[1], ev[2]
+    g_ok, s_ok = metric(greedy.answer, refs) >= 0.5, metric(final, refs) >= 0.5
+    print(f"{ex_id}: greedy={greedy.answer!r} ({g_ok}) simit={final!r} ({s_ok}) demos={len(demos)} "
+          f"{time.time() - t0:.0f}s refs={refs[:3]}", flush=True)
+    if args.keep == "improved" and not (s_ok and not g_ok):
+        continue
+    d = out / ex_id
+    if custom and d.exists():          # the user's own files stay; only earlier demo images are replaced
+        for f in d.glob("demo_*.png"):
+            f.unlink()
+    else:
+        shutil.rmtree(d, ignore_errors=True)
+        d.mkdir()
+    if custom and Path(args.image).resolve().parent == d.resolve():
+        query_file = Path(args.image).name
+    else:
+        shown = image.convert("RGB")
+        shown.thumbnail((1024, 1024))      # display copy (the run used the original)
+        shown.save(d / "query.jpg", quality=90)
+        query_file = "query.jpg"
+    stored = []
+    for i, dm in enumerate(demos):
+        dm.image.save(d / f"demo_{i}.png")
+        stored.append({"image": f"demo_{i}.png", "question": dm.question, "answer": dm.answer, "skill": dm.skill,
+                       "verify_score": dm.verify_score, "confidence": dm.confidence})
+    entry = {"id": ex_id, "query_file": query_file, "model": args.model, "budget": args.budget,
+             "always": args.always, "question": question, "reference": refs[0],
+             "greedy": greedy.answer, "p0": greedy.confidence, "simit": final, "greedy_ok": g_ok, "simit_ok": s_ok,
+             "demos": stored, "elapsed": info.get("elapsed", time.time() - t0), "source": source}
+    index = [e for e in index if e["id"] != ex_id] + [entry]
+    index_file.write_text(json.dumps(index, indent=1))

tools/results/latency_bagel.json

@@ -0,0 +1,882 @@
+[
+ {
+  "demos": [
+   26.002389907836914
+  ],
+  "greedy": 6.846283435821533,
+  "total": 26.74322533607483,
+  "setting": "think50",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   12.16827392578125
+  ],
+  "greedy": 6.742192268371582,
+  "total": 12.760454177856445,
+  "setting": "think50",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   17.263279914855957
+  ],
+  "greedy": 6.47596001625061,
+  "total": 17.900477647781372,
+  "setting": "think50",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   33.245251417160034
+  ],
+  "greedy": 6.677622079849243,
+  "total": 34.04009461402893,
+  "setting": "think50",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   21.499918460845947
+  ],
+  "greedy": 6.547531366348267,
+  "total": 22.275030374526978,
+  "setting": "think50",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [],
+  "greedy": 6.649888515472412,
+  "total": 57.07210087776184,
+  "setting": "think50",
+  "k": 1,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   14.020853996276855,
+   15.747642755508423
+  ],
+  "greedy": 6.6902241706848145,
+  "total": 16.58476734161377,
+  "setting": "think50",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   14.245272397994995,
+   15.97743272781372
+  ],
+  "greedy": 6.408905506134033,
+  "total": 16.649274826049805,
+  "setting": "think50",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   14.016824960708618,
+   27.909880876541138
+  ],
+  "greedy": 6.539390563964844,
+  "total": 28.65375280380249,
+  "setting": "think50",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   19.595810413360596,
+   20.93066096305847
+  ],
+  "greedy": 6.817545413970947,
+  "total": 21.97844171524048,
+  "setting": "think50",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   21.255191326141357,
+   23.45720887184143
+  ],
+  "greedy": 6.85756778717041,
+  "total": 24.484545946121216,
+  "setting": "think50",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   17.32177734375,
+   29.950555562973022
+  ],
+  "greedy": 6.608810186386108,
+  "total": 30.91857123374939,
+  "setting": "think50",
+  "k": 2,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   19.331388235092163,
+   20.449653387069702,
+   20.78369402885437,
+   24.906179904937744
+  ],
+  "greedy": 6.746614217758179,
+  "total": 25.91768527030945,
+  "setting": "think50",
+  "k": 4,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   19.866760969161987,
+   21.37235426902771,
+   25.51891589164734,
+   25.743648767471313
+  ],
+  "greedy": 6.74257230758667,
+  "total": 26.58999800682068,
+  "setting": "think50",
+  "k": 4,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   18.303645849227905,
+   19.48818588256836,
+   19.763233184814453,
+   21.081400394439697
+  ],
+  "greedy": 6.863516330718994,
+  "total": 21.978867053985596,
+  "setting": "think50",
+  "k": 4,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   18.599497079849243,
+   47.75624179840088,
+   48.84687924385071
+  ],
+  "greedy": 6.99056339263916,
+  "total": 123.67395544052124,
+  "setting": "think50",
+  "k": 4,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   16.050732851028442,
+   17.049232006072998,
+   17.45568537712097,
+   17.470638751983643
+  ],
+  "greedy": 6.753719091415405,
+  "total": 18.682735443115234,
+  "setting": "think50",
+  "k": 4,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   17.657182216644287,
+   17.93081760406494,
+   28.86550760269165,
+   48.43350315093994
+  ],
+  "greedy": 6.711917877197266,
+  "total": 49.75816011428833,
+  "setting": "think50",
+  "k": 4,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   12.481452226638794
+  ],
+  "greedy": 6.7068750858306885,
+  "total": 13.248512744903564,
+  "setting": "nothink50",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   12.028658628463745
+  ],
+  "greedy": 6.6228907108306885,
+  "total": 12.620982646942139,
+  "setting": "nothink50",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   12.994443416595459
+  ],
+  "greedy": 6.687414646148682,
+  "total": 13.635976076126099,
+  "setting": "nothink50",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [],
+  "greedy": 6.697295427322388,
+  "total": 31.110349416732788,
+  "setting": "nothink50",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   26.508880138397217
+  ],
+  "greedy": 6.976494312286377,
+  "total": 27.29017186164856,
+  "setting": "nothink50",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   41.90188384056091
+  ],
+  "greedy": 6.816613435745239,
+  "total": 42.66916561126709,
+  "setting": "nothink50",
+  "k": 1,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   15.123422622680664,
+   15.220020294189453
+  ],
+  "greedy": 7.205089807510376,
+  "total": 16.069376230239868,
+  "setting": "nothink50",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   16.50620174407959,
+   17.572696685791016
+  ],
+  "greedy": 7.146985292434692,
+  "total": 18.24684238433838,
+  "setting": "nothink50",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   25.19501304626465,
+   32.75030159950256
+  ],
+  "greedy": 7.005849599838257,
+  "total": 33.467132806777954,
+  "setting": "nothink50",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   40.2707622051239,
+   67.3787031173706
+  ],
+  "greedy": 7.045415639877319,
+  "total": 68.47934651374817,
+  "setting": "nothink50",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   16.73797917366028,
+   26.571113348007202
+  ],
+  "greedy": 7.133496999740601,
+  "total": 27.498658895492554,
+  "setting": "nothink50",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   18.925209283828735,
+   45.47921824455261
+  ],
+  "greedy": 7.158755540847778,
+  "total": 46.42701768875122,
+  "setting": "nothink50",
+  "k": 2,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   18.92365074157715,
+   20.19171667098999,
+   20.651800394058228,
+   24.54624581336975
+  ],
+  "greedy": 6.934945344924927,
+  "total": 25.56127405166626,
+  "setting": "nothink50",
+  "k": 4,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   20.397387266159058,
+   21.484426259994507,
+   21.654940843582153,
+   39.95197796821594
+  ],
+  "greedy": 6.822445392608643,
+  "total": 40.79026508331299,
+  "setting": "nothink50",
+  "k": 4,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   17.846131563186646,
+   31.462411880493164,
+   31.96941876411438,
+   42.511390209198
+  ],
+  "greedy": 7.063258647918701,
+  "total": 43.39737248420715,
+  "setting": "nothink50",
+  "k": 4,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   31.60993504524231,
+   32.77745318412781,
+   46.97655916213989,
+   67.76566171646118
+  ],
+  "greedy": 7.10393762588501,
+  "total": 69.44919919967651,
+  "setting": "nothink50",
+  "k": 4,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   13.910687446594238,
+   18.546985864639282,
+   19.083553314208984,
+   19.979618787765503
+  ],
+  "greedy": 6.712128639221191,
+  "total": 21.244775533676147,
+  "setting": "nothink50",
+  "k": 4,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   27.11043953895569,
+   36.58712029457092,
+   41.17549681663513,
+   51.479180574417114
+  ],
+  "greedy": 6.7352728843688965,
+  "total": 52.88075041770935,
+  "setting": "nothink50",
+  "k": 4,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   15.192901134490967
+  ],
+  "greedy": 6.91644811630249,
+  "total": 15.934749126434326,
+  "setting": "nothink25",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   14.546730279922485
+  ],
+  "greedy": 6.616936206817627,
+  "total": 15.137541770935059,
+  "setting": "nothink25",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   10.399396657943726
+  ],
+  "greedy": 6.794123411178589,
+  "total": 11.040837049484253,
+  "setting": "nothink25",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [],
+  "greedy": 6.857384443283081,
+  "total": 36.08436679840088,
+  "setting": "nothink25",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   17.479740619659424
+  ],
+  "greedy": 6.688612699508667,
+  "total": 18.236947536468506,
+  "setting": "nothink25",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   59.99591255187988
+  ],
+  "greedy": 6.7639524936676025,
+  "total": 60.8376727104187,
+  "setting": "nothink25",
+  "k": 1,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   12.354958057403564,
+   19.66935706138611
+  ],
+  "greedy": 6.810552358627319,
+  "total": 20.500991344451904,
+  "setting": "nothink25",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   13.075764179229736,
+   13.986951351165771
+  ],
+  "greedy": 6.526533842086792,
+  "total": 14.659094333648682,
+  "setting": "nothink25",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   11.25294542312622,
+   12.738617658615112
+  ],
+  "greedy": 6.653463363647461,
+  "total": 13.45282244682312,
+  "setting": "nothink25",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   52.51059865951538
+  ],
+  "greedy": 6.916888952255249,
+  "total": 69.37151217460632,
+  "setting": "nothink25",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   16.357442378997803,
+   22.164559364318848
+  ],
+  "greedy": 6.922367811203003,
+  "total": 23.112447500228882,
+  "setting": "nothink25",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   31.607569217681885,
+   64.07539129257202
+  ],
+  "greedy": 7.111017942428589,
+  "total": 65.04016399383545,
+  "setting": "nothink25",
+  "k": 2,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   14.144917011260986,
+   14.450908422470093,
+   14.969972133636475,
+   22.031485557556152
+  ],
+  "greedy": 6.845429420471191,
+  "total": 23.0630886554718,
+  "setting": "nothink25",
+  "k": 4,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   15.669585943222046,
+   19.648891925811768,
+   20.805134534835815,
+   29.888363122940063
+  ],
+  "greedy": 6.622054576873779,
+  "total": 30.800080060958862,
+  "setting": "nothink25",
+  "k": 4,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   15.39557409286499,
+   16.348963260650635,
+   19.663447618484497,
+   20.33902072906494
+  ],
+  "greedy": 6.637497663497925,
+  "total": 21.22013282775879,
+  "setting": "nothink25",
+  "k": 4,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   15.362420558929443,
+   15.921366691589355,
+   24.226126670837402,
+   45.1386661529541
+  ],
+  "greedy": 6.875400066375732,
+  "total": 46.43324565887451,
+  "setting": "nothink25",
+  "k": 4,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   15.788510084152222,
+   16.097867727279663,
+   19.899868726730347,
+   24.75997519493103
+  ],
+  "greedy": 7.031396865844727,
+  "total": 26.027679443359375,
+  "setting": "nothink25",
+  "k": 4,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   20.376034021377563,
+   25.979355096817017,
+   39.69830060005188,
+   47.3364462852478
+  ],
+  "greedy": 7.084344863891602,
+  "total": 48.586371660232544,
+  "setting": "nothink25",
+  "k": 4,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   10.09440541267395
+  ],
+  "greedy": 7.1284401416778564,
+  "total": 10.858816146850586,
+  "setting": "nocritic25",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   9.839585781097412
+  ],
+  "greedy": 6.931880235671997,
+  "total": 10.431512594223022,
+  "setting": "nocritic25",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   9.52512812614441
+  ],
+  "greedy": 6.723513603210449,
+  "total": 10.16373586654663,
+  "setting": "nocritic25",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   10.409209489822388
+  ],
+  "greedy": 6.779156446456909,
+  "total": 11.201088905334473,
+  "setting": "nocritic25",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   13.039257526397705
+  ],
+  "greedy": 6.719432830810547,
+  "total": 13.820397138595581,
+  "setting": "nocritic25",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   13.39191722869873
+  ],
+  "greedy": 6.744472503662109,
+  "total": 13.756275415420532,
+  "setting": "nocritic25",
+  "k": 1,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   10.131142616271973,
+   10.941117525100708
+  ],
+  "greedy": 6.75618314743042,
+  "total": 11.777562379837036,
+  "setting": "nocritic25",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   9.828729629516602,
+   10.91353988647461
+  ],
+  "greedy": 6.636314868927002,
+  "total": 11.586650609970093,
+  "setting": "nocritic25",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   10.159326076507568,
+   11.352014064788818
+  ],
+  "greedy": 7.0619635581970215,
+  "total": 12.062993288040161,
+  "setting": "nocritic25",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   12.485306739807129,
+   37.578511238098145
+  ],
+  "greedy": 6.796743392944336,
+  "total": 38.57980442047119,
+  "setting": "nocritic25",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   13.873923778533936,
+   15.028538465499878
+  ],
+  "greedy": 6.755588054656982,
+  "total": 15.933726072311401,
+  "setting": "nocritic25",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   13.108296394348145,
+   32.831605195999146
+  ],
+  "greedy": 6.732863426208496,
+  "total": 33.64893388748169,
+  "setting": "nocritic25",
+  "k": 2,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   10.41283893585205,
+   12.226398944854736,
+   13.141302824020386,
+   13.47109079360962
+  ],
+  "greedy": 6.717079401016235,
+  "total": 14.473924160003662,
+  "setting": "nocritic25",
+  "k": 4,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   10.277020931243896,
+   12.27582597732544,
+   13.355355024337769,
+   14.00516128540039
+  ],
+  "greedy": 6.814017057418823,
+  "total": 14.8407723903656,
+  "setting": "nocritic25",
+  "k": 4,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   11.423258543014526,
+   13.428069591522217,
+   14.622854709625244,
+   16.18061876296997
+  ],
+  "greedy": 6.668004751205444,
+  "total": 17.06151270866394,
+  "setting": "nocritic25",
+  "k": 4,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   10.60582423210144,
+   16.0261173248291,
+   16.700007915496826,
+   33.09624981880188
+  ],
+  "greedy": 6.843786954879761,
+  "total": 34.594066858291626,
+  "setting": "nocritic25",
+  "k": 4,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   13.488800764083862,
+   17.300123929977417,
+   18.207062244415283,
+   18.684130907058716
+  ],
+  "greedy": 6.735860824584961,
+  "total": 19.977572679519653,
+  "setting": "nocritic25",
+  "k": 4,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   12.050382614135742,
+   14.343275308609009,
+   15.249944925308228,
+   15.570221662521362
+  ],
+  "greedy": 6.8102805614471436,
+  "total": 16.96952247619629,
+  "setting": "nocritic25",
+  "k": 4,
+  "query": 5,
+  "subset": "docvqa_val"
+ }
+]

tools/results/latency_lance.json

@@ -0,0 +1,446 @@
+[
+ {
+  "demos": [
+   7.727382183074951
+  ],
+  "greedy": 4.262204647064209,
+  "total": 8.35633897781372,
+  "setting": "steps30",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   5.997106552124023
+  ],
+  "greedy": 4.037919521331787,
+  "total": 6.613722324371338,
+  "setting": "steps30",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   7.009830474853516
+  ],
+  "greedy": 4.070275545120239,
+  "total": 7.65753698348999,
+  "setting": "steps30",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   14.579021215438843
+  ],
+  "greedy": 3.9722044467926025,
+  "total": 15.173923015594482,
+  "setting": "steps30",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   7.009682893753052
+  ],
+  "greedy": 4.298367500305176,
+  "total": 7.6930906772613525,
+  "setting": "steps30",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   9.488945960998535
+  ],
+  "greedy": 4.068331241607666,
+  "total": 10.129173040390015,
+  "setting": "steps30",
+  "k": 1,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   8.336636781692505,
+   10.828430652618408
+  ],
+  "greedy": 4.046021223068237,
+  "total": 11.519325733184814,
+  "setting": "steps30",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   8.213298082351685,
+   10.596574068069458
+  ],
+  "greedy": 4.0463056564331055,
+  "total": 11.278587102890015,
+  "setting": "steps30",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   8.57367467880249,
+   10.886900663375854
+  ],
+  "greedy": 4.036629676818848,
+  "total": 11.565510511398315,
+  "setting": "steps30",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   11.43087124824524,
+   13.566925525665283
+  ],
+  "greedy": 3.9821929931640625,
+  "total": 14.221354246139526,
+  "setting": "steps30",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   9.571406841278076,
+   10.061784029006958
+  ],
+  "greedy": 4.069300174713135,
+  "total": 10.802971124649048,
+  "setting": "steps30",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   8.667640209197998,
+   11.335434436798096
+  ],
+  "greedy": 4.0576910972595215,
+  "total": 12.056833982467651,
+  "setting": "steps30",
+  "k": 2,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   9.209861993789673,
+   10.122817754745483,
+   11.41738224029541,
+   11.986809253692627
+  ],
+  "greedy": 4.213792562484741,
+  "total": 13.086158275604248,
+  "setting": "steps30",
+  "k": 4,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   8.620089054107666,
+   11.2899010181427,
+   12.844002962112427,
+   13.42651104927063
+  ],
+  "greedy": 4.100450038909912,
+  "total": 14.261920928955078,
+  "setting": "steps30",
+  "k": 4,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   7.910981893539429,
+   10.294513940811157,
+   10.294527292251587,
+   11.984448909759521
+  ],
+  "greedy": 4.040694236755371,
+  "total": 12.851784229278564,
+  "setting": "steps30",
+  "k": 4,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   8.954644203186035,
+   12.185572862625122,
+   12.858001470565796,
+   15.00017762184143
+  ],
+  "greedy": 3.9677560329437256,
+  "total": 15.794921159744263,
+  "setting": "steps30",
+  "k": 4,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   7.161853790283203,
+   10.45381784439087,
+   11.22620153427124,
+   13.835843801498413
+  ],
+  "greedy": 4.065622329711914,
+  "total": 14.741939306259155,
+  "setting": "steps30",
+  "k": 4,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   9.17311716079712,
+   9.261759042739868,
+   9.878037929534912,
+   10.475617170333862
+  ],
+  "greedy": 4.253077268600464,
+  "total": 11.615146398544312,
+  "setting": "steps30",
+  "k": 4,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   9.40204405784607
+  ],
+  "greedy": 4.1254284381866455,
+  "total": 10.027106285095215,
+  "setting": "steps20",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   7.260383367538452
+  ],
+  "greedy": 4.095474481582642,
+  "total": 7.883589267730713,
+  "setting": "steps20",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   8.731263637542725
+  ],
+  "greedy": 4.012105703353882,
+  "total": 9.378830432891846,
+  "setting": "steps20",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   11.81164813041687
+  ],
+  "greedy": 4.001217842102051,
+  "total": 12.418342351913452,
+  "setting": "steps20",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   6.208858489990234
+  ],
+  "greedy": 4.017943859100342,
+  "total": 6.826139211654663,
+  "setting": "steps20",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   9.363977670669556
+  ],
+  "greedy": 4.0784912109375,
+  "total": 10.00316596031189,
+  "setting": "steps20",
+  "k": 1,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   8.895264863967896,
+   11.027926921844482
+  ],
+  "greedy": 4.093302011489868,
+  "total": 11.739497900009155,
+  "setting": "steps20",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   8.46868896484375,
+   8.78795862197876
+  ],
+  "greedy": 4.05513858795166,
+  "total": 9.470828771591187,
+  "setting": "steps20",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   7.852817535400391,
+   8.184361457824707
+  ],
+  "greedy": 4.0655388832092285,
+  "total": 8.493193626403809,
+  "setting": "steps20",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   7.7487473487854,
+   9.187946081161499
+  ],
+  "greedy": 4.10131573677063,
+  "total": 9.855353355407715,
+  "setting": "steps20",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   8.990880966186523,
+   9.044522285461426
+  ],
+  "greedy": 4.088159084320068,
+  "total": 9.791136264801025,
+  "setting": "steps20",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   6.566630125045776,
+   8.522355079650879
+  ],
+  "greedy": 4.04767918586731,
+  "total": 9.243301391601562,
+  "setting": "steps20",
+  "k": 2,
+  "query": 5,
+  "subset": "docvqa_val"
+ },
+ {
+  "demos": [
+   8.488969802856445,
+   10.305099725723267,
+   10.305113554000854,
+   10.661272048950195
+  ],
+  "greedy": 3.986391544342041,
+  "total": 11.507165431976318,
+  "setting": "steps20",
+  "k": 4,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   8.554109811782837,
+   10.155081510543823,
+   10.1550931930542,
+   10.528407335281372
+  ],
+  "greedy": 4.004652976989746,
+  "total": 11.368907690048218,
+  "setting": "steps20",
+  "k": 4,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   7.51558256149292,
+   8.081176280975342,
+   9.272027015686035,
+   9.97708535194397
+  ],
+  "greedy": 4.020036935806274,
+  "total": 10.843841552734375,
+  "setting": "steps20",
+  "k": 4,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   7.948291540145874,
+   8.477357387542725,
+   10.59636402130127,
+   12.149571418762207
+  ],
+  "greedy": 3.9414212703704834,
+  "total": 12.953968048095703,
+  "setting": "steps20",
+  "k": 4,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   7.972762107849121,
+   9.431615591049194,
+   9.854741096496582,
+   12.329490423202515
+  ],
+  "greedy": 4.051168918609619,
+  "total": 13.170568704605103,
+  "setting": "steps20",
+  "k": 4,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   8.014616250991821,
+   9.337750434875488,
+   9.904663562774658,
+   13.779886245727539
+  ],
+  "greedy": 4.087887763977051,
+  "total": 14.651642084121704,
+  "setting": "steps20",
+  "k": 4,
+  "query": 5,
+  "subset": "docvqa_val"
+ }
+]

tools/results/latency_qwen.json

@@ -0,0 +1,347 @@
+[
+ {
+  "demos": [
+   30.282771110534668
+  ],
+  "greedy": 12.503609418869019,
+  "total": 31.559006452560425,
+  "setting": "natural",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   24.56417155265808
+  ],
+  "greedy": 12.057083129882812,
+  "total": 25.436967611312866,
+  "setting": "natural",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   27.104141235351562
+  ],
+  "greedy": 11.834922552108765,
+  "total": 28.232157707214355,
+  "setting": "natural",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   30.065296173095703
+  ],
+  "greedy": 11.83328104019165,
+  "total": 31.18956470489502,
+  "setting": "natural",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   28.98874258995056
+  ],
+  "greedy": 12.290349006652832,
+  "total": 30.4953191280365,
+  "setting": "natural",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   49.92998290061951,
+   49.93000912666321
+  ],
+  "greedy": 12.474754333496094,
+  "total": 51.23621082305908,
+  "setting": "natural",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   39.91589426994324,
+   39.91593027114868
+  ],
+  "greedy": 11.602417945861816,
+  "total": 40.8339729309082,
+  "setting": "natural",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   44.16777563095093,
+   44.16781187057495
+  ],
+  "greedy": 11.90004849433899,
+  "total": 45.30206370353699,
+  "setting": "natural",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   39.3218777179718,
+   39.32190752029419
+  ],
+  "greedy": 12.019229888916016,
+  "total": 40.492233753204346,
+  "setting": "natural",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   102.98197150230408,
+   102.9820077419281
+  ],
+  "greedy": 11.957572221755981,
+  "total": 104.37802076339722,
+  "setting": "natural",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   34.93216681480408
+  ],
+  "greedy": 12.076660633087158,
+  "total": 36.01313900947571,
+  "setting": "skills",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   26.665636777877808
+  ],
+  "greedy": 11.77315878868103,
+  "total": 27.53944206237793,
+  "setting": "skills",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   30.65152359008789
+  ],
+  "greedy": 11.954692363739014,
+  "total": 31.91823172569275,
+  "setting": "skills",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   32.95020818710327
+  ],
+  "greedy": 11.763028144836426,
+  "total": 35.805758476257324,
+  "setting": "skills",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   30.482857704162598
+  ],
+  "greedy": 12.288076162338257,
+  "total": 31.9945011138916,
+  "setting": "skills",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   47.012778997421265,
+   47.012813568115234
+  ],
+  "greedy": 11.944422960281372,
+  "total": 48.55853629112244,
+  "setting": "skills",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   43.62376880645752,
+   44.03544592857361
+  ],
+  "greedy": 11.559141635894775,
+  "total": 44.94946813583374,
+  "setting": "skills",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   44.78076720237732,
+   44.78083920478821
+  ],
+  "greedy": 12.2094886302948,
+  "total": 45.610127210617065,
+  "setting": "skills",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   39.1291127204895,
+   39.12914180755615
+  ],
+  "greedy": 11.88507342338562,
+  "total": 40.32528471946716,
+  "setting": "skills",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   90.45457100868225,
+   90.8869194984436
+  ],
+  "greedy": 12.611650466918945,
+  "total": 92.93272161483765,
+  "setting": "skills",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   75.17404270172119
+  ],
+  "greedy": 12.442910432815552,
+  "total": 76.42014288902283,
+  "setting": "critic",
+  "k": 1,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   70.33047604560852
+  ],
+  "greedy": 11.583081483840942,
+  "total": 71.14328861236572,
+  "setting": "critic",
+  "k": 1,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   71.25398755073547
+  ],
+  "greedy": 12.159323930740356,
+  "total": 72.36944890022278,
+  "setting": "critic",
+  "k": 1,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   349.18033623695374
+  ],
+  "greedy": 11.960462808609009,
+  "total": 350.2297682762146,
+  "setting": "critic",
+  "k": 1,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   79.8816294670105
+  ],
+  "greedy": 12.26366114616394,
+  "total": 81.34397840499878,
+  "setting": "critic",
+  "k": 1,
+  "query": 4,
+  "subset": "chartqa"
+ },
+ {
+  "demos": [
+   214.83966445922852,
+   246.23991990089417
+  ],
+  "greedy": 11.864954471588135,
+  "total": 247.5113034248352,
+  "setting": "critic",
+  "k": 2,
+  "query": 0,
+  "subset": "vizwiz_vqa_val"
+ },
+ {
+  "demos": [
+   111.94581985473633,
+   112.84730768203735
+  ],
+  "greedy": 11.496987581253052,
+  "total": 113.7428183555603,
+  "setting": "critic",
+  "k": 2,
+  "query": 1,
+  "subset": "vqav2_val"
+ },
+ {
+  "demos": [
+   106.2651150226593,
+   107.3609848022461
+  ],
+  "greedy": 11.693379402160645,
+  "total": 108.45507740974426,
+  "setting": "critic",
+  "k": 2,
+  "query": 2,
+  "subset": "ok_vqa_val2014"
+ },
+ {
+  "demos": [
+   97.19141793251038,
+   402.90441393852234
+  ],
+  "greedy": 11.953308343887329,
+  "total": 404.0378088951111,
+  "setting": "critic",
+  "k": 2,
+  "query": 3,
+  "subset": "textvqa_val"
+ },
+ {
+  "demos": [
+   125.35352444648743,
+   164.42774605751038
+  ],
+  "greedy": 11.947067499160767,
+  "total": 166.28691792488098,
+  "setting": "critic",
+  "k": 2,
+  "query": 4,
+  "subset": "chartqa"
+ }
+]