Per-step requests (fresh ZeroGPU token each GPU call)

app.py

@@ -284,10 +284,47 @@ def read_script(workdir: Path) -> str:
     return p.read_text(errors="replace") if p.is_file() else ""
 
 
-def run_agent(prompt: str, max_steps: int = 12, temperature: float = 1.0, top_p: float = 0.95,
-              max_new_tokens: int = 16384, request: gr.Request = None):
-    """Turn a scene description into an animated 3D "draft video": Video-Coder-9B writes a Blender script, renders it,
-    inspects the frames, fixes problems and submits.
+def new_state(prompt: str, max_steps: int, temperature: float, top_p: float, max_new_tokens: int) -> dict:
+    workdir = (RUNS_DIR / uuid.uuid4().hex).resolve()
+    workdir.mkdir(parents=True)
+    return {
+        "workdir": str(workdir),
+        "messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": INSTANCE_TEMPLATE.replace("__TASK__", prompt.strip())},
+        ],
+        "image_urls": [],  # frames shown to the model, in order (JPEG data URLs)
+        "chat": [],
+        "video": None,
+        "preview": None,
+        "last_ok": None,  # (script hash, video) of the last run_blender that reported RUN OK
+        "last_key": None,
+        "repeats": 0,
+        "n_format_errors": 0,
+        "step": 1,
+        "raw": "",  # current (possibly unfinished) model turn
+        "used": 0,
+        "max_steps": int(max_steps),
+        "temperature": float(temperature),
+        "top_p": float(top_p),
+        "max_new_tokens": int(max_new_tokens),
+        "done": False,
+        "status": "Starting…",
+        "t_start": time.time(),
+    }
+
+
+def ui(state: dict, tick=None):
+    workdir = Path(state["workdir"])
+    msg = f"**{state['status']}** · {time.time() - state['t_start']:.0f}s elapsed"
+    return (state, state["chat"], state["video"], read_script(workdir), state["preview"], msg,
+            gr.update() if tick is None else tick)
+
+
+def start(prompt: str, max_steps: int = 12, temperature: float = 1.0, top_p: float = 0.95,
+          max_new_tokens: int = 16384):
+    """Start a Video-Coder agent run for a scene description. The run then proceeds turn by turn: call `advance`
+    with the returned state (same session) until its status says it finished.
 
     Args:
         prompt: description of the scene / video to build.
