File size: 29,623 Bytes
eac7784
 
 
 
cb5eccc
eac7784
cb5eccc
eac7784
cb5eccc
 
eac7784
 
 
cb5eccc
eac7784
 
 
 
 
 
 
 
 
 
 
cb5eccc
eac7784
 
cb5eccc
 
eac7784
cb5eccc
 
 
 
 
 
 
 
 
 
eac7784
 
 
 
 
cb5eccc
 
eac7784
cb5eccc
 
eac7784
 
 
cb5eccc
 
 
 
 
 
eac7784
cb5eccc
 
 
 
 
 
eac7784
 
 
 
cb5eccc
 
 
 
 
 
 
 
 
eac7784
cb5eccc
 
 
 
 
 
 
 
 
eac7784
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
 
 
 
 
 
 
 
 
 
cb5eccc
 
 
 
 
 
eac7784
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
 
cb5eccc
eac7784
 
cb5eccc
eac7784
2977eaa
cb5eccc
eac7784
 
cb5eccc
 
 
 
 
 
 
 
 
eac7784
 
cb5eccc
eac7784
cb5eccc
 
 
 
 
eac7784
 
 
cb5eccc
 
eac7784
cb5eccc
 
 
 
 
 
 
eac7784
cb5eccc
 
 
eac7784
 
 
cb5eccc
 
eac7784
 
cb5eccc
 
 
eac7784
 
 
 
 
 
cb5eccc
 
 
 
 
eac7784
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
 
 
 
 
 
cb5eccc
 
eac7784
 
cb5eccc
eac7784
 
cb5eccc
 
eac7784
 
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
cb5eccc
 
 
eac7784
cb5eccc
 
 
eac7784
 
 
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
 
cb5eccc
 
 
 
 
eac7784
 
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7f4191e
 
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7f4191e
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3c27bc5
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3c27bc5
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3c27bc5
 
cb5eccc
 
 
 
 
 
 
eac7784
 
 
 
 
 
cb5eccc
eac7784
 
 
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
 
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
cb5eccc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eac7784
 
cb5eccc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must come before torch)

import base64  # noqa: E402
import difflib  # noqa: E402
import hashlib  # noqa: E402
import io  # noqa: E402
import json  # noqa: E402
import re  # noqa: E402
import shutil  # noqa: E402
import signal  # noqa: E402
import subprocess  # noqa: E402
import sys  # noqa: E402
import tempfile  # noqa: E402
import threading  # noqa: E402
import time  # noqa: E402
import uuid  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from jinja2 import StrictUndefined, Template  # noqa: E402
from PIL import Image  # noqa: E402
from transformers import AutoModelForImageTextToText, AutoProcessor, TextIteratorStreamer  # noqa: E402

ROOT = Path(__file__).resolve().parent
HARNESS_ROOT = ROOT / "harness"
sys.path.insert(0, str(HARNESS_ROOT))
for _tool in (HARNESS_ROOT / "bin").iterdir():  # the Hub does not keep the executable bit
    _tool.chmod(0o755)

from video_coder_harness.blender_exec import execute as blender_execute  # noqa: E402
from video_coder_harness.guard import (  # noqa: E402
    BLOCKED,
    REPEAT_WARNING,
    _normalize_output,
    harness_violation,
)
from video_coder_harness.static_check import check_file  # noqa: E402

MODEL_ID = "Video-Coder/Video-Coder-9B"

# --------------------------------------------------------------------------------------------------------------------
# Blender 4.5 (bpy 4.5.14 on Python 3.11, same as the authors' bpy-agent.sif sandbox), rendered on CPU with Mesa
# llvmpipe. Installed into its own venv at startup because the Space itself runs Python 3.12.
# --------------------------------------------------------------------------------------------------------------------
BPY_ENV = Path(os.environ.get("BPY_ENV", str(Path.home() / "bpyenv")))
BPY_PY = BPY_ENV / "bin" / "python"


