Running Stable Diffusion in Browser with JavaScript & WebGPU

Why run Stable Diffusion directly in the browser?

Imagine a user uploading a text prompt and receiving a generated image within seconds—no server, no API keys, no latency caused by network hops. In 2024, WebGPU implementations in Chrome and Edge have reached 85% of desktop GPUs, making client‑side AI feasible for millions of visitors. Bringing Stable Diffusion to the browser preserves privacy, reduces cost, and opens interactive creative tools that run offline.

Prerequisites and system requirements

The target environment is a modern desktop browser that supports WebGPU (Chrome 113+, Edge 113+, or Safari 16.5+ with the experimental flag). Users need at least 4 GB of VRAM; otherwise, the model must be quantized to 8‑bit. On the development side, Node.js 20+ is recommended for bundling, and the project should use a module‑aware build tool such as Vite or Webpack.

Installing ONNX Runtime Web

ONNX Runtime Web (ORT‑Web) provides a thin JavaScript wrapper that can dispatch tensor operations to WebGPU. The installation is a single npm command, followed by a small configuration tweak to enable the WebGPU execution provider.

npm install onnxruntime-web

After installation, add the following polyfill to your entry file so that the browser knows to use WebGPU when available:

import ort from 'onnxruntime-web';
ort.env.wasm.wasmPaths = {"ort-wasm-simd.wasm":"https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort-wasm-simd.wasm"};
ort.env.webgpu.enabled = true;

Loading a Stable Diffusion model with WebGPU

Stable Diffusion is distributed as an ONNX graph that contains the UNet, VAE, and text encoder. For a browser demo, the “sd‑v1‑4‑unet.onnx” file (≈200 MB) is typically hosted on a CDN. Loading the model is asynchronous; you should display a progress bar because the download can take 5–10 seconds on a 50 Mbps connection.

const session = await ort.InferenceSession.create('https://cdn.example.com/sd-v1-4-unet.onnx', {
  executionProviders: ['webgpu'],
  graphOptimizationLevel: 'all'
});
console.log('ONNX session ready');

The same approach applies to the text encoder and VAE; keep them in separate sessions to reuse weights across multiple generations.

Generating an image from a prompt

The inference pipeline mirrors the Python reference: tokenize the prompt, run the CLIP text encoder, sample latent noise, execute the UNet for 50 denoising steps, and finally decode with the VAE. Below is a minimal JavaScript function that runs a single denoising step. Real‑world apps would loop over 50 steps and accumulate the result.

async function denoiseStep(latent, textEmbedding, step) {
  const feeds = {
    'latent_input': latent,
    'text_embeddings': textEmbedding,
    'timestep': new ort.Tensor('float32', [step])
  };
  const results = await session.run(feeds);
  return results['latent_output'];
}
// Example usage
const initLatent = ort.Tensor.fromArray(new Float32Array(64*64*4), [1,4,64,64]);
const textEmb = await encodePrompt('A futuristic city at sunset');
let latent = initLatent;
for (let i=0; i<50; i++) {
  latent = await denoiseStep(latent, textEmb, i);
}
const imageTensor = await decodeVAE(latent);
displayImage(imageTensor);

All tensors are allocated in GPU memory, so the browser never copies large buffers back to the CPU until the final image is ready. Converting the output tensor to an ImageBitmap can be done with ort.Tensor.toImageData() followed by createImageBitmap.

Performance tips and common pitfalls

1. Quantize the ONNX graph to 8‑bit integer (use onnxruntime-tools offline) – this cuts VRAM usage by ~60% and speeds up each step by 1.8× on mid‑range GPUs.
2. Reuse the same InferenceSession object; creating a new session per generation adds a 200 ms overhead.
3. Limit the image resolution to 512×512 for the first prototype; scaling to 768×768 multiplies the compute cost roughly by 2.25.
4. On mobile devices, WebGPU is still experimental; fallback to the WebAssembly execution provider with executionProviders: ['wasm'] to avoid crashes.

Conclusion

Running Stable Diffusion in the browser is no longer a research curiosity. By pairing ONNX Runtime Web with the native WebGPU API, developers can deliver privacy‑first, low‑latency AI experiences using pure JavaScript. The key steps are installing ORT‑Web, loading quantized ONNX models, and orchestrating the diffusion loop entirely on the GPU. With the code snippets above, you can prototype a text‑to‑image widget in under an hour and iterate on performance without ever touching a backend server.

Sources

• ONNX Runtime Web official documentation
• WebGPU specification (W3C)
• Stable Diffusion v1.4 model card (CompVis)

Author: Mahmut Sarıkaya — sarikayadev.com