Live 3D preview while generating, spinner + timing stats

app.py

@@ -1,8 +1,8 @@
 import ctypes
 import glob
-import html
 import re
 import site
+import time
 
 import gradio as gr
 import spaces
@@ -61,13 +61,6 @@ def extract_html(text):
     return m.group(0) if m else text.strip()
 
 
-def iframe(code):
-    return (
-        f'<iframe sandbox="allow-scripts" srcdoc="{html.escape(code, quote=True)}" '
-        'style="width:100%;height:640px;border:0;border-radius:8px;background:#000"></iframe>'
-    )
-
-
 # prefill the code fence so the model skips chatty preambles ("I cannot create files...")
 NO_THINK = "<think>\n\n</think>\n\n```html\n"
 
@@ -76,16 +69,19 @@ NO_THINK = "<think>\n\n</think>\n\n```html\n"
 @spaces.GPU(duration=107)
 def generate(prompt, thinking, max_tokens, temperature, top_p, top_k):
     global llm
+    t0 = time.perf_counter()
     if llm is None:
         from llama_cpp import Llama  # import inside GPU worker so CUDA is available
 
         llm = Llama(MODEL_PATH, n_gpu_layers=-1, n_ctx=32768, flash_attn=True, verbose=False)
+    t1 = time.perf_counter()
     # ponytail: hand-built ChatML so thinking can be toggled; switch to create_chat_completion if template kwargs land
     text = (
         f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
         + ("<think>\n" if thinking else NO_THINK)
     )
     out = "<think>\n" if thinking else "```html\n"
+    n, t2 = 0, None
     for chunk in llm(
         text,
         max_tokens=int(max_tokens),
@@ -97,18 +93,104 @@ def generate(prompt, thinking, max_tokens, temperature, top_p, top_k):
         stream=True,
     ):
         out += chunk["choices"][0]["text"]
-        yield out
-    yield out
+        n += 1
+        t2 = t2 or time.perf_counter()
+        yield out, f"{n} tokens, {n / (time.perf_counter() - t2 + 1e-9):.1f} tok/s"
+    t3 = time.perf_counter()
+    t2 = t2 or t3
+    # logged so the @spaces.GPU duration can be sized from real runs
+    stats = (
+        f"GPU time {t3 - t0:.1f}s: model load {t1 - t0:.1f}s, prefill {t2 - t1:.1f}s, "
+        f"{n} tokens in {t3 - t2:.1f}s ({n / (t3 - t2 + 1e-9):.1f} tok/s)"
+    )
+    print(stats, flush=True)
+    yield out, stats
 
 
 GPU, WEB = "ZeroGPU (server)", "WebGPU (your browser)"
 
