This project enables Stable Diffusion image generation directly within web browsers, eliminating the need for server-side infrastructure. It targets web developers and users seeking privacy-preserving, cost-effective AI image generation capabilities, leveraging client-side hardware for computation.

How It Works

The project utilizes Apache TVM Unity, a machine learning compilation framework, to convert Stable Diffusion models (specifically Runway's v1-5) into a WebAssembly runtime. It employs TorchDynamo and Torch FX for model capture, TVM's TensorIR and MetaSchedule for shader optimization, and Emscripten for WebAssembly compilation. This approach allows models to run natively on client GPUs via WebGPU, with optimized shaders and static memory planning for efficient browser execution.

Quick Start & Requirements

• Install TVM Unity: pip3 install mlc-ai-nightly -f https://mlc.ai/wheels or build from source.
• Prerequisites: Emscripten, Rust, wasm-pack, Jekyll, Chrome Canary.
• Build & Deploy:
  • python3 build.py (for local GPU) or python3 build.py --target webgpu (for WebGPU).
  • ./scripts/local_deploy_site.sh for web deployment.
• Demo: Run localhost:8888 in Chrome Canary with specific flags.
• Documentation: Walkthrough Notebook

Highlighted Details

• World's first Stable Diffusion implementation running entirely in the browser.
• Leverages WebGPU for native GPU acceleration in web environments.
• Python-first development workflow for model optimization and deployment.
• Supports local GPU runtime deployment for comparison and native use.

Maintenance & Community

• Actively developed within the Apache TVM ecosystem.
• Benefits from contributions from the Hugging Face, PyTorch, WebAssembly, and Rust communities.
• Related project: Web LLM.

Licensing & Compatibility

• Primarily licensed under Apache-2.0 (TVM Unity).
• Compatible with commercial use and closed-source linking, as the core runtime is WebAssembly.

Limitations & Caveats

• WebGPU is still evolving and may require specific browser versions (Chrome Canary) and flags (--enable-dawn-features=disable_robustness) for optimal performance.
• Requires a GPU with at least 8GB of RAM.
• Performance on Apple Silicon via WebGPU shows a ~3x degradation compared to native GPU without specific flags due to robustness checks.
• Only tested on Apple Silicon; broader platform compatibility is not guaranteed.