Web Stable Diffusion is an open-source project by mlc-ai on GitHub, with the goal of bringing the Stable Diffusion model into web browsers. According to the project description, all computation is done inside the browser without server support; the project calls it "the world's first Stable Diffusion that runs entirely in the browser." The project focuses on moving some or all computation to the client side to bring potential benefits such as reducing service provider costs and enhancing personalization and privacy protection. The project also mentions that if you are interested in deploying LLM-based chatbots to the browser, you can refer to Web LLM.
The project uses machine learning compilation (MLC) technology. The overall solution is built on the open-source ecosystem, including PyTorch, Hugging Face diffusers and tokenizers, Rust, WASM, and WebGPU; the main workflow is based on Apache TVM Unity. The project uses Runway's stable diffusion v1-5 model (from the Hugging Face diffuser library), and uses TorchDynamo and Torch FX to capture key model components as IRModule in TVM. Each function in the TVM IRModule can be further transformed and generate runnable code, enabling universal deployment in any environment that supports the minimal TVM runtime (including JavaScript). The project uses TensorIR and MetaSchedule to build a solution for automatically generating optimized programs; these transformations are tuned on specific devices through the native GPU runtime and then used to generate optimized GPU shaders; the project also provides a database that records these transformations, so new builds do not need to be re-tuned. The project also builds static memory planning optimization to reuse memory across multiple layers; it uses Emscripten and TypeScript to build the TVM Web runtime that can deploy the generated modules, and utilizes a wasm port of Hugging Face's rust tokenizers library. Except for the approximately 400 lines of JavaScript application used at the end to connect the components, all parts of the workflow are done in Python.
The project provides Jupyter notebooks that walk through all stages step by step,including: explaining the key points of Web ML model deployment and how the project meets these key points, importing the stable diffusion model, optimizing the model, building the model, deploying the model locally using the native GPU runtime, and deploying the model on a web page using the WebGPU runtime.
If you want to execute these steps in the command line, you need to first install TVM Unity: you can use pip3 install mlc-ai-nightly -f https://mlc.ai/wheels Install the TVM Unity wheel, or build from source according to the TVM documentation; when building from source, after git clone use git checkout origin/unity to switch to TVM Unity. When importing, optimizing, and building the stable diffusion model, run python3 build.py; by default build.py uses apple/m2-gpu as the build target, and the CUDA target can also be specified with python3 build.py --target cuda. When deploying the model locally with the native GPU runtime, run python3 deploy.py --prompt " A photo of an astronaut riding a horse on mars. ", and you can replace the prompt with your own content, and you can specify a negative prompt with --negative-prompt "Your negative prompt".
When deploying the model on a web page using the WebGPU runtime, the last section of the walkthrough notebook, "Deploy on web", lists the complete instructions. The project also provides a plain-text explanation: first install all prerequisite dependencies. emscripten is an LLVM-based compiler that can compile C/C++ source code into WebAssembly; install the latest emsdk according to the installation instructions, and load emsdk_env.sh with source path/to/emsdk_env.sh so that emcc is accessible from PATH and the command emcc is available. Rust; wasm-pack, which helps build WebAssembly generated by Rust and is used in this project for the tokenizer; install jekyll according to the official guide, which is the tool the project uses for the website; install jekyll-remote-theme with the command gem install jekyll-remote-theme; install Chrome Canary, which is the developer version of Chrome that supports using WebGPU. You can try emcc, jekyll, and wasm-pack in the terminal respectively to verify that the installation was successful. Then prepare all dependencies required for the Web builddependencies: ./scripts/prep_deps.sh. Next, you can build the model to the WebGPU backend by running python3 build.py --target webgpu, and export the executable file to disk in the WebAssembly file format. Finally, use ./scripts/local_deploy_site.sh to set up the site. After the site is set up, you can visit localhost:8888/ in Chrome Canary to experience the demo locally on your machine. The project specifically notes that when starting Chrome Canary you need to use /Applications/Google \ Chrome \ Canary.app/Contents/MacOS/Google \ Chrome \ Canary --enable-dawn-features=disable_robustness to disable Chrome's robustness check.
The project also provides an explanation of comparisons with the native GPU runtime, limitations, and opportunities. In addition to the WebGPU runtime, the project also provides native deployment options using the local GPU runtime; these options can serve both as tools for deploying in a native environment and as a reference point for comparing native GPU driver performance with WebGPU. WebGPU works by translating WGSL (WebGPU Shading Language) shaders into native shaders, so in theory a zero gap can be achieved between the WebGPU runtime and the native environment. But the project points out that if you directly view the current demo in Chrome on Apple silicon, you will find a performance drop of about 3 times, because Chrome's WebGPU implementation inserts bounds clipping for all array index accesses, turning a[i] into a[min(i, a.size)]. Ideally, the downstream shader compiler should be able to optimize away the bounds clipping, but currently this is not the case. Once the WebGPU implementation is more mature and can check the range of index accesses and remove such clipping, this gap can be fixed. The project mentions that this problem can be bypassed by starting Chrome with a special flag; after exiting Chrome, enter the corresponding command on the command line, and you will find that execution speed is as fast as in the native GPU environment, and it is expected that this problem will be resolved as WebGPU matures.
The project believes that WebGPU is still evolving, is provided only through Chrome Canary, and may be unstable; it also has limitations, such as supporting only FP32 (the FP16 shader extension is already in the specification but not yet implemented). Stable diffusion here requires a GPU with sufficient memory (8GB). So far, the project has tested its approach only on Apple silicon. The project also mentions opportunities to support advanced optimizations such as FlashAttention and quantization to further improve system performance. The project believes these opportunities can be addressed in the near future, and individual components in the solution may still be useful, for example, only the text encoder part of the model can be deployed; the same Python-first development and general deployment workflow can also be used to bring ML models to other environments, such as new hardware or mobile scenarios; the same machine learning compilation stack is also shared with server-like use cases and can be used to optimize server workloads.
