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

以上是把Stable Diffusion模型塞進iPhone裡，做成APP一分鐘出圖的詳細內容。更多資訊請關注PHP中文網其他相關文章！