Web 稳定扩散

本项目将稳定扩散模型引入到网页浏览器中。所有计算都在浏览器内部完成，无需服务器支持。 据我们所知，这是全球首个完全在浏览器上运行的稳定扩散模型。请访问我们的演示页面体验。

如果您对将基于 LLM 的聊天机器人部署到浏览器感兴趣，也欢迎查看 Web LLM。

近年来，AI 模型取得了令人瞩目的进展。得益于开源社区的努力，开发者现在可以轻松地组合各种开源模型，完成许多令人惊叹的任务。稳定扩散模型能够根据文本输入自动生成逼真的图像以及各种风格的艺术作品。然而，这些模型通常体积庞大、计算密集，因此在开发基于它们的 Web 应用时，我们必须将所有的计算请求转发到 GPU 服务器上处理。此外，大多数工作负载还需要在特定类型的 GPU 上运行，而这些 GPU 上往往已经预装了流行的深度学习框架。

本项目旨在改变这一现状，为生态系统带来更多多样性。将部分（或全部）计算转移到客户端有许多理由。这样做不仅能够降低服务提供商的成本，还能提升个性化体验和隐私保护水平。个人电脑（甚至移动设备）的发展正朝着支持这种可能性的方向迈进——客户端设备的功能越来越强大。

为这些应用构建专用的客户端应用程序是一种选择（我们也支持这种方式），但如果我们只需打开浏览器，就能直接将 AI 原生地集成到你的浏览器标签页中，岂不是更加令人兴奋？目前，生态系统已经具备了一定的基础：WebAssembly 允许我们将更多底层运行时移植到 Web 平台上。为了解决计算问题，WebGPU 技术也在不断成熟，使得在浏览器中进行原生 GPU 计算成为可能。

我们正目睹客户端侧所需的各项要素逐渐就位，无论是硬件层面还是浏览器生态方面。尽管如此，我们仍需克服诸多挑战，例如：

• 我们需要在没有相关 GPU 加速 Python 框架的情况下运行这些模型。
• 大多数 AI 框架高度依赖由硬件厂商维护的优化计算库。我们需要从零开始构建。为了获得最佳性能，我们可能还需要针对不同的客户端环境生成不同的版本。
• 必须精心规划内存使用，以确保模型能够完整加载到内存中。

我们并不希望只为单一模型实现这一点。相反，我们希望提供一种可重复、可扩展、可组合的工作流程，使任何人都能在“Python 优先”的环境中轻松开发和优化这些模型，并将其普遍部署到包括 Web 在内的任何平台上。

开始使用

我们提供了一个 Jupyter 笔记本，它将引导您完成整个流程，包括：

• 详细阐述 Web ML 模型部署的关键点，以及我们如何满足这些要求；
• 导入稳定扩散模型；
• 对模型进行优化；
• 构建模型；
• 使用原生 GPU 运行时在本地部署模型；以及
• 使用 WebGPU 运行时在 Web 上部署模型。

如果您希望通过命令行完成这些步骤，请按照以下命令操作：

命令

• 安装 TVM Unity。您可以：

  • 使用 pip3 install mlc-ai-nightly -f https://mlc.ai/wheels 安装 TVM Unity 的 wheel 包，或者
  • 参照 TVM 官方文档 从源码编译。请注意，在 git clone 之后，请使用 git checkout origin/unity 切换到 TVM Unity 分支。

• 导入、优化并构建稳定扩散模型：

  python3 build.py
  

  默认情况下，build.py 会将 apple/m2-gpu 作为构建目标。您也可以通过以下命令指定 CUDA 目标：

  python3 build.py --target cuda
  

• 使用原生 GPU 运行时在本地部署模型：

  python3 deploy.py --prompt "A photo of an astronaut riding a horse on mars."
  

  您可以替换提示词为自己想要的内容，还可以选择性地使用 --negative-prompt "Your negative prompt" 来指定负面提示词。

• 要使用 WebGPU 运行时在 Web 上部署模型，演示笔记本 的最后一节“Deploy on web”提供了完整的说明，您也可以参考此处的简明指令：

  说明

  首先，让我们安装所有必要的依赖项：

  1. emscripten。这是一个基于 LLVM 的编译器，可以将 C/C++ 源代码编译成 WebAssembly。
     • 请按照 安装说明 安装最新版本的 emsdk。
     • 执行 source path/to/emsdk_env.sh 以激活 emsdk_env.sh，从而使 emcc 能够通过 PATH 访问，并确保 emcc 命令正常运行。
  2. Rust。
  3. wasm-pack。它用于构建由 Rust 生成的 WebAssembly，而在本例中则用于分词器。
  4. 按照 官方指南 安装 Jekyll。这是我们用来搭建网站的工具包。
  5. 使用以下命令安装 jekyll-remote-theme：

     gem install jekyll-remote-theme
     

  6. 安装 Chrome Canary。它是 Chrome 的开发者版本，支持 WebGPU 功能。

  您可以通过分别在终端中尝试运行 emcc、jekyll 和 wasm-pack 来验证是否成功安装。

  接下来，准备 Web 构建所需的所有依赖项：

  ./scripts/prep_deps.sh
  

  现在，我们可以将模型构建为 WebGPU 后端，并以 WebAssembly 文件格式导出可执行文件：

  python3 build.py --target webgpu
  

  最后一步是设置站点：

  ./scripts/local_deploy_site.sh
  

  站点搭建完成后，您可以在 Chrome Canary 中访问 localhost:8888/，在本地机器上试用演示。请务必使用以下命令启动 Chrome Canary，以关闭 Chrome 的健壮性检查：

  /Applications/Google\ Chrome\ Canary.app/Contents/MacOS/Google\ Chrome\ Canary --enable-dawn-features=disable_robustness
  

