在 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。