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

选自 liuliu.me

机器之心编译

Stable Diffusion 可能很快就会在移动端普及。

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

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

应用商店链接：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. 图像解码器，从潜在表示中解码图像。

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

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

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

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

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

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

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

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

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

参考链接：

https://news.ycombinator.com/item?id=33539192

https://liuliu.me/eyes/stretch-iphone-to-its-limit-a-2gib-model-that-can-draw-everything-in-your-pocket/