def setup_bpy() -> None:
    if BPY_PY.exists():
        r = subprocess.run([str(BPY_PY), "-c", "import bpy; print(bpy.app.version_string)"],
                           capture_output=True, text=True)
        if r.returncode == 0:
            print("bpy ready:", r.stdout.strip())
            return
    t0 = time.time()
    subprocess.run([sys.executable, "-m", "uv", "venv", "-q", "-p", "3.11", "--clear", str(BPY_ENV)], check=True)
    subprocess.run([sys.executable, "-m", "uv", "pip", "install", "-q", "-p", str(BPY_PY), "bpy==4.5.14"],
                   check=True)
    r = subprocess.run([str(BPY_PY), "-c", "import bpy; print(bpy.app.version_string)"],
                       capture_output=True, text=True)
    print(f"bpy install took {time.time() - t0:.0f}s:", r.stdout.strip(), r.stderr[-2000:])


setup_bpy()

# Harness settings: configs/teacher.yaml + configs/uni_container_visual.yaml (the RL / inference setting).
HARNESS = {
    "blender_version": "4.5",
    "blender_available": True,
    "bpy_python": str(BPY_PY),
    "bpy_libs": "/usr",
    "exec_timeout": 600,
    "summary_samples": 6,
    "content_samples": 8,
    "visual_feedback": True,
    "preview_frames": 8,
    "preview_width": 640,
    "judge": {"enabled": False},
    "output_video": "output.mp4",
    "fps": 24,
    "frame_count": 120,
    "max_frame_count": 240,
    "resolution": [1280, 720],
    "script_name": "main.py",
}
SCRIPT = HARNESS["script_name"]
CMD_TIMEOUT = 900
REPEAT_WARN, REPEAT_STOP, REPEAT_SIM = 3, 6, 0.9
MAX_FORMAT_ERRORS = 3
RUNS_DIR = Path(tempfile.gettempdir()) / "video_coder_runs"
RUNS_DIR.mkdir(parents=True, exist_ok=True)
_SECRET_ENV = re.compile(r"(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL)", re.I)
HARNESS_MARKERS = [str(HARNESS_ROOT), str(BPY_ENV)]