-# Same prompt/extract/iframe flow as the server path, but llama.cpp runs in the browser via wllama.
+# Shared by both backends: window.wgLive(rawOutput) re-renders the partial scene into #wg-live every few seconds.
+# Unfinished JS is cut back to the last complete-looking line, open brackets are closed and a fallback render loop
+# is appended; the new iframe only replaces the old one once that patched script actually ran (it posts 'wg-ok').
+LIVE_JS = r"""
+window.wgCode ??= (out) => {
+  const body = out.replace(/<think>[\s\S]*?(<\/think>|$)/g, '');
+  let m = body.match(/```html\s*([\s\S]*?)(```|$)/);
+  return m ? m[1].trim() : ((m = body.match(/<!DOCTYPE html[\s\S]*?(<\/html>|$)/i)) ? m[0] : body.trim());
+};
+window.wgLive ??= (() => {
+  const TAIL = "\n;try{const R=renderer,S=scene,C=camera;(function f(){try{R.render(S,C)}catch(e){}"
+    + "requestAnimationFrame(f)})()}catch(e){}\nparent.postMessage('wg-ok','*');\n</scr" + "ipt>";
+  const closers = (s) => {
+    const st = [];
+    for (let i = 0; i < s.length; i++) {
+      const c = s[i];
+      if (c === '/' && s[i + 1] === '/') { i = s.indexOf('\n', i); if (i < 0) break; }
+      else if (c === '/' && s[i + 1] === '*') { i = s.indexOf('*/', i + 2) + 1; if (i <= 0) break; }
+      else if (c === '"' || c === "'" || c === '`') { while (++i < s.length && s[i] !== c) if (s[i] === '\\') i++; }
+      else if ('([{'.includes(c)) st.push({ '(': ')', '[': ']', '{': '}' }[c]);
+      else if (')]}'.includes(c)) st.pop();
+    }
+    return st.reverse().join('');
+  };
+  // -> [html, isPartial], or null when there is no script to run yet
+  const patch = (code) => {
+    const i = code.lastIndexOf('<script'), j = code.indexOf('>', i) + 1;
+    if (i < 0 || !j) return null;
+    if (/<\/script>/i.test(code.slice(j))) return [code, false];
+    const lines = code.slice(j).split('\n');
+    lines.pop();
+    while (lines.length && !/[;{},]\s*$/.test(lines.at(-1))) lines.pop();
+    const js = lines.join('\n');
+    return [code.slice(0, j) + js + '\n' + closers(js) + TAIL, true];
+  };
+  let latest = '', last = 0, timer = 0, pending = null, shown = '';
+  const box = () => document.querySelector('#wg-live');
+  const show = (f) => {
+    if (pending !== f) return;
+    pending = null; shown = f.srcdoc;
+    for (const c of [...box().children]) if (c !== f) c.remove();
+    f.style.visibility = 'visible';
+  };
+  addEventListener('message', (e) => { if (pending && e.source === pending.contentWindow && e.data === 'wg-ok') show(pending); });
+  const render = () => {
+    timer = 0;
+    const p = patch(window.wgCode(latest));
+    if (!p || p[0] === shown || p[0] === pending?.srcdoc) return;
+    if (pending && p[1]) { timer = setTimeout(render, 1000); return; }  // let the in-flight partial finish first
+    pending?.remove();
+    last = performance.now();
+    const f = pending = document.createElement('iframe');
+    f.sandbox = 'allow-scripts';
+    f.style.cssText = 'position:absolute;inset:0;width:100%;height:100%;border:0;visibility:hidden';
+    if (!p[1]) f.onload = () => show(f);  // final scene: show it even if it errors
+    setTimeout(() => { if (pending === f) { f.remove(); pending = null; } }, 15000);
+    f.srcdoc = p[0];
+    box().append(f);
+  };
+  return (raw) => {
+    latest = raw;
+    if (!raw) { clearTimeout(timer); timer = 0; pending = null; shown = ''; box().replaceChildren(); return; }
+    timer ||= setTimeout(render, Math.max(0, 5000 - (performance.now() - last)));
+  };
+})();
+"""
+
+# Same prompt/extract flow as the server path, but llama.cpp runs in the browser via wllama.
 WEBGPU_JS = r"""
 async (prompt, thinking, maxTokens, temperature, topP, topK) => {
+  __LIVE__
   const st = document.querySelector('#wg-status');
-  const say = (h) => { st.innerHTML = h; };
   const esc = (s) => s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/"/g, '&quot;');
+  const t0 = performance.now(), sec = (a, b = performance.now()) => ((b - a) / 1000).toFixed(1);
+  const S = window.wgS = { phase: 'Loading wllama runtime', tail: '' };
+  const cpu = navigator.gpu ? '' : '<br><b>No WebGPU in this browser, running on CPU (very slow).</b> Use Chrome/Edge 113+.';
+  let k = 0, out = '', n = 0, tLoad, tFirst;
+  const tick = setInterval(() => {
+    st.innerHTML = `<b>${'⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'[k++ % 10]} ${S.phase}</b> · ${sec(t0)}s elapsed${cpu}`
+      + (S.tail && `<pre style="white-space:pre-wrap;max-height:200px;overflow:auto">${esc(S.tail)}</pre>`);
+  }, 100);
+  window.wgLive('');
   try {
     window.wg ??= (async () => {
       const V = 'https://cdn.jsdelivr.net/npm/@wllama/wllama@3.8.1/esm/';
@@ -116,58 +198,66 @@ async (prompt, thinking, maxTokens, temperature, topP, topK) => {
       const w = new Wllama({ default: V + 'wasm/wllama.wasm' });
       const load = () => w.loadModelFromHF({ repo: '__REPO__', quant: 'Q4_K_M' }, {
         n_ctx: 18432, n_gpu_layers: 999,
-        progressCallback: ({ loaded, total }) => say(
-          `Downloading model (cached after first run): ${(loaded / 1e9).toFixed(2)} / ${(total / 1e9).toFixed(2)} GB`),
+        progressCallback: ({ loaded, total }) => { window.wgS.phase = loaded < total
+          ? `Downloading model (cached after first run): ${(loaded / 1e9).toFixed(2)} / ${(total / 1e9).toFixed(2)} GB`
+          : 'Loading model onto GPU'; },
       });
       // an interrupted download leaves a partial cache that wllama won't repair; wipe it and retry once
-      try { await load(); } catch (e) { say('Model cache incomplete, re-downloading...'); await w.cacheManager.clear(); await load(); }
+      try { await load(); } catch (e) { window.wgS.phase = 'Model cache incomplete, re-downloading'; await w.cacheManager.clear(); await load(); }
       return w;
     })().catch((e) => { window.wg = null; throw e; });
-    if (!navigator.gpu) say('<b>No WebGPU in this browser, falling back to CPU (very slow).</b> Use Chrome/Edge 113+.');
     const w = await window.wg;
+    tLoad = performance.now();
+    S.phase = 'Reading prompt (prefill)';
     const text = `<|im_start|>system\n__SYSTEM__<|im_end|>\n<|im_start|>user\n${prompt}<|im_end|>\n<|im_start|>assistant\n`
       + (thinking ? '<think>\n' : '__NO_THINK__');
-    let out = thinking ? '<think>\n' : '```html\n', n = 0;
-    const t0 = performance.now();
+    out = thinking ? '<think>\n' : '```html\n';
     const it = await w.createCompletion({ prompt: text, max_tokens: maxTokens, temperature, top_p: topP, top_k: topK,
                                           penalty_repeat: 1.0, stop: ['<|im_end|>'], stream: true });
     for await (const c of it) {
       out += c.choices[0].text;
-      if (++n % 8 === 0) say(`Generating in your browser: ${n} tokens, ${(n * 1000 / (performance.now() - t0)).toFixed(1)} tok/s`
-                             + `<pre style="white-space:pre-wrap;max-height:200px;overflow:auto">${esc(out.slice(-800))}</pre>`);
+      tFirst ??= performance.now();
+      S.phase = `Generating in your browser: ${++n} tokens, ${(n * 1000 / (performance.now() - tFirst + 1)).toFixed(1)} tok/s`;
+      S.tail = out.slice(-800);
+      window.wgLive(out);
     }
-    const body = out.replace(/<think>[\s\S]*?(<\/think>|$)/g, '');
-    let m = body.match(/```html\s*([\s\S]*?)(```|$)/);
-    const code = m ? m[1].trim() : ((m = body.match(/<!DOCTYPE html[\s\S]*?(<\/html>|$)/i)) ? m[0] : body.trim());
-    const url = URL.createObjectURL(new Blob([code], { type: 'text/html' }));
-    say(`Done: ${n} tokens. <a href="${url}" download="scene.html">Download scene.html</a>`);
-    return [out, `<iframe sandbox="allow-scripts" srcdoc="${esc(code)}" `
-                 + 'style="width:100%;height:640px;border:0;border-radius:8px;background:#000"></iframe>'];
+    const tEnd = performance.now();
+    tFirst ??= tEnd;
+    const msg = `Done in ${sec(t0, tEnd)}s: model load ${sec(t0, tLoad)}s, prefill ${sec(tLoad, tFirst)}s, `
+      + `${n} tokens in ${sec(tFirst, tEnd)}s (${(n * 1000 / (tEnd - tFirst + 1)).toFixed(1)} tok/s)`;
+    console.log('[3D World Generator]', msg);
+    const url = URL.createObjectURL(new Blob([window.wgCode(out)], { type: 'text/html' }));
+    clearInterval(tick);
+    st.innerHTML = `${msg}${cpu}<br><a href="${url}" download="scene.html">Download scene.html</a>`;
   } catch (e) {
-    say(`<b>WebGPU error:</b> ${esc(String(e))}. Try the ZeroGPU backend.`);
-    return ['', ''];
+    clearInterval(tick);
+    st.innerHTML = `<b>WebGPU error:</b> ${esc(String(e))}. Try the ZeroGPU backend.`;
   }
+  return [out];
 }
-""".replace("__REPO__", "ReallyFloppyPenguin/Qwen3.5-4B-3D-Coder-GGUF-split").replace(
+""".replace("__LIVE__", LIVE_JS).replace("__REPO__", "ReallyFloppyPenguin/Qwen3.5-4B-3D-Coder-GGUF-split").replace(
     "__SYSTEM__", SYSTEM_PROMPT.replace("\n", "\\n")
 ).replace("__NO_THINK__", NO_THINK.replace("\n", "\\n").replace("`", "\\`"))
 
 
 def run(prompt, thinking, max_tokens, temperature, top_p, top_k):
-    raw = ""
-    for raw in generate(prompt, thinking, max_tokens, temperature, top_p, top_k):
-        yield raw, gr.update(), None
-    code = extract_html(raw)
+    t0 = time.perf_counter()
+    elapsed = lambda: f"{time.perf_counter() - t0:.0f}s elapsed"  # noqa: E731
+    yield "", None, f"<b>⏳ Waiting for ZeroGPU, loading model</b> · {elapsed()}"
+    raw, info = "", ""
+    for raw, info in generate(prompt, thinking, max_tokens, temperature, top_p, top_k):
+        yield raw, None, f"<b>⏳ Generating on ZeroGPU: {info}</b> · {elapsed()}"
     with open("scene.html", "w", encoding="utf-8") as f:
-        f.write(code)
-    yield raw, iframe(code), "scene.html"
+        f.write(extract_html(raw))
+    yield raw, "scene.html", f"Done in {time.perf_counter() - t0:.1f}s total. {info}"
 
 
 with gr.Blocks(title="3D World Generator (WebGPU)") as demo:
     gr.Markdown(
         "# 3D World Generator (WebGPU)\nDescribe a world; "
         "[ccore/Qwen3.5-4B-3D-Coder](https://huggingface.co/ccore/Qwen3.5-4B-3D-Coder) writes a self-contained "
-        "Three.js scene. Drag/scroll in the preview to move the camera (controls are model-written, so they vary).\n\n"
+        "Three.js scene. Drag/scroll in the preview to move the camera (controls are model-written, so they vary). "
+        "The preview rebuilds every few seconds while the scene is being written.\n\n"
         "**WebGPU** runs the model in your browser (2.7 GB download once, then cached; needs Chrome/Edge 113+). "
         "**ZeroGPU** runs it on the server."
     )
@@ -186,13 +276,19 @@ with gr.Blocks(title="3D World Generator (WebGPU)") as demo:
                 top_k = gr.Slider(1, 100, value=20, step=1, label="Top K")
             gr.Examples(EXAMPLES, inputs=prompt, example_labels=EXAMPLE_LABELS, label="Examples from the model card")
         with gr.Column(scale=2):
-            preview = gr.HTML(label="Preview")
+            # filled client-side by window.wgLive for both backends
+            gr.HTML(
+                '<div id="wg-live" style="position:relative;height:640px;border-radius:8px;overflow:hidden;'
+                'background:#000"></div>'
+            )
             download = gr.File(label="Download scene.html")
             with gr.Accordion("Raw model output", open=False):
                 code = gr.Code(language="html", label="Output")
     inputs = [prompt, thinking, max_tokens, temperature, top_p, top_k]
-    btn.click(run, inputs, [code, preview, download])
-    web_btn.click(None, inputs, [code, preview], js=WEBGPU_JS)
+    btn.click(None, js="() => {" + LIVE_JS + "window.wgLive('');}")
+    btn.click(run, inputs, [code, download, status])
+    code.change(None, code, js="(c) => {" + LIVE_JS + "window.wgLive(c);}")
+    web_btn.click(None, inputs, [code], js=WEBGPU_JS)
     backend.change(lambda b: (gr.update(visible=b == GPU), gr.update(visible=b == WEB)), backend, [btn, web_btn])
 
 demo.queue().launch()