web-stable-diffusion - 为网络浏览器带来稳定的扩散模型。一切都在浏览器内运行,没有服务器支持。

Created at: 2023-03-06 08:50:51
Language: Jupyter Notebook
License: Apache-2.0

卷筒纸稳定扩散

该项目将稳定的扩散模型引入Web浏览器。一切都在浏览器中运行,没有服务器支持。据我们所知,这是世界上第一个完全在浏览器上运行的稳定扩散。请查看我们的演示网页进行尝试。

浏览器截图

我们最近通过人工智能模型看到了惊人的进步。由于开源的努力,开发人员现在可以轻松地将开源模型组合在一起,以产生惊人的任务。稳定的扩散可以自动创建逼真的图像以及基于文本输入的各种样式的图像。这些模型通常很大且计算量很大,这意味着在基于这些模型开发 Web 应用程序时,我们必须将所有计算请求通过管道传输到 (GPU) 服务器。此外,大多数工作负载必须在特定类型的GPU上运行,在这些GPU上,流行的深度学习框架很容易获得。

该项目迈出了改变现状并为生态系统带来更多多样性的一步。有很多原因将部分(或全部)计算提供给客户端。有许多可能的好处,例如服务提供商方面的成本降低,以及个性化和隐私保护的增强。个人电脑(甚至移动设备)的发展正朝着实现这种可能性的方向发展。客户端变得非常强大。例如,最新的MacBook Pro可以拥有高达96GB的统一RAM,可用于存储模型权重和相当强大的GPU来运行许多工作负载。

为这些应用程序构建特殊的客户端应用程序是一种选择(我们也支持),但如果我们可以简单地打开浏览器并直接将 AI 本机带到你的浏览器选项卡,那不是更令人惊讶吗?生态系统中有一定程度的准备。WebAssembly 允许我们将更多较低级别的运行时移植到 Web 上。为了解决计算问题,WebGPU 最近越来越成熟,并支持在浏览器上执行本机 GPU。

我们只是看到客户端的必要元素在硬件和浏览器生态系统方面结合在一起。尽管如此,仍有很大的障碍需要跨越,仅举几例:

  • 我们需要将模型带到没有相关 GPU 加速 Python 框架的地方。
  • 大多数 AI 框架都严重依赖由硬件供应商维护的优化计算库。我们需要从零开始。为了获得最大收益,我们可能还需要为每个客户端环境生成变体。
  • 仔细规划内存使用,以便我们可以将模型放入内存中。

我们不想只对一个模型这样做。相反,我们想提供一个可重复的、可破解的、可组合的工作流程,使任何人都可以在 Python 优先的环境中轻松开发和优化这些模型,并将它们普遍部署到任何地方,包括 Web。

开始

我们有一个Jupyter笔记本,可以引导你完成所有阶段,包括

  • 详细说明 Web ML 模型部署的关键点以及我们如何满足这些要点,
  • 导入稳定扩散模型,
  • 优化模型,
  • 构建模型,
  • 使用本机 GPU 运行时在本地部署模型,以及
  • 使用 WebGPU 运行时在 Web 上部署模型。

如果你想在命令行中完成这些步骤,请按照以下步骤操作:

命令
  • 安装 TVM Unity。你可以

    • 用于安装 TVM Unity 滚轮,或
      pip3 install mlc-ai-nightly -f https://mlc.ai/wheels
    • 按照 TVM 的文档从源代码构建。请在 git 克隆后使用 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 上部署”列出了可以参考的完整说明。我们还在此处提供了相同的简单说明列表:

    指示

    首先,让我们安装所有先决条件:

    1. Emscripten。它是一个基于 LLVM 的编译器,可将 C/C++ 源代码编译到 WebAssembly。
      • 按照安装说明安装最新的 emsdk。
      • Source by ,以便可以从 PATH 访问,并且该命令有效。
        emsdk_env.sh
        source path/to/emsdk_env.sh
        emcc
        emcc
    2. 生锈
    3. 瓦斯姆包。它有助于构建 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/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中令人兴奋的持续发展。

  • 我们从拥抱脸扩散器库中获取 Runway 的稳定扩散 v1-5 模型。
  • 我们使用TorchDynamoTorch FX将关键模型组件捕获到TVM中的IRModule中。
  • TVM IRModule中的每个函数都可以使用可运行的代码进一步转换和生成,这些代码可以普遍部署在最小TVM运行时支持的任何环境中(javascript就是其中之一)。
  • TensorIRMetaSchedule用于构建自动化解决方案以生成优化的程序。这些转换通过本机 GPU 运行时在特定设备上进行调整,然后用于生成优化的 GPU 着色器。我们提供了一个记录这些转换的数据库,因此无需调整即可完成新的构建。
  • 我们构建静态内存规划优化,以跨多个层重用内存。
  • 我们使用Emscripten和打字稿来构建一个可以部署生成的模块的TVM Web运行时。
  • 我们还利用了 rust 标记器库wasm 端口来自拥抱面。

工作流程

这个工作流程的所有部分都是用Python完成的,当然,除了最后一部分,它构建了一个400-loc的JavaScript应用程序,将事物连接在一起。这也是一个有趣的互动开发过程,带来了新的模型。

所有这些都是通过我们利用的开源生态系统实现的。具体来说,我们大量使用 TVM Unity,这是 TVM 项目中令人兴奋的最新发展,它支持这种 Python 优先的交互式 MLC 开发体验,使我们能够轻松地编写新的优化,全部使用 Python,并逐步将我们的应用程序带到 Web。TVM Unity 还提供了一种在生态系统中组合新解决方案的简单方法。例如,我们将来可以轻松地将其他 WebGPU 着色器生成器或着色器库引入此工作流程。

与本机 GPU 运行时、限制和机会的比较

除了 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 模型引入其他环境,例如新硬件或移动案例。最后,相同的机器学习编译堆栈也与服务器类用例共享,也可用于优化服务器工作负载。

确认

该项目之所以成为可能,要归功于与以下机构的合作:

清迈大学计算机科学学院 催化剂 MLC 八角星

这个项目之所以成为可能,要归功于我们肩负的开源生态系统。我们要感谢 Apache TVM 社区和 TVM Unity 工作的开发人员。我们要感谢公开这些模型的开源 ML 社区成员,以及使这些模型易于访问的 PyTorch、Hugging Face 社区。我们要感谢Mithril Security的tokenizer wasm port。我们还要感谢 WebAssembly、Emscripten、Rust 和 WebGPU 社区。最后,感谢 Dawn 开发人员,他们及时回答了 Chrome 上的问题。