SYSTEM_PROMPT = (
    "You are an expert Blender technical artist and Python programmer. You write Blender Python (bpy) scripts that "
    "build and animate 3D scenes and render them as untextured white-model (clay) videos. You work in a Linux shell "
    "through a bash tool."
)
# The exact user prompt of the SFT trajectories (Video-Coder/Train-Data sft_data.json), with the scene as __TASK__.
INSTANCE_TEMPLATE = (HARNESS_ROOT / "instance_template.txt").read_text()
TOOLS = [{
    "type": "function",
    "function": {
        "name": "bash",
        "description": "Execute a bash command",
        "parameters": {
            "type": "object",
            "properties": {"command": {"type": "string", "description": "The bash command to execute"}},
            "required": ["command"],
        },
    },
}]
OBSERVATION_TEMPLATE = Template("""{%- if output.output | length < 10000 -%}
{
  "returncode": {{ output.returncode }},
  "output": {{ output.output | tojson }}
  {%- if output.exception_info %}, "exception_info": {{ output.exception_info | tojson }}{% endif %}
}
{%- else -%}
{
  "returncode": {{ output.returncode }},
  "output_head": {{ output.output[:5000] | tojson }},
  "output_tail": {{ output.output[-5000:] | tojson }},
  "elided_chars": {{ output.output | length - 10000 }},
  "warning": "Output too long."
  {%- if output.exception_info %}, "exception_info": {{ output.exception_info | tojson }}{% endif %}
}
{%- endif -%}
""", undefined=StrictUndefined)
FORMAT_ERROR_TEMPLATE = Template("""{% if finish_reason is defined and (finish_reason == "length" or (finish_reason == "tool_calls" and not has_tool_calls)) -%}
Your previous response reached the output token limit (finish_reason={{ finish_reason }}) before you produced a tool call, so it was cut off. Respond more concisely and finish with exactly one bash tool call.
{%- else -%}
Tool call error:

<error>
{{error}}
</error>

Every response must call the 'bash' tool exactly once with {"command": "..."}.
To finish, run `echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` on its own.
{%- endif %}
""")
MM_TAG = re.compile(r"(?s)<MSWEA_MULTIMODAL_CONTENT><CONTENT_TYPE>(.+?)</CONTENT_TYPE>(.+?)</MSWEA_MULTIMODAL_CONTENT>")
TOOL_CALL_RE = re.compile(r"<tool_call>\s*<function=([^>\s]+)>(.*?)</function>\s*</tool_call>", re.S)
PARAM_RE = re.compile(r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>", re.S)
END_TOKENS = ("<|im_end|>", "<|endoftext|>")

# --------------------------------------------------------------------------------------------------------------------
# Model: Video-Coder-9B (Qwen3.5-9B, vision-language) in bf16, fully on the GPU.
# --------------------------------------------------------------------------------------------------------------------
processor = AutoProcessor.from_pretrained(MODEL_ID)
tokenizer = processor.tokenizer
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda").eval()

GPU_CHUNK_SECONDS = 90  # one GPU call generates for at most this long; long turns continue in the next call


@spaces.GPU(duration=120)
def llm_chunk(messages: list, images: list, prefix: str = "", max_new_tokens: int = 16384,
              temperature: float = 1.0, top_p: float = 0.95):
    """Continue the agent's current turn for up to GPU_CHUNK_SECONDS; yields (text so far, finished, n_tokens)."""
    prompt = tokenizer.apply_chat_template(messages, tools=TOOLS, add_generation_prompt=True, tokenize=False)
    prompt += prefix
    inputs = processor(text=[prompt], images=images or None, return_tensors="pt").to("cuda")
    n_in = inputs["input_ids"].shape[1]
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
    result = {}

    def _run():
        try:
            out = model.generate(**inputs, streamer=streamer, max_new_tokens=int(max_new_tokens), do_sample=True,
                                 temperature=float(temperature), top_p=float(top_p), top_k=0, use_cache=True,
                                 max_time=GPU_CHUNK_SECONDS)
            result["n"] = int(out.shape[1] - n_in)
        except Exception as e:  # surfaced below
            result["error"] = f"{type(e).__name__}: {e}"
            streamer.end()

    th = threading.Thread(target=_run, daemon=True)
    th.start()
    text = ""
    last = 0.0
    for piece in streamer:
        text += piece
        if time.time() - last > 0.4:
            last = time.time()
            yield text, False, None
    th.join()
    if "error" in result:
        raise gr.Error(result["error"])
    finished = any(t in text for t in END_TOKENS)
    yield text, finished, result.get("n", 0)


# --------------------------------------------------------------------------------------------------------------------
# Agent loop: mini-swe-agent DefaultAgent + video_coder_harness GuardedAgent semantics, re-implemented without the
# litellm/vLLM gateway (the model is called in-process).
# --------------------------------------------------------------------------------------------------------------------
def agent_env() -> dict:
    env = dict(os.environ)
    env.update({k: "" for k in os.environ if _SECRET_ENV.search(k)})
    env.update({"PAGER": "cat", "MANPAGER": "cat", "PIP_PROGRESS_BAR": "off", "TQDM_DISABLE": "1"})
    env["PATH"] = f"{HARNESS_ROOT / 'bin'}:{os.environ.get('PATH', '')}"
    env["VCSFT_PYTHON"] = sys.executable
    env["VCSFT_HARNESS"] = json.dumps(HARNESS)
    return env


def run_command(command: str, cwd: str) -> dict:
    """mini-swe-agent LocalEnvironment.execute."""
    proc = subprocess.Popen(command, shell=True, text=True, cwd=cwd, env=agent_env(), encoding="utf-8",
                            errors="replace", stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            start_new_session=True)
    try:
        stdout, _ = proc.communicate(timeout=CMD_TIMEOUT)
        return {"output": stdout, "returncode": proc.returncode, "exception_info": ""}
    except subprocess.TimeoutExpired as e:
        os.killpg(proc.pid, signal.SIGKILL)
        stdout, _ = proc.communicate()
        return {"output": stdout or "", "returncode": -1,
                "exception_info": f"An error occurred while executing the command: {e}"}


def is_submit(output: dict) -> bool:
    lines = output.get("output", "").lstrip().splitlines()
    return bool(lines) and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output["returncode"] == 0


def split_images(text: str):
    """video_coder_harness.models.split_images: cut frame images out of the tool output."""
    images = []

    def repl(m):
        if m.group(1).strip() != "image_url":
            return m.group(0)
        before = text[: m.start()].rstrip().splitlines()
        caption = before[-1].strip() if before else ""
        images.append((caption, m.group(2).strip()))
        return f"[image {len(images)} attached below]"

    return MM_TAG.sub(repl, text), images


def data_url_to_pil(url: str) -> Image.Image:
    return Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).convert("RGB")


