flux-klein.js

FLUX.2 Klein 4B text-to-image and image editing in the browser. The transformer runs in hand-written WGSL on WebGPU, the text encoder and VAE on onnxruntime-web's wasm backend. Two engines behind one entry point:

| engine | weights | where | notes | | ------- | ----------------------------------------------------------- | ------------------------------------------------------------------ | ------------------------- | | desktop | int8, resident on the GPU (~4 GB) | Chrome / Edge / Safari on a laptop or desktop | fastest | | mobile | int4, streamed from disk through a small GPU ring (~0.6 GB) | phones and low-memory devices; built for the iPhone's 2 GB tab cap | one heavy thing at a time |

The package is the code plus the tokenizer (~12 MB). The weights (~5 GB desktop, ~3.3 GB mobile) are fetched from a Hugging Face repo on first use and kept in the browser's origin-private file system, so they download once per origin.

Architecture

One generate() is four stages. The transformer is the only part on the engine's own WebGPU device; the ONNX graphs run on onnxruntime-web's wasm backend (CPU threads).

| stage | runs on | desktop | mobile | | ----------------------------------------------- | --------------- | -------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- | | tokenizer (Qwen3, chat template) | JS, main thread | @huggingface/transformers | same | | text encoder (distilled 0.6B, q8, ~1 GB) | ORT wasm | ephemeral session on the main thread, weights kept in RAM | one throw-away Worker per prompt, embeddings memoised | | transformer (FLUX.2 Klein 4B, 1–8 steps) | WGSL on WebGPU | int8 resident on the GPU (~4 GB); optional LoRA (e.g. SANA-Sprint) applied at run time | int4 streamed from OPFS through a 3-slot GPU ring every step (~0.6 GB resident) | | VAE decoder (and reference encoder for editing) | ORT wasm | ephemeral session, overlapping square tiles of 320/256/128 px | one Worker per image, tiles of 256/128 px |

decoder: "tiny" swaps the tiled decoder for fal's 2.9M-parameter FLUX.2 Tiny AutoEncoder (one dynamic-shape pass, ~20× less decode work, slightly softer textures) — see Tiny decoder.

The two engines share the kernels (src/kernels.js, src/models/flux2.js) and the sampler; they differ in where the weights live (src/flux2.js resident vs src/flux2-mobile.js streaming) and in memory tactics. Attention is one fused kernel (online softmax, no scores buffer), so output size is bounded by the widest activation buffer (tokens × 27648, f16): klein.limits.maxArea is computed from the device's storage-binding limit, and iPhones are additionally capped at 512².

Why the ONNX parts stay on wasm: on desktop the page has one GPU budget (~5 GB in Chrome) and the transformer already takes most of it; on iOS the whole tab is capped at 2 GB, and ORT's WebGPU execution provider was measured slower than its wasm threads and got the tab killed during decode. A terminated Worker is also the only way to hand a wasm heap back to the OS, which is why the mobile pipeline runs every ONNX job in a fresh Worker. On iPhone (oneThing) the pipeline further keeps one heavy thing alive at a time: the transformer is dropped before the decoder runs, big allocations wait 10 s after a Worker ends, and an idle tab releases its GPU weights after 25 s (src/limits.js has every budget with the measurement that set it).

Install

npm install flux-klein.js onnxruntime-web @huggingface/transformers

onnxruntime-web (1.27+) runs the ONNX graphs; @huggingface/transformers provides the tokenizer class. Both are peer dependencies. You can pass your own copies instead (ort, PreTrainedTokenizer options) and skip installing them.

@huggingface/transformers also pulls in onnxruntime-node and sharp for its Node.js code paths. This package never executes them in the browser, but npm audit flags their current pinned versions (adm-zip < 0.6.0, sharp < 0.35.0). Until upstream bumps them, silence the report by adding to your package.json:

"overrides": {
  "adm-zip": "^0.6.0",
  "sharp": "^0.35.4"
}

(pnpm: pnpm.overrides; yarn: resolutions.)

Use

import {
  createFluxKlein,
  formatProgress,
  progressFraction,
  toImageData,
} from "flux-klein.js";

const klein = await createFluxKlein({
  onEvent: (ev) => (status.textContent = formatProgress(ev)), // downloads, uploads, denoise steps…
});