怎么做？

这里的关键技术是机器学习编译（MLC）。我们的解决方案建立在开源生态系统的肩膀之上，包括 PyTorch、Hugging Face diffusers 和 tokenizers、Rust、Wasm 以及 WebGPU。主要流程基于 Apache TVM Unity 构建，这是 Apache TVM 中一项令人振奋的持续开发工作。

• 我们从 Hugging Face 的 diffuser 库中获取 Runway 的稳定扩散 v1-5 模型。
• 我们使用 TorchDynamo 和 Torch FX 将关键模型组件捕获为 TVM 中的 IRModule。
• TVM 的 IRModule 中的每个函数都可以进一步转换并生成可运行的代码，这些代码可以在任何支持最小 TVM 运行时的环境中通用部署（JavaScript 就是其中之一）。
• 我们使用 TensorIR 和 MetaSchedule 来构建自动化方案以生成优化后的程序。这些转换会在特定设备上通过原生 GPU 运行时进行调优，然后用于生成优化的 GPU 着色器。我们提供一个记录这些转换的数据库，以便在无需再次调优的情况下完成新的构建。
• 我们构建了静态内存规划优化，以实现跨多层的内存复用。
• 我们使用 Emscripten 和 TypeScript 构建了一个 TVM Web 运行时，用于部署生成的模块。
• 我们还利用了来自 Hugging Face 的 rust tokenizers 库 的 Wasm 移植版。

这个工作流的所有部分都是用 Python 完成的，当然，最后一部分除外——它构建了一个约 400 行的 JavaScript 应用程序来将各个部分连接起来。这也是一个有趣的交互式开发过程，可以不断引入新模型。

这一切都得益于我们所依赖的开源生态系统。特别是，我们大量使用了 TVM Unity，这是 TVM 项目中一项令人兴奋的最新进展，它支持这种以 Python 为主导的交互式 MLC 开发体验，使我们能够轻松地用 Python 组合新的优化方法，并逐步将我们的应用推向 Web 平台。TVM Unity 还提供了一种简便的方式，可以在生态系统中组合新的解决方案。例如，未来我们可以轻松地将其他 WebGPU 着色器生成器或着色器库引入到这个工作流中。

与原生 GPU 运行时的对比、局限性及机遇

除了 WebGPU 运行时之外，我们还提供了使用本地 GPU 运行时进行原生部署的选项。这些选项既可以作为在原生环境中部署的工具，也可以作为比较原生 GPU 驱动性能与 WebGPU 性能的参考点。

WebGPU 通过将 WGSL（WebGPU 着色语言）着色器转换为原生着色器来工作。因此，理论上，WebGPU 运行时与原生环境之间可以实现零差距。然而，如果我们直接使用 Chrome 在 Apple Silicon 上查看当前演示，会发现性能下降了约 3 倍。这是因为 Chrome 的 WebGPU 实现会对所有数组索引访问插入边界检查，例如将 a[i] 替换为 a[min(i, a.size)]。理想情况下，下游的着色器编译器应该能够优化掉这些边界检查，但遗憾的是，目前并未做到。随着 WebGPU 实现的逐渐成熟，能够更好地检查索引访问范围并移除此类检查后，这一差距将会得到解决。

您可以通过使用特殊标志启动 Chrome 来绕过这个问题（感谢 Dawn 开发人员提供的提示），即完全退出 Chrome，然后在命令行中输入：

/path/to/chrome-canary --enable-dawn-features=disable_robustness

这样您就会发现执行速度与原生 GPU 环境一样快。我们预计随着 WebGPU 的成熟，这个问题将会得到解决。

我们正目睹一场我们认为即将爆发的变革的曙光。WebGPU 仍在不断发展（尽管今年已接近正式发布），目前仅可通过 Chrome Canary 使用，且可能不够稳定。它仍然存在一些限制，比如仅支持 FP32（FP16 着色器扩展已在规范中提出，但尚未实现）。这里的稳定扩散模型需要配备具有足够显存的 GPU（至少 8GB）。到目前为止，我们只在 Apple Silicon 上测试过我们的解决方案。此外，还有机会支持诸如 FlashAttention 和量化等高级优化技术，以进一步提升系统性能。

