把Stable Diffusion模型塞进iPhone里，做成APP一分钟出图

在 iPhone 上运行 Stable Diffusion 到底难不难？今天我们要介绍的这篇文章，作者给出了答案：不难，而且 iPhone 还剩余 50% 的性能。

众所周知，每年苹果都会推出一款声称在各方面都更快、更好的新款 iPhone，这主要得益于新的视觉模型和图像传感器的快速发展。就拿拍照来说，如果回到 10 年前，你能用 iPhone 拍出高质量的图片吗，答案是不能，因为技术的发展是渐进式的，10 年时间，足够提高手机拍照技术。

由于技术的这种发展模式（渐进式），在一段时间里有些程序即使运行在最好的计算设备上，也几乎无法使用。但是这些带有新启用场景的新程序吸引了一些用户的注意力，人们愿意研究它。

本文的作者就是被吸引的其中之一，在过去 3 周里，作者开发了一个应用程序，可以通过 Stable Diffusion 来生成（summon）图像，然后按你喜欢的方式编辑它。该应用在最新的 iPhone 14 Pro 上生成图片仅需一分钟，使用大约 2GiB 的应用内存，另外还需要下载大约 2GiB 的初始数据才能开始使用。

应用商店链接：https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820

这一结果引来众多网友讨论，有人开始担心手机耗电问题，并开玩笑的说：这很酷，不过这看起来是一个消耗手机电池的好方法。

「我从来没有像现在这样开心地感受 iPhone 的热度。」

「这个寒冬，可以把手机当暖手器用了。」

不过在大家调侃手机发热问题的同时，他们也给与这项工作极高的评价。

「这简直不可思议。在我的 iPhone SE3 上生成一张完整的图像大约需要 45 秒——这几乎和我的 M1 Pro macbook 用原始版本生成的速度一样快！」

内存、硬件同时优化

这是如何做到的呢？接下来我们看看作者的实现过程：

想要完成在 iPhone 上运行 Stable Diffusion，还能结余 50% 的性能，面临的一大挑战是需要在 6GiB RAM 的 iPhone 设备上将程序运行起来。6GiB 听起来很多，但如果你在 6GiB 设备上使用超过 2.8GiB，或在 4GiB 设备上使用超过 2GiB，iOS 就会杀死你的应用程序。

那么 Stable Diffusion 模型究竟需要多少内存来进行推理？

这还要从模型的结构说起。通常 Stable Diffusion 模型包含 4 个部分：1. 文本编码器，它生成文本特征向量以指导图像生成；2. 可选的图像编码器，将图像编码到潜在空间 (用于图像到图像生成)；3. 降噪器模型，它从噪声中缓慢地去噪图像的潜在表示；4. 图像解码器，从潜在表示中解码图像。

第 1、第 2 和第 4 个模块在推理过程中运行一次，最大需要约 1GiB。而降噪器模型占用了大约 3.2GiB（全浮点数），并且还需要执行多次，因此作者想让该模块在 RAM 中保存得更久。

最初的 Stable Diffusion 模型需要接近 10GiB 才能执行单个图像推理。在单个输入（2x4x64x64）与输出（2x4x64x64）之间，其中夹杂着许多输出层。并不是所有层的输出都可以被立即复用，它们中一部分必须保留一些参数以供后续使用（残差网络）。

一段时间以来，研究者围绕 PyTorch Stable Diffusion 进行了一番优化，对 PyTorch 用到的 NVIDIA CUDNN 和 CUBLAS 库，他们保留了暂存空间，这些优化都是为了降低内存使用量，因此 Stable Diffusion 模型可以用低至 4GiB 的卡运行。

但这仍然超出了作者的预期。因此作者开始专注于苹果硬及优化。

起初作者考虑的是 3.2GiB 或 1.6GiB 半浮点数，如果不想触发苹果的 OOM（Out of Memory，指的是 App 占用的内存达到了 iOS 系统对单个 App 占用内存上限后，而被系统强杀掉的现象），作者大约有 500MiB 的空间可以使用。

第一个问题，每个中间输出的大小到底是多少？

事实证明，它们中的大多数都相对较小，每个都低于 6MiB (2x320x64x64)。作者使用的框架 (s4nnc) 可以合理地将它们打包到小于 50MiB，以备复用。

值得一提的是，降噪器有一个自注意机制，它以自己的图像潜在表示作为输入。在自注意力计算期间，有一个大小为 16x4096x4096 的批处理矩阵，对该矩阵应用 softmax 后，大约是 FP16 中的 500MiB，并且可以「inplace」完成，这意味着它可以安全地重写其输入而不会损坏。幸运的是，Apple 和 NVIDIA 低级库都提供了 inplace softmax 实现，然而 PyTorch 等更高级的库中没有。

那么是否真的使用 550MiB + 1.6GiB 左右的内存就能完成？

在 Apple 硬件上，实现神经网络后端的一个常用选择是使用 MPSGraph 框架。于是作者首先尝试使用 MPSGraph 实现了所有的神经网络操作。在 FP16 精度下峰值内存使用量大约是 6GiB，显然比预期的内存使用量多太多，这是怎么回事？

作者详细分析了原因，首先他没有按照常见的 TensorFlow 方式使用 MPSGraph。MPSGraph 需要对整个计算图进行编码，然后使用输入 / 输出张量，进而处理内部分配，并让用户提交整个图以供执行。

而作者使用 MPSGraph 的方式很像 PyTorch 的做法——当作一个操作执行引擎。为了执行推理任务，许多已编译的 MPSGraphExecutable 在 Metal 命令队列上执行，它们中的每一个都可能持有一些中间分配内存。如果一次性提交，那么所有这些命令都持有分配内存，直到完成执行。

一种解决这个问题的简单方法是调整提交速度，没有必要一次性提交所有命令。实际上，Metal 的每个队列有 64 个并发提交的限制。作者尝试改成一次提交 8 个操作，峰值内存就降低到了 4GiB。

然而，这仍然比 iPhone 能承受的多 2 GiB。

为了使用 CUDA 计算自注意力，原始 Stable Diffusion 代码实现中有一个常见技巧：使用置换而不是转置。这个技巧很有效，因为 CUBLAS 可以直接处理置换的跨步（strided）张量，避免使用专用内存来转置张量。

但是 MPSGraph 没有跨步张量支持，一个置换的张量无论如何都会在内部被转置，这需要中间分配内存。通过显式转置，分配将由更高级别的层处理，避免了 MPSGraph 内部效率低下。利用这个技巧，内存使用量将接近 3GiB。

事实证明，从 iOS 16.0 开始，MPSGraph 不能再为 softmax 做出最优分配决策。即使输入和输出张量都指向相同的数据，MPSGraph 也会分配一个额外的输出张量，然后将结果复制到指向的位置。

作者发现使用 Metal Performance Shaders 替代方案完全符合要求，并将内存使用量降至 2.5GiB，而不会出现任何性能下降。

另一方面，MPSGraph 的 GEMM 内核需要内部转置。显式转置在此也无济于事，因为这些转置不是更高级别层的「inplace」操作，对于特定的 500MiB 大小的张量，这种额外的分配是不可避免的。通过切换到 Metal Performance Shaders，项目作者又回收了 500MiB，性能损失约为 1%，最终将内存使用量减到了理想的 2GiB。