const result = await klein.generate({
  prompt: "a red bicycle leaning on a whitewashed wall, morning light",
  width: 256,
  height: 256, // multiples of 16, 128…1024 per side, area ≤ klein.limits.maxArea
  steps: 2, // Klein's shifted schedule; 2 is fast and good, 4 a little sharper (see `lora: "sprint"` for 1–2 steps)
  seed: 42,
});
canvas.getContext("2d").putImageData(toImageData(result), 0, 0);

Image editing: pass a reference (anything drawImage accepts, an ImageData, a Blob). It is cover-fitted to a square, encoded with the VAE, and the prompt describes the edit:

const ref = await klein.encodeReference(imageBitmap, {
  width: 256,
  height: 256,
}); // once
for (const prompt of prompts)
  await klein.generate({ prompt, reference: ref, width: 256, height: 256 });

For a live loop, encode the upcoming prompts first and render afterwards, so the text encoder never runs between frames (the pattern the Jetson Orin demo uses):

for (const p of nextPrompts) await klein.encodePrompt(p); // memoised
for (const p of nextPrompts)
  show(await klein.generate({ prompt: p, steps: 2 }));

SANA-Sprint LoRA (few-step model)

lora: "sprint" applies the project's SANA-Sprint distillation of Klein 4B — a rank-256 LoRA trained for 1–2 step generation, the model behind the sub-second demos — on top of the base weights at run time. It is one extra download (~0.39 GB, cached like everything else), fetched the first time generate() runs, and it is applied at whatever step count you ask for. Nothing is loaded unless you set it; the default is the stock distilled weights at every step count.

const klein = await createFluxKlein({ lora: "sprint" });
await klein.generate({ prompt, steps: 1 }); // or 2, 3, 4… — the LoRA runs at every count

Per call, generate({ lora: false }) runs the stock weights and lora: true the configured set, whatever the step count — an instant A/B, since the base never reloads.

Any lora-i8 set works the same way: lora: "my_lora_dir" names a directory under base (see WEIGHTS for the format; the exporter that produced the sprint set is distill/export_lora_web.py in flux-klein-tiny). To mix models by step count, loraDirs: { "1step": "a", "2step": "b", "4step": "" } names a set per count — an entry wins over lora for that count, and "" means the stock weights there. The same sets load on both the desktop int8 and the mobile int4 base; swapping is cheap (the base never reloads).

The fastest path is the sprint LoRA with the tiny decoder:

const klein = await createFluxKlein({ lora: "sprint", decoder: "tiny" });
show(await klein.generate({ prompt, steps: 1, width: 256, height: 256 }));

Tiny decoder (fast preview decode)

decoder: "tiny" swaps the VAE decoder for fal/FLUX.2-Tiny-AutoEncoder (Apache-2.0), a 2.9M-parameter autoencoder distilled for the FLUX.2 latent space: ~20× less decode work and an ~11 MB download instead of ~107 MB per decoder shape, at the cost of slightly softer textures (≈24 dB against the full decode on real Klein latents). Its ONNX export is dynamic-shape, so any output size decodes in one pass instead of a grid of overlapping tiles, and the latent normalisation is baked into the graph — it consumes the exact latent the full decoders do. The default (decoder: "full") is unchanged.

const klein = await createFluxKlein({ decoder: "tiny" }); // live loops, previews
// …or per call, e.g. tiny while iterating and full for the keeper:
await klein.generate({ prompt, decoder: "tiny" });

Cancel with an AbortSignal (generate({ signal })); the run stops at the next step. Call klein.destroy() to release the GPU and workers. klein.cacheInfo() / klein.clearCache() report and delete the cached weights.

Options