这些都为当前解决方案带来了数倍的性能提升机会。我们相信其中许多问题在不久的将来都能够得到解决。即使只是该解决方案中的单个组件，也仍然具有实用价值。例如，用户可以选择仅部署模型的文本编码器部分。此外，同样的以 Python 为主导的开发和通用部署流程也可以用于将机器学习模型引入其他环境，比如新型硬件或移动设备。最后，相同的机器学习编译堆栈也可用于服务器级应用场景，从而优化服务器工作负载。

致谢

本项目得以实现，离不开以下机构的合作：

本项目之所以能够实现，完全归功于我们所依托的开源生态系统。我们要感谢 Apache TVM 社区以及 TVM Unity 项目的开发者们。同时也要感谢那些公开提供这些模型的开源机器学习社区成员，以及让这些模型易于获取的 PyTorch 和 Hugging Face 社区。此外，还要感谢 Mithril Security 对分词器 Wasm 移植工作的贡献。我们也感谢 WebAssembly、Emscripten、Rust 和 WebGPU 社区的支持。最后，特别感谢 Dawn 开发人员，他们及时解答了我们在使用 Chrome 时遇到的问题。

Web Stable Diffusion 快速上手指南

Web Stable Diffusion 是一个开创性的开源项目，它让 Stable Diffusion 模型能够完全在浏览器中运行，无需任何服务器支持。所有计算均在客户端（浏览器）通过 WebGPU 完成。

环境准备

系统要求

• 操作系统: macOS (Apple Silicon 推荐), Linux, Windows
• 浏览器: Chrome Canary (必须，用于启用 WebGPU 支持)
• 硬件: 建议配备至少 8GB 显存 的 GPU。目前主要在 Apple Silicon 上经过充分测试。

前置依赖

请确保已安装以下工具：

1. Python 环境: 建议 Python 3.8+
2. TVM Unity: 机器学习编译后端
3. Emscripten: 用于将代码编译为 WebAssembly
4. Rust & wasm-pack: 用于构建分词器 (Tokenizer)
5. Jekyll: 用于搭建本地演示网站

安装步骤

1. 安装 TVM Unity

你可以选择通过 pip 安装预编译包，或从源码构建（推荐高级用户）。

方式 A：使用 pip 安装 (推荐)

pip3 install mlc-ai-nightly -f https://mlc.ai/wheels

方式 B：从源码构建

git clone --recursive https://github.com/apache/tvm tvm
cd tvm
git checkout origin/unity
# 后续请参考 TVM 官方文档进行编译配置

2. 安装 Web 构建工具链

按照以下顺序安装必要的环境工具：

• Emscripten: 下载并安装最新版的 emsdk，然后激活环境：

  source path/to/emsdk_env.sh
  

  (验证：终端输入 emcc 应能正常响应)

• Rust: 访问 rust-lang.org 安装。

• wasm-pack:

  curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
  

• Jekyll 及相关插件:

  gem install jekyll
  gem install jekyll-remote-theme
  

  (验证：终端输入 jekyll 和 wasm-pack 应能正常响应)

• Chrome Canary: 下载并安装 Google Chrome Canary。

3. 准备依赖与构建模型

克隆本项目后，执行以下命令：

步骤 A: 准备 Web 构建依赖

./scripts/prep_deps.sh

步骤 B: 构建模型 (针对 WebGPU) 默认构建目标为 apple/m2-gpu，若需部署到网页，请指定 webgpu：

python3 build.py --target webgpu

(注：若仅需在本地原生环境测试，可使用 --target cuda 或不加参数)

步骤 C: 搭建本地演示站点

./scripts/local_deploy_site.sh

基本使用

1. 启动演示服务

运行上述 local_deploy_site.sh 脚本后，本地服务器将启动。

2. 启动浏览器

重要：必须使用特殊参数启动 Chrome Canary 以禁用鲁棒性检查，否则性能会大幅下降（约慢 3 倍）。

在终端执行以下命令启动浏览器（macOS 示例）：

/Applications/Google\ Chrome\ Canary.app/Contents/MacOS/Google\ Chrome\ Canary --enable-dawn-features=disable_robustness

(Windows/Linux 用户请替换为对应的 Chrome Canary 可执行文件路径)

3. 生成图像

1. 在浏览器中访问 http://localhost:8888/。
2. 在文本框中输入提示词（Prompt），例如：

   A photo of an astronaut riding a horse on mars.
   

3. 点击生成按钮，即可在浏览器端直接看到由 GPU 加速生成的图像。

4. 命令行本地测试 (可选)

如果不通过浏览器，也可以直接在本地原生 GPU 环境下测试生成的模型：

python3 deploy.py --prompt "A photo of an astronaut riding a horse on mars."

支持添加负面提示词：

python3 deploy.py --prompt "..." --negative-prompt "blurry, low quality"