Spaces:
Running
Running
| <html lang="en"> | |
| <head> | |
| <meta charset="utf-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1"> | |
| <title>kohra — text diffusion in the browser</title> | |
| <!-- inline fog mark (on-brand SVG favicon) — kills the favicon.ico 404 on Pages + the HF Space --> | |
| <link rel="icon" type="image/svg+xml" href="data:image/svg+xml,%3Csvg%20xmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20viewBox%3D%220%200%2064%2064%22%3E%3Crect%20width%3D%2264%22%20height%3D%2264%22%20rx%3D%2214%22%20fill%3D%22%231b1e25%22%2F%3E%3Cg%20fill%3D%22none%22%20stroke-linecap%3D%22round%22%20stroke-width%3D%225%22%3E%3Cpath%20d%3D%22M13%2023q9-6%2018%200t18%200%22%20stroke%3D%22%237aa2f7%22%20opacity%3D%22.9%22%2F%3E%3Cpath%20d%3D%22M11%2035q9-6%2018%200t18%200%22%20stroke%3D%22%239ece6a%22%20opacity%3D%22.6%22%2F%3E%3Cpath%20d%3D%22M15%2046q8-5%2016%200t14%200%22%20stroke%3D%22%237aa2f7%22%20opacity%3D%22.4%22%2F%3E%3C%2Fg%3E%3C%2Fsvg%3E"> | |
| <style> | |
| :root { --fg:#e8e6e3; --bg:#14161a; --dim:#8a8f98; --acc:#7aa2f7; --mask:#3b4261; --fresh:#9ece6a; } | |
| * { box-sizing: border-box; } | |
| body { margin:0; background:var(--bg); color:var(--fg); font:15px/1.5 ui-monospace, SFMono-Regular, Menlo, monospace; } | |
| main { max-width: 880px; margin: 0 auto; padding: 24px 16px 80px; } | |
| h1 { font-size: 18px; font-weight: 600; } | |
| h1 small { color: var(--dim); font-weight: 400; } | |
| /* discreet loader: a small status pill with a spinner (loading/ready/error) */ | |
| .status-badge { | |
| display: inline-flex; align-items: center; gap: 7px; | |
| padding: 4px 11px; border-radius: 999px; | |
| font-size: 12px; border: 1px solid transparent; | |
| transition: background .4s, border-color .4s, color .4s; | |
| } | |
| .status-badge.loading { background: rgba(122,162,247,.12); border-color: rgba(122,162,247,.35); color: var(--acc); } | |
| .status-badge.ready { background: rgba(158,206,106,.12); border-color: rgba(158,206,106,.35); color: var(--fresh); } | |
| .status-badge.error { background: rgba(247,118,142,.12); border-color: rgba(247,118,142,.4); color: #f7768e; } | |
| .spinner { | |
| width: 10px; height: 10px; flex: none; | |
| border: 2px solid currentColor; border-top-color: transparent; | |
| border-radius: 50%; animation: spin .75s linear infinite; | |
| } | |
| .status-badge:not(.loading) .spinner { display: none; } | |
| @keyframes spin { to { transform: rotate(360deg); } } | |
| textarea { width:100%; background:#1b1e25; color:var(--fg); border:1px solid #2a2e37; border-radius:8px; padding:10px; font:inherit; min-height:64px; } | |
| .row { display:flex; gap:12px; align-items:center; flex-wrap:wrap; margin:10px 0; } | |
| label { color:var(--dim); font-size:13px; } | |
| input[type=number] { width:72px; background:#1b1e25; color:var(--fg); border:1px solid #2a2e37; border-radius:6px; padding:4px 6px; font:inherit; } | |
| select { background:#1b1e25; color:var(--fg); border:1px solid #2a2e37; border-radius:6px; padding:4px 6px; font:inherit; cursor:pointer; } | |
| button { background:var(--acc); color:#10131a; border:0; border-radius:8px; padding:8px 18px; font:inherit; font-weight:600; cursor:pointer; } | |
| button:disabled { opacity:.4; cursor:default; } | |
| #canvas { background:#1b1e25; border:1px solid #2a2e37; border-radius:8px; padding:14px; margin-top:14px; min-height:120px; white-space:pre-wrap; word-break:break-word; } | |
| .tok-mask { color:var(--mask); } | |
| .tok-fresh { color:var(--fresh); } | |
| #perf { color:var(--dim); font-size:13px; margin-top:8px; } | |
| /* bottom pill: repo + portfolio links (LocalMind-style, centered, fixed) */ | |
| .footer-nav { | |
| position: fixed; bottom: 12px; left: 50%; transform: translateX(-50%); | |
| display: flex; gap: 20px; white-space: nowrap; | |
| background: rgba(27,30,37,.9); border: 1px solid #2a2e37; | |
| padding: 6px 16px; border-radius: 999px; | |
| box-shadow: 0 2px 12px rgba(0,0,0,.35); backdrop-filter: blur(6px); | |
| font-size: 12px; | |
| } | |
| .footer-nav a { color: var(--acc); text-decoration: none; transition: opacity .2s; } | |
| .footer-nav a:hover { opacity: .65; } | |
| </style> | |
| </head> | |
| <body> | |
| <main> | |
| <h1>kohra <small>— masked-diffusion text generation, fully in your browser (WebGPU)</small></h1> | |
| <p><span id="status" class="status-badge loading"><span class="spinner"></span><span id="statusText">Loading…</span></span></p> | |
| <textarea id="prompt">Lily runs 12 km/h for 4 hours. How far in 8 hours?</textarea> | |
| <div class="row"> | |
| <label>model <select id="modelSel"> | |
| <option value="mdlm">MDLM · fp16</option> | |
| <option value="mdlm-q4">MDLM · q4</option> | |
| <option value="bd3lm">BD3LM · fp16 · block diffusion</option> | |
| <option value="bd3lm-q4">BD3LM · q4 · block diffusion</option> | |
| </select></label> | |
| <label>new tokens <input id="maxNew" type="number" value="128" step="32" min="16"></label> | |
| <label>steps <input id="steps" type="number" value="128" step="16" min="4"></label> | |
| <label>block <input id="blockSize" type="number" value="32" step="8" min="8"></label> | |
| <label>temp <input id="temp" type="number" value="0" step="0.1" min="0"></label> | |
| <label title="Fast-dLLM: unmask all tokens above this confidence per step (0 = off, use steps)">conf≥ <input id="threshold" type="number" value="0" step="0.05" min="0" max="1"></label> | |
| <button id="go" disabled>Generate</button> | |
| </div> | |
| <div id="canvas"></div> | |
| <div id="perf"></div> | |
| </main> | |
| <!-- bottom pill: portfolio + source --> | |
| <div class="footer-nav"> | |
| <a href="https://naklitechie.github.io/" target="_blank" rel="noopener">⬅ More Projects</a> | |
| <a href="https://github.com/NakliTechie/kohra" target="_blank" rel="noopener">⌨ Source</a> | |
| </div> | |
| <script type="module"> | |
| // Thin DOM harness over the kohra.js library — all the diffusion logic lives there. | |
| // Model + tokenizer load from Hugging Face by default, so this page works as-is on | |
| // GitHub Pages / a static HF Space. For local dev against a local export, override: | |
| // ?model=models/qwen3-0.6b-mdlm-onnx/model_fp16_fused.onnx | |
| import { DiffusionLM } from './kohra.js'; | |
| // fused fp16 (G1d): RMSNorm fused to SimplifiedLayerNormalization so WebGPU's native | |
| // fp16 kernels don't overflow the decomposed Pow(x,2) -> all-zero. ~1.4GB, ~9.8 tok/s. | |
| const $ = (id) => document.getElementById(id); | |
| // state: 'loading' (spinner) | 'ready' | 'error' | |
| const status = (t, state = 'loading') => { | |
| $('statusText').textContent = t; | |
| $('status').className = 'status-badge ' + state; | |
| }; | |
| // Model registry — each entry is a self-contained HF repo (graph + tokenizer + bundled | |
| // kohra.js). bd3lm (block diffusion) drives the block-causal sampler; q4 variants are a | |
| // smaller download. The dropdown picks one (reloads with ?arch=<key>); ?model= (+ optional | |
| // ?arch=, ?tokenizer=, ?ort=, ?opt=) overrides for local/dev. MDLM-fp16 stays the default. | |
| const REPO = { | |
| mdlm: 'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX/resolve/main', | |
| bd3lm: 'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-bd3lm-ONNX/resolve/main', | |
| }; | |
| const TOKENIZER = { | |
| mdlm: 'naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX', | |
| bd3lm: 'naklitechie/Qwen3-0.6B-diffusion-bd3lm-ONNX', | |
| }; | |
| // q4 (RTN MatMulNBits) needs the ORT-web dev build with the WebGPU q4 fix + full graph opt. | |
| const ORT_Q4 = '1.26.0-dev.20260416-b7804b056c'; | |
| const REGISTRY = { | |
| 'mdlm': { label: 'MDLM · fp16', repo: 'mdlm', file: 'model_fp16_fused.onnx', blockCausal: false, q4: false }, | |
| 'mdlm-q4': { label: 'MDLM · q4 (smaller)', repo: 'mdlm', file: 'model_q4f16_rtn_sym.onnx', blockCausal: false, q4: true }, | |
| 'bd3lm': { label: 'BD3LM · fp16 · block diff', repo: 'bd3lm', file: 'model_fp16_fused.onnx', blockCausal: true, q4: false }, | |
| 'bd3lm-q4': { label: 'BD3LM · q4 · block diff', repo: 'bd3lm', file: 'model_q4f16_rtn_sym.onnx', blockCausal: true, q4: true }, | |
| }; | |
| const params = new URLSearchParams(location.search); | |
| const modelOverride = params.get('model'); | |
| const reqArch = params.get('arch') || (modelOverride?.includes('bd3lm') ? 'bd3lm' : 'mdlm'); | |
| const arch = REGISTRY[reqArch] ? reqArch : 'mdlm'; // unknown key → safe default | |
| const preset = REGISTRY[arch]; | |
| const modelUrl = modelOverride ?? `${REPO[preset.repo]}/onnx/${preset.file}`; | |
| const tokenizerId = params.get('tokenizer') ?? TOKENIZER[preset.repo]; | |
| const blockCausal = preset.blockCausal; | |
| const ortVersion = params.get('ort') || (preset.q4 ? ORT_Q4 : undefined); | |
| const graphOptimizationLevel = params.get('opt') || (preset.q4 ? 'all' : 'disabled'); | |
| let lm; | |
| // Visualize the "fog lifting": masked positions as ▒, freshly committed tokens in green. | |
| function renderCanvas({ x, P, fresh }) { | |
| const el = $('canvas'); el.textContent = ''; | |
| for (let i = P; i < x.length; i++) { | |
| const span = document.createElement('span'); | |
| if (x[i] === lm.maskId) { span.className = 'tok-mask'; span.textContent = '▒'; } | |
| else { | |
| span.textContent = lm.tokenizer.decode([Number(x[i])]); | |
| if (fresh.has(i)) span.className = 'tok-fresh'; | |
| } | |
| el.appendChild(span); | |
| } | |
| } | |
| async function generate() { | |
| $('go').disabled = true; | |
| try { | |
| const out = await lm.generate($('prompt').value, { | |
| maxNewTokens: +$('maxNew').value, | |
| steps: +$('steps').value, | |
| blockSize: +$('blockSize').value, | |
| temperature: +$('temp').value, | |
| threshold: +$('threshold').value || null, // Fast-dLLM confidence-threshold decoding | |
| blockCausal, // bd3lm: static block-causal attention mask | |
| stripThink: true, | |
| onStep: (s) => { | |
| renderCanvas(s); | |
| status(`block ${s.block + 1}/${s.numBlocks} · forward ${s.forward} · ${(s.elapsedMs / 1000).toFixed(1)}s`); | |
| }, | |
| }); | |
| $('canvas').textContent = out.text; | |
| $('perf').textContent = | |
| `${out.tokens} tokens · ${out.forwards} forwards · ${out.seconds.toFixed(1)}s · ${out.tokensPerSecond.toFixed(1)} tok/s`; | |
| status('Done.', 'ready'); | |
| } catch (e) { | |
| console.error(e); | |
| status('Generation failed: ' + e.message, 'error'); | |
| } finally { | |
| $('go').disabled = false; | |
| } | |
| } | |
| async function init() { | |
| try { | |
| if (!navigator.gpu) { status('This demo needs WebGPU — try Chrome/Edge 121+.', 'error'); return; } | |
| status(`Loading model… first run downloads ${preset.q4 ? '~0.7GB' : '~1.5GB'} + compiles WebGPU shaders (cached after).`); | |
| lm = await DiffusionLM.from_pretrained({ model: modelUrl, tokenizer: tokenizerId, ortVersion, graphOptimizationLevel }); | |
| window.kohra = lm; // handy for the console / preview probing | |
| status(`Ready — ${preset.q4 ? 'q4f16' : 'fused fp16'} Qwen3-0.6B-${blockCausal ? 'BD3LM (block diffusion)' : 'MDLM'} on WebGPU. Press Generate.`, 'ready'); | |
| $('go').disabled = false; | |
| } catch (e) { | |
| console.error(e); | |
| status('FAILED: ' + e.message, 'error'); | |
| } | |
| } | |
| addEventListener('error', e => status('Uncaught: ' + (e.error?.message ?? e.message), 'error')); | |
| addEventListener('unhandledrejection', e => status('Rejection: ' + (e.reason?.message ?? e.reason), 'error')); | |
| $('go').addEventListener('click', generate); | |
| // Model picker: reflect the active arch, and reload with the chosen preset on change | |
| // (drop any ?model override so the registry URL is used). | |
| $('modelSel').value = arch; | |
| $('modelSel').addEventListener('change', () => { location.search = '?arch=' + $('modelSel').value; }); | |
| init(); | |
| </script> | |
| </body> | |
| </html> | |