await createFluxKlein({
  mode: "auto", // 'desktop' | 'mobile'; auto = mobile on a mobile user agent
  base: "https://huggingface.co/radames/flux2-klein-edge-web/resolve/main", // or your own host
  device, // a GPUDevice you already created (see Limits)
  ort,
  ortUrl, // an onnxruntime-web module, or the URL of one
  wasmPaths: "/ort/", // where ORT's .wasm files are served (default: jsDelivr CDN)
  threads: 4, // ORT wasm threads (default: 8 if cross-origin isolated, else 1)
  PreTrainedTokenizer, // from @huggingface/transformers, if you import it yourself
  tokenizerUrl: "/tokenizer/", // if you copy tokenizer/ somewhere (default: inside the package)
  cacheDir: "my-app-v1", // OPFS directory (default 'flux-klein-v1')
  oneThing: true, // mobile: strict memory tactics (default: on iOS)
  promptCacheBytes: 64 << 20, // memoised prompt embeddings, LRU (3.75 MiB each)
  lora: "sprint", // opt-in LoRA at every step count: 'sprint' (SANA-Sprint) or a directory under base (default: none)
  loraDirs: { "4step": "" }, // opt-in: a LoRA (or "" = stock) per step count; wins over `lora` for that count
  decoder: "full", // 'tiny' = fal's 2.9M TAE: ~20× less decode work, slightly softer textures
  onEvent,
  log, // progress events, or sentences
});

Calls to generate() and encodeReference() run one at a time per instance: a second call made while one is running waits its turn (its signal is honoured when it starts). Size, step count and reference tokens are checked before anything is loaded or encoded.

Every option is documented in index.d.ts.

Deployment checklist

• WebGPU with 2 GB storage buffers: Chrome/Edge 121+, Safari 18+ (the mobile engine is what runs on iPhones).
• Cross-origin isolation for ORT threads: serve your page with Cross-Origin-Opener-Policy: same-origin and Cross-Origin-Embedder-Policy: require-corp. Without them everything still works, single-threaded (a few seconds slower per prompt).
• ORT wasm files: bundlers do not copy them. By default they load from jsDelivr for the installed version; to self-host, copy node_modules/onnxruntime-web/dist/*.wasm and set wasmPaths. With COEP on, a CDN needs CORS (jsDelivr sends it).
• Weight hosting: the default is a public Hugging Face repo. To self-host, mirror the repo and serve with CORS, Accept-Ranges: bytes (HTTP 206 for Range) and an ETag; the cache validates against the ETag and resumes interrupted downloads with Range requests.
• Vite: add optimizeDeps: { exclude: ['flux-klein.js', 'onnxruntime-web'] } so the workers and import.meta.url assets resolve untouched. Other bundlers: if the ONNX worker cannot be located, pass workerUrl pointing at a copy of src/ort-worker.js.
• Memory: desktop needs ~5 GB of GPU memory (the transformer and the text encoder take turns); mobile ~0.6 GB GPU + ~1 GB transient CPU during the text encoder. The mobile engine needs ~3.3 GB of OPFS quota.

Lower-level API

flux-klein.js/pipeline and flux-klein.js/mobile export the two pipelines (createPipeline, createPipelineMobile) with every knob of the engine; flux-klein.js/engine and flux-klein.js/engine-mobile the WebGPU forward itself (createEngine, loadWeights, ropeTables, …) for benches and other models; flux-klein.js/cache the OPFS cache; flux-klein.js/progress the event vocabulary; flux-klein.js/limits every tuned budget with the measurement that set it. Every entry ships type declarations (index.d.ts, types/): the pipeline and engine objects, the weight sets, and a per-stage ProgressEvent union that narrows on stage. npm run check:types compiles test/types.test-d.ts against them.

Development

npm install
npm test            # pure parts: sampling, latent plumbing, f16, quantisation, tokenizer
npm run serve       # http://127.0.0.1:8940/examples/vanilla/  (COOP/COEP + Range)
npm run serve -- --host   # bind 0.0.0.0 for other devices on the LAN (or --host=<addr> / HOST=<addr>)
npm run bench       # Deno WebGPU forward on a local clone of the model repo (KLEIN_MODEL_DIR)
# bench/attn.html (served): the fused attention kernels, every form the device admits, vs the
# old materialised path; bench/hf_kernels_deno.mjs: the same plus Hugging Face's Hub kernels
npm run pack-check  # what would be published, and that it stays under 20 MB

License

The code is MIT. The FLUX.2 Klein 4B weights are distributed separately by Black Forest Labs under their own license; using this package means fetching them under those terms. The optional tiny decoder is an ONNX export of fal/FLUX.2-Tiny-AutoEncoder (Apache-2.0).