def parse_response(raw: str):
    """Split a Qwen3.5 response into (reasoning, content, tool calls) like vLLM's qwen3 reasoning + qwen3_coder parsers."""
    text = raw
    for t in END_TOKENS:
        text = text.split(t)[0]
    reasoning = ""
    if "</think>" in text:
        reasoning, text = text.split("</think>", 1)
        reasoning = reasoning.replace("<think>", "").strip()
    calls = []
    for m in TOOL_CALL_RE.finditer(text):
        args = {k: v for k, v in PARAM_RE.findall(m.group(2))}
        calls.append({"name": m.group(1).strip(), "arguments": args})
    content = TOOL_CALL_RE.sub("", text)
    content = content.split("<tool_call>")[0].strip()
    return reasoning, content, calls


def sha(path: Path) -> str | None:
    return hashlib.sha256(path.read_bytes()).hexdigest() if path.is_file() else None


def short(text: str, n: int = 60) -> str:
    lines = text.splitlines()
    return text if len(lines) <= n else "\n".join(lines[:n]) + f"\n... ({len(lines) - n} more lines)"


def read_script(workdir: Path) -> str:
    p = workdir / SCRIPT
    return p.read_text(errors="replace") if p.is_file() else ""


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 = 16, 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.
        max_steps: maximum number of agent turns (one bash command each).
        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.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a scene description.")
    state = new_state(prompt, max_steps, temperature, top_p, max_new_tokens)
    return ui(state, tick=time.time())


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
    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
    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:
        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())


# --------------------------------------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1280px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    "The video captures a small bird perched on a tree branch. The bird, with its brown and white feathers, yellow "
    "beak, and black and white head, is the main focus of the video. The bird is seen in three different positions on "
    "the branch, each time facing a different direction. The background of the video is filled with green foliage, "
    "providing a natural and serene setting for the bird. The style of the video is a close-up shot.",
    "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 "
    "background features a city skyline with buildings of various heights, and the sky is a gradient of orange and "
    "yellow, suggesting either sunrise or sunset. The style of the video is realistic with a focus on the cars and "
    "their movement.",
    "a man standing in front of a machine.",
    "The video captures a lively scene on the ocean, with three boats in motion. The first boat, a large white yacht, "
    "is seen in the distance, sailing towards the horizon. The second and third boats are closer to the viewer, both "
    "moving at high speeds, leaving trails of white water behind them. The ocean is a deep blue, and the sky above is "
    "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.",
]

with gr.Blocks(title="Video-Coder") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎬 Video-Coder: agentic coding for video generation\n"
            "**Video-Coder-9B** writes and runs a Blender Python program for your scene description, inspects the "
            "rendered frames, fixes what is wrong and submits a **draft video** (white-model 3D render with the "
            "layout, motion and camera of the scene) — the structure/motion signal the paper feeds to a downstream "
            "video generator.\n\n"
            "[📖 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: 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'")
            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):
                video_out = gr.Video(label="Draft video (Blender render)", autoplay=True)
                preview_out = gr.Image(label="Latest render: frames 1, ⅓, ⅔, last", type="filepath")
            with gr.Column(scale=1):
                with gr.Tabs():
                    with gr.Tab("Agent trajectory"):
                        chat_out = gr.Chatbot(label="Video-Coder harness", height=640, show_label=False)
                    with gr.Tab("main.py"):
                        code_out = gr.Code(label="Blender script written by the agent", language="python",
                                           lines=30)
        with gr.Accordion("Advanced settings", open=False):
            max_steps = gr.Slider(2, 40, value=16, step=1, label="Max agent steps (one bash command each)")
            temperature = gr.Slider(0.1, 1.5, value=1.0, step=0.05, label="Temperature")
            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")

        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__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)