कोहरा means fog. Generation starts as a fully masked canvas and clears, pass by pass, into text.

Install

| Path | How | |---|---| | Try it, nothing to install | Open the live demo (also on Hugging Face Spaces) | | Use it in your page | cp kohra.js your-app/ from this repo | | Import from a CDN | import { pipeline } from 'https://cdn.jsdelivr.net/npm/[email protected]/kohra.js' | | npm | npm install kohra.js |

The first load pulls the model (~1.5 GB fp16, ~0.7 GB q4) and compiles WebGPU shaders. Your browser caches both, so later loads start in seconds. kohra.js loads onnxruntime-web from a CDN and the tokenizer from Hugging Face itself, so there is no build step. A complete page:

<button id="go">Generate</button><pre id="out"></pre>
<script type="module">
import { pipeline } from './kohra.js';
const generate = await pipeline('text-diffusion', {
  model: 'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX/resolve/main/onnx/model_fp16_fused.onnx',
  tokenizer: 'naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX',
});
document.getElementById('go').onclick = async () => {
  const { text } = await generate('Explain WebGPU in one sentence.', { maxNewTokens: 128, steps: 128, stripThink: true });
  document.getElementById('out').textContent = text;
};
</script>

Serve it over https or localhost, because WebGPU needs a secure context: python3 -m http.server 8000. No config file and no API key. kohra imports onnxruntime-web and transformers.js from jsDelivr at runtime, so a bundler must leave those two https: imports external.

Why

You want to try text diffusion, the alternative to token-by-token generation, and every runtime you reach for is autoregressive. Transformers.js, onnxruntime-web and WebLLM all ship only a left-to-right decode loop. The published diffusion models need a Python server and a datacenter GPU.

kohra is the missing piece for the browser. It runs a few hundred lines of JS sampler over raw ONNX forward passes, plus ONNX exports of two small diffusion LMs that run on WebGPU. On a 121-token math answer, BD3LM with the KV cache now beats autoregressive decoding of the same 0.6B base: 3.0 s against 3.5–4.0 s. Diffusion always denoises its full canvas, so short answers still favour autoregressive decoding (benchmark in KOHRA.md). kohra is for running, measuring and building on text diffusion in the browser.

Watch the fog lift

DiffusionLM exposes the sampler and a callback for each step. x is the full token canvas, where masked positions equal lm.maskId. fresh holds the positions revealed in this step: colour them to animate the fog clearing.

import { DiffusionLM } from './kohra.js';
const lm = await DiffusionLM.from_pretrained({ model, tokenizer });
const out = await lm.generate(prompt, {
  maxNewTokens: 128, steps: 128, blockSize: 32,
  temperature: 0,      // 0 = argmax; >0 = Gumbel sampling
  threshold: 0.8,      // Fast-dLLM: reveal every position above this confidence; null = fixed steps
  onStep: ({ x, P, fresh, forward }) => render(x, P, fresh),
});
// out: { text, tokenIds, tokens, forwards, seconds, tokensPerSecond }

threshold: 0.8 is the one speed setting to turn on, and the demo's default. Across 3 prompts it cuts MDLM from 384 forward passes to 194 and BD3LM from 328 to 137, with the same math answer and fluent text. At 0.7 MDLM breaks on the math prompt. index.html is the live demo, built only on this API.

Pick a model

You get two models, each in fp16 and q4. The demo's model picker switches between all four.

| Model | Attention | Call with | Notes | |---|---|---|---| | MDLM | bidirectional | default | the original masked-diffusion checkpoint | | BD3LM | block-causal | blockCausal: true | higher scores: GSM8K 46.3 vs 29.3, HumanEval 46.3 vs 30.5; KV-cache graph ~2× faster |

For BD3LM, load the KV-cache graph. Each step then runs only the current block against a cache of finished blocks. Output is identical to the plain graph and about twice as fast:

const HF = 'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-bd3lm-ONNX/resolve/main/onnx';
const lm = await DiffusionLM.from_pretrained({
  model: `${HF}/model_kv_fp16_fused.onnx`,          // or model_kv_q4f16_rtn_sym.onnx + graphOptimizationLevel: 'all'
  tokenizer: 'naklitechie/Qwen3-0.6B-diffusion-bd3lm-ONNX', kvCache: true,   // cache stays on the GPU
});
const { text } = await lm.generate(prompt, { blockCausal: true, threshold: 0.8 });

The plain q4 graphs are onnx/model_q4f16_rtn_sym.onnx; load them with graphOptimizationLevel: 'all'. They run on the stable onnxruntime-web 1.30.0 that kohra.js loads by default. q4 halves the download, and BD3LM q4 slips on arithmetic that fp16 gets right. To host your own export, serve the .onnx and .onnx.data side by side with permissive CORS. kohra finds the external-data file without configuration.

Commands

The Python export scripts need uv venv .venv && uv pip install --python .venv/bin/python torch "transformers==4.57.0" onnx onnxruntime onnxscript numpy huggingface_hub, plus git clone --depth 1 https://github.com/ZHZisZZ/dllm vendor/dllm for the BD3LM model code.

python3 -m http.server 8791                              # serve the demo and harnesses at localhost:8791
open http://localhost:8791/?arch=bd3lm                   # demo on a chosen model: mdlm | mdlm-q4 | bd3lm | bd3lm-q4
open http://localhost:8791/web/bench.html?mode=ar        # autoregressive Qwen3-0.6B baseline, same browser
.venv/bin/python scripts/export_onnx.py --fp16           # export MDLM to ONNX + parity check (export_bd3lm.py for BD3LM)
.venv/bin/python scripts/optimize_onnx.py                # fuse RMSNorm, then fp16 (required for WebGPU)
.venv/bin/python scripts/export_bd3lm_kv.py              # BD3LM block KV-cache graph: export, fuse, fp16, parity at each stage
.venv/bin/python scripts/push_to_hf.py --model mdlm --stage meta   # publish graph + tokenizer + card + kohra.js

Verify it yourself

npm test                                                 # 13 sampler tests on a fake ONNX session (Node >= 22, no install)
open http://localhost:8791/web/probe.html?model=<url>    # WebGPU logits vs the fp32 CPU ground truth
open http://localhost:8791/web/bench.html?mode=diff      # forwards, seconds and text per configuration (add &arch=bd3lm)

The probe fails a graph that returns non-finite or all-zero logits on WebGPU, or whose argmax disagrees with the fp32 reference. That is how the unfused fp16 graph was caught. The bench page produced the AR-vs-diffusion table in KOHRA.md, re-run on 2026-10-08 in thermal state nominal; the threshold and KV-cache tables come from the same kohra.js API run in that page. npm test runs the real kohra.js sampler against a fake session. It fails on a wrong commit order, a broken block-causal mask, a missing EOS trim, or a loop that stops making progress.

License

Apache-2.0. The published models carry the licences of their dllm-collection upstreams.

KOHRA.md (why, gate ladder, benchmark, gotchas) · reference/MDLM-algorithm.md (export recipe and WebGPU forensics)