@@ -295,185 +332,211 @@ def run_agent(prompt: str, max_steps: int = 12, temperature: float = 1.0, top_p:
         temperature: sampling temperature (the authors sample with 1.0).
         top_p: nucleus sampling (the authors use 0.95).
         max_new_tokens: token budget of a single agent turn.
-
-    Returns:
-        (agent trajectory, draft video, main.py, latest render preview, status)
     """
     if not prompt or not prompt.strip():
         raise gr.Error("Please enter a scene description.")
-    try:
-        _tok = request.headers.get("x-ip-token") if request is not None else None
-        if _tok:
-            _p = json.loads(base64.urlsafe_b64decode(_tok.split(".")[1] + "=="))
-            print("x-ip-token payload:", {k: v for k, v in _p.items() if k in ("exp", "iat", "nbf")}, "now", time.time())
-        else:
-            print("no x-ip-token")
-    except Exception as e:
-        print("token decode failed", e)
-    workdir = (RUNS_DIR / uuid.uuid4().hex).resolve()
-    workdir.mkdir(parents=True)
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {"role": "user", "content": INSTANCE_TEMPLATE.replace("__TASK__", prompt.strip())},
-    ]
-    images: list = []
-    chat: list = []
-    video = None
-    preview = None
-    last_ok = None  # (script hash, video path) of the last run_blender that reported RUN OK
-    last_key = None
-    repeats = 0
-    n_format_errors = 0
-    exit_status = "StepLimit"
-    t_start = time.time()
-
-    def status(msg):
-        return f"**{msg}** · {time.time() - t_start:.0f}s elapsed"
-
-    def ui(msg):
-        return chat, video, read_script(workdir), preview, status(msg)
-
-    yield ui("Starting…")
-    for step in range(1, int(max_steps) + 1):
-        # ---- model turn (possibly split over several GPU calls) ----
-        raw, used = "", 0
-        live = {"role": "assistant", "content": "", "metadata": {"title": f"💭 Step {step}: thinking…",
-                                                                 "status": "pending"}}
-        chat.append(live)
-        while True:
-            finished, n = False, None
-            base = raw
-            for text, finished, n in llm_chunk(messages, images, base, max_new_tokens - used, temperature, top_p):
-                raw = base + text
-                live["content"] = raw[-3000:].replace("<think>", "")
-                yield ui(f"Step {step}/{max_steps}: Video-Coder is writing ({len(raw)} chars)")
-            used += n or 0
-            if finished or used >= max_new_tokens:
-                break
-        chat.pop()
-        reasoning, content, calls = parse_response(raw)
-        if reasoning:
-            chat.append({"role": "assistant", "content": reasoning,
-                         "metadata": {"title": f"💭 Step {step}: reasoning", "status": "done"}})
-        msg = {"role": "assistant", "content": content}
-        if reasoning:
-            msg["reasoning_content"] = reasoning
-        if calls:
-            msg["tool_calls"] = [{"type": "function", "function": {"name": c["name"], "arguments": c["arguments"]}}
-                                 for c in calls]
-        messages.append(msg)
-
-        # ---- format errors ----
-        error = None
-        if not calls:
-            error = "No tool calls found in the response. Every response MUST include at least one tool call."
-        elif any(c["name"] != "bash" or "command" not in c["arguments"] for c in calls):
-            error = "Unknown tool or missing 'command' argument."
-        if error:
-            n_format_errors += 1
-            finish_reason = "length" if not any(t in raw for t in END_TOKENS) else "stop"
-            err = FORMAT_ERROR_TEMPLATE.render(error=error, finish_reason=finish_reason, has_tool_calls=bool(calls))
-            if content:
-                chat.append({"role": "assistant", "content": content})
-            chat.append({"role": "user", "content": f"⚠️ {err}"})
-            messages.append({"role": "user", "content": err})
-            yield ui(f"Step {step}: format error")
-            if n_format_errors >= MAX_FORMAT_ERRORS:
-                exit_status = "FormatError"
-                break
-            continue
-        n_format_errors = 0
-
-        # ---- execute the bash command(s) ----
-        submitted = False
-        stop = False
-        for call in calls:
-            command = call["arguments"]["command"]
-            shown = f"{content}\n\n" if content else ""
-            chat.append({"role": "assistant", "content": f"{shown}```bash\n{short(command)}\n```"})
-            yield ui(f"Step {step}: running `{command.strip().splitlines()[0][:60]}`")
-            reason = harness_violation(command, HARNESS_MARKERS, str(workdir))
-            if reason:
-                out = {"output": BLOCKED.format(reason=reason), "returncode": 1, "exception_info": ""}
-            else:
-                out = run_command(command, str(workdir))
-                if is_submit(out):
-                    submitted = True
-                    break
-            key = (command, _normalize_output(out.get("output", "")))
-            same = (last_key is not None and key[1] == last_key[1]
-                    and difflib.SequenceMatcher(None, key[0], last_key[0]).ratio() >= REPEAT_SIM)
-            repeats = repeats + 1 if same else 1
-            last_key = key
-            if repeats >= REPEAT_WARN:
-                out["output"] = (out.get("output") or "") + REPEAT_WARNING.format(n=repeats, stop=REPEAT_STOP)
-            if repeats >= REPEAT_STOP:
-                stop = True
-
-            text, imgs = split_images(out.get("output") or "")
-            obs = OBSERVATION_TEMPLATE.render(output={**out, "output": text})
-            if imgs:
-                parts = [{"type": "text", "text": obs}]
-                for k, (caption, url) in enumerate(imgs, 1):
-                    parts += [{"type": "text", "text": f"image {k}: {caption}" if caption else f"image {k}:"},
-                              {"type": "image"}]
-                    images.append(data_url_to_pil(url))
-                messages.append({"role": "tool", "content": parts})
-            else:
-                messages.append({"role": "tool", "content": obs})
-
-            if command.strip().startswith("run_blender"):
-                vid = workdir / HARNESS["output_video"]
-                if vid.is_file():
-                    keep = workdir / f".render_step{step}.mp4"
-                    shutil.copy(vid, keep)
-                    video = str(keep)
-                    if text.startswith("RUN OK"):
-                        last_ok = (sha(workdir / SCRIPT), str(keep))
-                if (workdir / "preview.png").is_file():
-                    keep_png = workdir / f".preview_step{step}.png"
-                    shutil.copy(workdir / "preview.png", keep_png)
-                    preview = str(keep_png)
-            chat.append({"role": "assistant", "content": f"```\n{short(text.strip(), 40)}\n```",
-                         "metadata": {"title": f"🛠️ Output (exit code {out['returncode']})", "status": "done"}})
-            yield ui(f"Step {step}/{max_steps} done")
-        if submitted:
-            exit_status = "Submitted"
-            break
-        if stop:
-            exit_status = "RepeatedActions"
-            break
+    state = new_state(prompt, max_steps, temperature, top_p, max_new_tokens)
+    return ui(state, tick=time.time())
 
-    # ---- acceptance: static check + the final script re-rendered from scratch (harness final_execution) ----
+
+def _finish_turn(state: dict):
+    """Parse the finished model turn, run its command(s) in the harness; a generator of UI updates."""
+    workdir = Path(state["workdir"])
+    step, raw = state["step"], state["raw"]
+    chat, messages = state["chat"], state["messages"]
+    if chat and chat[-1].get("metadata", {}).get("status") == "pending":
+        chat.pop()
+    reasoning, content, calls = parse_response(raw)
+    state["raw"], state["used"] = "", 0
+    if reasoning:
+        chat.append({"role": "assistant", "content": reasoning,
+                     "metadata": {"title": f"💭 Step {step}: reasoning", "status": "done"}})
+    msg = {"role": "assistant", "content": content}
+    if reasoning:
+        msg["reasoning_content"] = reasoning
+    if calls:
+        msg["tool_calls"] = [{"type": "function", "function": {"name": c["name"], "arguments": c["arguments"]}}
+                             for c in calls]
+    messages.append(msg)
+
+    error = None
+    if not calls:
+        error = "No tool calls found in the response. Every response MUST include at least one tool call."
+    elif any(c["name"] != "bash" or "command" not in c["arguments"] for c in calls):
+        error = "Unknown tool or missing 'command' argument."
+    if error:
+        state["n_format_errors"] += 1
+        finish_reason = "length" if not any(t in raw for t in END_TOKENS) else "stop"
+        err = FORMAT_ERROR_TEMPLATE.render(error=error, finish_reason=finish_reason, has_tool_calls=bool(calls))
+        if content:
+            chat.append({"role": "assistant", "content": content})
+        chat.append({"role": "user", "content": f"⚠️ {err}"})
+        messages.append({"role": "user", "content": err})
+        if state["n_format_errors"] >= MAX_FORMAT_ERRORS:
+            return "FormatError"
+        return None
+    state["n_format_errors"] = 0
+
+    for call in calls:
+        command = call["arguments"]["command"]
+        shown = f"{content}\n\n" if content else ""
+        chat.append({"role": "assistant", "content": f"{shown}```bash\n{short(command)}\n```"})
+        first = command.strip().splitlines()[0][:60] if command.strip() else ""
+        state["status"] = f"Step {step}: running `{first}`"
+        yield
+        reason = harness_violation(command, HARNESS_MARKERS, str(workdir))
+        if reason:
+            out = {"output": BLOCKED.format(reason=reason), "returncode": 1, "exception_info": ""}
+        else:
+            out = run_command(command, str(workdir))
+            if is_submit(out):
+                return "Submitted"
+        key = (command, _normalize_output(out.get("output", "")))
+        last = state["last_key"]
+        same = (last is not None and key[1] == last[1]
+                and difflib.SequenceMatcher(None, key[0], last[0]).ratio() >= REPEAT_SIM)
+        state["repeats"] = state["repeats"] + 1 if same else 1
+        state["last_key"] = key
+        if state["repeats"] >= REPEAT_WARN:
+            out["output"] = (out.get("output") or "") + REPEAT_WARNING.format(n=state["repeats"], stop=REPEAT_STOP)
+
+        text, imgs = split_images(out.get("output") or "")
+        obs = OBSERVATION_TEMPLATE.render(output={**out, "output": text})
+        if imgs:
+            parts = [{"type": "text", "text": obs}]
+            for k, (caption, url) in enumerate(imgs, 1):
+                parts += [{"type": "text", "text": f"image {k}: {caption}" if caption else f"image {k}:"},
+                          {"type": "image"}]
+                state["image_urls"].append(url)
+            messages.append({"role": "tool", "content": parts})
+        else:
+            messages.append({"role": "tool", "content": obs})
+
+        if command.strip().startswith("run_blender"):
+            vid = workdir / HARNESS["output_video"]
+            if vid.is_file():
+                keep = workdir / f"render_step{step}.mp4"
+                shutil.copy(vid, keep)
+                state["video"] = str(keep)
+                if text.startswith("RUN OK"):
+                    state["last_ok"] = (sha(workdir / SCRIPT), str(keep))
+            if (workdir / "preview.png").is_file():
+                keep_png = workdir / f"preview_step{step}.png"
+                shutil.copy(workdir / "preview.png", keep_png)
+                state["preview"] = str(keep_png)
+        chat.append({"role": "assistant", "content": f"```\n{short(text.strip(), 40)}\n```",
+                     "metadata": {"title": f"🛠️ Output (exit code {out['returncode']})", "status": "done"}})
+        if state["repeats"] >= REPEAT_STOP:
+            return "RepeatedActions"
+    return None
+
+
+def _accept(state: dict):
+    """Harness acceptance: static check + the submitted script re-rendered from scratch (final_execution)."""
+    workdir = Path(state["workdir"])
     script = workdir / SCRIPT
-    if exit_status == "Submitted" and script.is_file():
-        chat.append({"role": "assistant", "content": "✅ Submitted. Validating the final script…"})
-        yield ui("Validating the submitted script")
-        check = check_file(script)
-        if not check["ok"]:
-            yield ui(f"Submitted, but the static check failed: {check['errors']}")
+    state["chat"].append({"role": "assistant", "content": "✅ Submitted. Validating the final script…"})
+    if not script.is_file():
+        state["status"] = f"Submitted, but no {SCRIPT} was written"
+        return
+    check = check_file(script)
+    if not check["ok"]:
+        state["status"] = f"Submitted, but the static check failed: {check['errors']}"
+        return
+    last_ok = state["last_ok"]
+    if last_ok and last_ok[0] == sha(script):  # Workbench rendering is deterministic: reuse the identical render
+        state["video"] = last_ok[1]
+    else:
+        state["status"] = "Re-rendering the submitted script"
+        yield
+        final_dir = workdir / "final_run"
+        final_dir.mkdir(exist_ok=True)
+        shutil.copy(script, final_dir / SCRIPT)
+        report = blender_execute(final_dir / SCRIPT, final_dir, HARNESS)
+        if (final_dir / HARNESS["output_video"]).is_file():
+            state["video"] = str(final_dir / HARNESS["output_video"])
+        if (final_dir / "preview.png").is_file():
+            state["preview"] = str(final_dir / "preview.png")
+        if not report["ok"]:
+            state["status"] = f"Submitted, but the final render has problems: {report['problems']}"
             return
-        if last_ok and last_ok[0] == sha(script):  # rendering is deterministic: reuse the identical render
-            video = last_ok[1]
-        else:
-            yield ui("Re-rendering the submitted script")
-            final_dir = workdir / ".final_run"
-            final_dir.mkdir(exist_ok=True)
-            shutil.copy(script, final_dir / SCRIPT)
-            report = blender_execute(final_dir / SCRIPT, final_dir, HARNESS)
-            if (final_dir / HARNESS["output_video"]).is_file():
-                video = str(final_dir / HARNESS["output_video"])
-            if (final_dir / "preview.png").is_file():
-                preview = str(final_dir / "preview.png")
-            if not report["ok"]:
-                yield ui(f"Submitted, but the final render has problems: {report['problems']}")
-                return
-        yield ui("Accepted ✅ — draft video rendered from the agent's Blender script")
+    state["status"] = "Accepted ✅ — draft video rendered from the agent's Blender script"
+
+
+def advance(state: dict):
+    """Run the next piece of the agent: one GPU generation chunk of the current turn and, once the turn is complete,
+    its bash command in the Video-Coder harness. Returns a new tick while the run is still going.
+
+    Args:
+        state: the run state returned by `start` / the previous `advance`.
+    """
+    if not state or state.get("done"):
+        yield ui(state) if state else (state, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(),
+                                       gr.update())
+        return
+    step = state["step"]
+    chat = state["chat"]
+    if not (chat and chat[-1].get("metadata", {}).get("status") == "pending"):
+        chat.append({"role": "assistant", "content": "",
+                     "metadata": {"title": f"💭 Step {step}: thinking…", "status": "pending"}})
+    live = chat[-1]
+    state["status"] = f"Step {step}/{state['max_steps']}: Video-Coder is thinking"
+    yield ui(state)
+
+    # ---- one GPU call (the ZeroGPU token of a request is short-lived, so each request makes exactly one) ----
+    images = [data_url_to_pil(u) for u in state["image_urls"]]
+    base = state["raw"]
+    finished, n = False, None
+    t0 = time.time()
+    for text, finished, n in llm_chunk(state["messages"], images, base, state["max_new_tokens"] - state["used"],
+                                       state["temperature"], state["top_p"]):
+        state["raw"] = base + text
+        live["content"] = state["raw"][-3000:].replace("<think>", "")
+        state["status"] = f"Step {step}/{state['max_steps']}: Video-Coder is writing ({len(state['raw'])} chars)"
+        yield ui(state)
+    state["used"] += n or 0
+    print(f"step {step}: chunk {time.time() - t0:.1f}s, {n} tokens, finished={finished}")
+    if not finished and state["used"] < state["max_new_tokens"]:
+        yield ui(state, tick=time.time())  # continue this turn in the next request
+        return
+
+    # ---- the turn is complete: execute it ----
+    exit_status = None
+    gen = _finish_turn(state)
+    while True:
+        try:
+            next(gen)
+            yield ui(state)
+        except StopIteration as stop:
+            exit_status = stop.value
+            break
+    state["step"] += 1
+    if exit_status is None and state["step"] > state["max_steps"]:
+        exit_status = "StepLimit"
+    if exit_status is None:
+        state["status"] = f"Step {step}/{state['max_steps']} done"
+        yield ui(state, tick=time.time())
+        return
+    state["done"] = True
+    if exit_status == "Submitted":
+        for _ in _accept(state):
+            yield ui(state)
     else:
-        note = {"StepLimit": "Step limit reached before the agent submitted (showing the latest render)",
-                "RepeatedActions": "Stopped: the agent repeated the same command",
-                "FormatError": "Stopped after repeated format errors"}[exit_status]
-        yield ui(note)
+        state["status"] = {
+            "StepLimit": "Step limit reached before the agent submitted (showing the latest render)",
+            "RepeatedActions": "Stopped: the agent repeated the same command",
+            "FormatError": "Stopped after repeated format errors",
+        }[exit_status]
+    yield ui(state)
+
+
+def stop_run(state: dict):
+    """Stop the current run."""
+    if state:
+        state["done"] = True
+        state["status"] = "Stopped"
+        return ui(state)
+    return (state, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update())
 
 
 # --------------------------------------------------------------------------------------------------------------------
@@ -486,6 +549,8 @@ CSS = """
 
 EXAMPLES = [
     "a man standing in front of a machine.",
+    "A single polished walnut-wood chess pawn stands on a small square pedestal of white Carrara marble. No other "
+    "objects are in the scene. The camera makes one complete 360-degree orbit around the pawn while slowly rising.",
     "The video shows a sequence of three identical white cars driving on a road. The first car is in the foreground, "
     "the second car is slightly behind and to the right, and the third car is in the background. The cars are moving "
     "from left to right across the frame. The road appears to be a city street with a dark grey surface. The "
@@ -498,8 +563,6 @@ EXAMPLES = [
     "clear, suggesting a sunny day. The boats are the main focus of the video, with their movement and the water's "
     "reaction to their speed creating a dynamic and exciting scene. The overall style of the video is a dynamic "
     "aerial shot, capturing the movement and beauty of the ocean and the boats.",
-    "A single polished walnut-wood chess pawn stands on a small square pedestal of white Carrara marble. No other "
-    "objects are in the scene. The camera makes one complete 360-degree orbit around the pawn while slowly rising.",
 ]
 
 with gr.Blocks(title="Video-Coder") as demo:
@@ -513,14 +576,16 @@ with gr.Blocks(title="Video-Coder") as demo:
             "[📖 Paper](https://huggingface.co/papers/2611.09894) · [🌐 Project page](https://cywang735.github.io/"
             "Video-Coder/) · [💻 Code](https://github.com/CYWang735/Video-Coder) · [🤗 Model](https://huggingface.co/"
             "Video-Coder/Video-Coder-9B)\n\n"
-            "⏳ A full agent run takes several minutes (each turn generates a long script; each render runs Blender "
-            "on CPU). The trajectory streams live below."
+            "⏳ A full agent run takes several minutes: the first turn writes a long script, and every `run_blender` "
+            "renders the animation with Blender on CPU. The trajectory streams live; keep this tab open."
         )
         with gr.Row():
             prompt = gr.Textbox(label="Scene description", lines=3, scale=4,
                                 placeholder="Describe the scene and its motion, e.g. 'a red ball rolls down a ramp "
                                             "and knocks over three dominoes'")
-            run_btn = gr.Button("Generate draft video", variant="primary", scale=1)
+            with gr.Column(scale=1, min_width=180):
+                run_btn = gr.Button("Generate draft video", variant="primary")
+                stop_btn = gr.Button("Stop", variant="stop")
         status = gr.Markdown()
         with gr.Row():
             with gr.Column(scale=1):
@@ -539,12 +604,16 @@ with gr.Blocks(title="Video-Coder") as demo:
             top_p = gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top-p")
             max_new_tokens = gr.Slider(2048, 32768, value=16384, step=1024, label="Max tokens per agent turn")
 
-        outputs = [chat_out, video_out, code_out, preview_out, status]
-        run_btn.click(run_agent, inputs=[prompt, max_steps, temperature, top_p, max_new_tokens], outputs=outputs,
-                      api_name="generate", concurrency_limit=2)
-        prompt.submit(run_agent, inputs=[prompt, max_steps, temperature, top_p, max_new_tokens], outputs=outputs,
-                      api_visibility="private", concurrency_limit=2)
-        gr.Examples(examples=[[e] for e in EXAMPLES], inputs=[prompt], outputs=outputs, fn=run_agent,
+        state = gr.State(None)
+        tick = gr.Number(value=0, visible=False)
+        outputs = [state, chat_out, video_out, code_out, preview_out, status, tick]
+        inputs = [prompt, max_steps, temperature, top_p, max_new_tokens]
+        run_btn.click(start, inputs=inputs, outputs=outputs, api_name="start")
+        prompt.submit(start, inputs=inputs, outputs=outputs, api_visibility="private")
+        # Each agent step is its own request (fresh ZeroGPU token); the tick change chains the next one.
+        adv = tick.change(advance, inputs=[state], outputs=outputs, api_name="advance", concurrency_limit=None)
+        stop_btn.click(stop_run, inputs=[state], outputs=outputs, cancels=[adv], api_name="stop")
+        gr.Examples(examples=[[e] for e in EXAMPLES], inputs=[prompt], outputs=outputs, fn=start,
                     cache_examples=False, run_on_click=True)
 
 if __name__ == "__main__":