Web Stable Diffusion
这个项目将stable diffusion模型带到了网络浏览器上。**所有操作都在浏览器内进行,无需服务器支持。**据我们所知,这是全球首个完全在浏览器上运行的稳定扩散。请查看我们的演示网页来试用。
如果您对在浏览器上部署基于LLM的聊天机器人感兴趣,我们也欢迎您查看Web LLM。
最近我们看到AI模型取得了惊人的进展。得益于开源努力,开发者现在可以轻松组合开源模型来完成惊人的任务。稳定扩散使基于文本输入的自动生成真实感图像以及各种风格的图像成为可能。这些模型通常很大且计算密集,这意味着我们在开发基于这些模型的网络应用程序时,必须将所有计算请求传输到(GPU)服务器。此外,大多数工作负载必须在特定类型的GPU上运行,因为那里有流行的深度学习框架。
这个项目旨在改变现状,为生态系统带来更多多样性。将部分(或全部)计算任务转移到客户端有很多理由。可能的好处包括降低服务提供商的成本,以及提升个性化和隐私保护。个人计算机(甚至移动设备)的发展方向正是提供这样的可能性。客户端变得非常强大。
为这些应用构建特殊的客户端应用程序是一个选择(我们也支持),但如果我们能简单地打开浏览器,将AI直接本地带到浏览器标签中,这岂不是更棒?生态系统中已有一定的准备就绪。WebAssembly使我们能够将更多的低级运行时移植到网络上。为了解决计算问题,WebGPU最近逐渐成熟,并开启了浏览器内的原生GPU执行。
我们在客户端看到了硬件和浏览器生态系统的必要元素逐渐汇集。不过,仍有许多障碍需要克服,例如:
- 我们需要在没有相关GPU加速Python框架的地方引入模型。
- 大多数AI框架严重依赖由硬件供应商维护的优化计算库。我们需要从零开始。为了获得最大收益,我们可能还需要针对不同的客户端环境生成变体。
- 需要仔细规划内存使用,使模型能够适配内存。
我们不仅仅希望为单个模型做到这些。而是希望展示一个可重复、可破解、可组合的工作流程,使得任何人都能在Python优先的环境中轻松开发和优化这些模型,并且能够在包括网络在内的所有地方通用部署这些模型。
入门
我们提供了一个Jupyter notebook,涵盖了所有阶段,包括
- 详细说明网络ML模型部署的重要点以及我们如何满足这些点,
- 导入stable diffusion模型,
- 优化模型,
- 构建模型,
- 使用原生GPU运行时在本地部署模型,及
- 使用WebGPU运行时在网络上部署模型。
如果您希望通过命令行完成这些步骤,请按照以下命令操作:
命令
-
安装TVM Unity。您可以
- 使用
pip3 install mlc-ai-nightly -f https://mlc.ai/wheels
安装TVM Unity软件包,或 - 按照TVM文档从源代码构建。请在git clone后使用
git checkout origin/unity
切换到TVM Unity。
- 使用
-
导入、优化并构建stable diffusion模型:
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运行时在网络上部署模型,请参考walkthrough notebook的最后部分“部署到网络”中的完整说明。我们也在这里提供相同的说明内容:
说明
首先,让我们安装所有前置条件:
- emscripten。这是一个基于LLVM的编译器,将C/C++源代码编译为WebAssembly。
- 按照安装说明安装最新的emsdk。
- 通过
source path/to/emsdk_env.sh
来源emsdk_env.sh
,使得emcc
可从PATH访问并使命令emcc
有效。
- Rust。
wasm-pack
。它帮助构建Rust生成的WebAssembly,在我们的案例中用于Tokenizer。- 按顺序安装jekyll。这是我们用于网站的包。
- 通过命令安装jekyll-remote-theme
gem install jekyll-remote-theme
- 安装Chrome Canary。这是Chrome的开发者版本,允许使用WebGPU。
我们可以通过分别在终端中尝试
emcc
、jekyll
和wasm-pack
来验证安装是否成功。然后,准备网络构建所需的所有依赖项:
./scripts/prep_deps.sh
现在我们可以将模型构建到WebGPU后端并将可执行文件导出到磁盘中的WebAssembly文件格式,运行
python3 build.py --target webgpu
最后一步是设置网站:
./scripts/local_deploy_site.sh
设置网站后,您可以在Chrome Canary中访问
localhost:8888/
来在本地机器上试用演示。不要忘记使用/Applications/Google\ Chrome\ Canary.app/Contents/MacOS/Google\ Chrome\ Canary --enable-dawn-features=disable_robustness
启动Chrome Canary以禁用Chrome的健壮性检查。
- emscripten。这是一个基于LLVM的编译器,将C/C++源代码编译为WebAssembly。
如何实现?
这里的关键技术是机器学习编译(MLC)。我们的解决方案建立在开源生态系统的基础上,包括PyTorch、Hugging Face的diffusers和tokenizers、rust、wasm和WebGPU。主要流程是建立在Apache TVM Unity之上,这是Apache TVM中一个令人兴奋的持续发展。
- 我们从 Hugging Face 的 diffuser 库中获取了Runway 的 stable diffusion 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 着色器扩展在规范中但尚未实现)。这里的稳定扩散模型需要一个拥有足够内存(8GB)的 GPU。我们目前只在 Apple silicon 上测试了我们的解决方案。还有机会支持诸如 FlashAttention 和量化等高级优化,以进一步提高系统性能。
这些都是为当前解决方案带来数倍性能改进的机会。我们相信很多问题能在不久的将来得到解决。这个解决方案的单一组件仍然是有用的。例如,可以选择仅部署模型的文本编码器部分。此外,相同的以 Python 为先的开发、通用部署工作流程也可以用于将 ML 模型引入其他环境,如新硬件或移动场景。最后,同样的机器学习编译栈也共享给服务器级用例,并可以用来优化服务器工作负载。
致谢
这个项目得益于以下合作的可能
这个项目得以实现要感谢我们所依靠的开源生态系统。我们要感谢 Apache TVM 社区和 TVM Unity 努力的开发人员。我们要感谢使这些模型公开的开源 ML 社区成员,以及使这些模型可访问的 PyTorch 和 Hugging Face 社区。我们还要感谢 Mithril Security 的 tokenizer wasm 端口。我们还要感谢 WebAssembly、Emscripten、Rust 和 WebGPU 社区。最后,感谢 Dawn 开发人员,他们提供了对 Chrome 问题的及时回答。