该项目将稳定的扩散模型引入Web浏览器。一切都在浏览器中运行,没有服务器支持。据我们所知,这是世界上第一个完全在浏览器上运行的稳定扩散。请查看我们的演示网页进行尝试。
我们最近通过人工智能模型看到了惊人的进步。由于开源的努力,开发人员现在可以轻松地将开源模型组合在一起,以产生惊人的任务。稳定的扩散可以自动创建逼真的图像以及基于文本输入的各种样式的图像。这些模型通常很大且计算量很大,这意味着在基于这些模型开发 Web 应用程序时,我们必须将所有计算请求通过管道传输到 (GPU) 服务器。此外,大多数工作负载必须在特定类型的GPU上运行,在这些GPU上,流行的深度学习框架很容易获得。
该项目迈出了改变现状并为生态系统带来更多多样性的一步。有很多原因将部分(或全部)计算提供给客户端。有许多可能的好处,例如服务提供商方面的成本降低,以及个性化和隐私保护的增强。个人电脑(甚至移动设备)的发展正朝着实现这种可能性的方向发展。客户端变得非常强大。例如,最新的MacBook Pro可以拥有高达96GB的统一RAM,可用于存储模型权重和相当强大的GPU来运行许多工作负载。
为这些应用程序构建特殊的客户端应用程序是一种选择(我们也支持),但如果我们可以简单地打开浏览器并直接将 AI 本机带到你的浏览器选项卡,那不是更令人惊讶吗?生态系统中有一定程度的准备。WebAssembly 允许我们将更多较低级别的运行时移植到 Web 上。为了解决计算问题,WebGPU 最近越来越成熟,并支持在浏览器上执行本机 GPU。
我们只是看到客户端的必要元素在硬件和浏览器生态系统方面结合在一起。尽管如此,仍有很大的障碍需要跨越,仅举几例:
我们不想只对一个模型这样做。相反,我们想提供一个可重复的、可破解的、可组合的工作流程,使任何人都可以在 Python 优先的环境中轻松开发和优化这些模型,并将它们普遍部署到任何地方,包括 Web。
我们有一个Jupyter笔记本,可以引导你完成所有阶段,包括
如果你想在命令行中完成这些步骤,请按照以下步骤操作:
安装 TVM Unity。你可以
pip3 install mlc-ai-nightly -f https://mlc.ai/wheels
git 结帐源/统一来签出到 TVM Unity。
要导入、优化和构建稳定的扩散模型,请执行以下操作:
python3 build.py
默认情况下,作为构建目标。你还可以通过以下方式指定 CUDA 目标
build.py
apple/m2-gpu
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 上部署模型,演练笔记本的最后一部分“在 Web 上部署”列出了可以参考的完整说明。我们还在此处提供了相同的简单说明列表:
首先,让我们安装所有先决条件:
emsdk_env.sh
source path/to/emsdk_env.sh
emcc
emcc
瓦斯姆包
。它有助于构建 Rust 生成的 WebAssembly,在我们这里的例子中用于标记器。gem install jekyll-remote-theme
我们可以通过尝试来验证安装是否成功,分别在终端中。
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/web-stable-diffusion/
/Applications/Google\ Chrome\ Canary.app/Contents/MacOS/Google\ Chrome\ Canary --enable-dawn-features=disable_robustness
以启动 Chrome Canary 以关闭来自 Chrome 的稳健性检查。
这里的关键技术是机器学习编译(MLC)。我们的解决方案建立在开源生态系统的肩膀上,包括 PyTorch、Hugging Face 扩散器和分词器、rust、wasm 和 WebGPU。主流建立在Apache TVM Unity上,这是Apache TVM中令人兴奋的持续发展。
这个工作流程的所有部分都是用Python完成的,当然,除了最后一部分,它构建了一个400-loc的JavaScript应用程序,将事物连接在一起。这也是一个有趣的互动开发过程,带来了新的模型。
所有这些都是通过我们利用的开源生态系统实现的。具体来说,我们大量使用 TVM Unity,这是 TVM 项目中令人兴奋的最新发展,它支持这种 Python 优先的交互式 MLC 开发体验,使我们能够轻松地编写新的优化,全部使用 Python,并逐步将我们的应用程序带到 Web。TVM Unity 还提供了一种在生态系统中组合新解决方案的简单方法。例如,我们将来可以轻松地将其他 WebGPU 着色器生成器或着色器库引入此工作流程。
除了 WebGPU 运行时之外,我们还提供本地 GPU 运行时的本机部署选项。这些选项既可以用作在本机环境中部署的工具,也可以用作比较本机 GPU 驱动程序性能和 WebGPU 的参考点。
WebGPU 的工作原理是将 WGSL(WebGPU 着色语言)着色器转换为原生着色器。因此,从理论上讲,我们可以在 WebGPU 运行时和本机环境之间达到零差距。但是,如果我们直接使用 Chrome 在 Apple 芯片上检查当前的演示,我们会发现性能下降(约 3 倍)。这是因为Chrome的WebGPU实现为所有数组索引访问插入绑定剪辑,因此变为.理想情况下,下游着色器编译器应该能够优化绑定裁剪,但不幸的是,情况并非如此。一旦 WebGPU 实现变得更加成熟,检查索引访问范围并删除此类裁剪,就可以修复此差距。
a[i]
a[min(i, a.size)]
你可以通过使用特殊标志启动Chrome(感谢Dawn开发人员提供指针)来解决此问题,方法是完全退出Chrome,然后在命令行中键入
/path/to/chrome-canary --enable-dawn-features=disable_robustness
然后你会发现执行速度和原生GPU环境一样快。我们预计随着 WebGPU 的成熟,这个问题将得到解决。
我们只是看到了我们认为是火山爆发的曙光。WebGPU仍在不断发展(尽管今年它接近出货),并且只能通过Chrome Canary获得,并且可能不稳定。它还具有限制,例如仅支持 FP32(FP16 着色器扩展在规范中,但尚未实现)。这里的稳定扩散需要具有相当数量的 RAM (8GB) 的 GPU。到目前为止,我们仅通过 Apple 芯片测试了我们的解决方案。还有机会支持高级优化,如FlashAttention和量化,以进一步提高系统的性能。
这些机会使当前解决方案的性能提高了数倍。我们相信,其中许多问题在不久的将来可以得到解决。此解决方案的单个组件仍然有用。例如,可以选择仅部署模型的文本编码器部分。此外,相同的 Python 优先开发通用部署工作流可用于将 ML 模型引入其他环境,例如新硬件或移动案例。最后,相同的机器学习编译堆栈也与服务器类用例共享,也可用于优化服务器工作负载。
该项目之所以成为可能,要归功于与以下机构的合作:
这个项目之所以成为可能,要归功于我们肩负的开源生态系统。我们要感谢 Apache TVM 社区和 TVM Unity 工作的开发人员。我们要感谢公开这些模型的开源 ML 社区成员,以及使这些模型易于访问的 PyTorch、Hugging Face 社区。我们要感谢Mithril Security的tokenizer wasm port。我们还要感谢 WebAssembly、Emscripten、Rust 和 WebGPU 社区。最后,感谢 Dawn 开发人员,他们及时回答了 Chrome 上的问题。