Core ML版Stable DiffusionをiOSで快適に動かす

Core ML版 Stable Diffusion を iOS で快適に動かす 堤 修一 @shu223

自己紹介 • 堤 修一 • @shu223 (GitHub, Qiita, Zenn, note, 𝕏, YouTube, Podcast, etc...) • 書籍（商業出版4冊、個人出版多数 @BOOTH）:

• https://shu223.booth.pm/items

• 1年間研究開発部のマネージャーにチャレンジ • 12年間iOSだけやってきたエンジニアのキャリアチェンジ • 2023年6月からエンジニアに復帰 • マネージャーからエンジニアに戻ります • 最近のお仕事 • タッチ名刺交換

• https://shu223.hatenablog.com/entry/2022/03/24/184840
• https://shu223.hatenablog.com/entry/2023/05/29/174918
• https://buildersbox.corp-sansan.com/entry/2023/08/04/120104

本日のお題について • ⭕ ❌ 「Stable Diffusion」という画像生成AIの解説 • 巨大なモデルをiOSで動かす際の考え方や勘所

Stable Diffusionとは • 画像生成AI • 入力テキストに応じて画像を自動 生成するtext-to-imageモデル • 2022.8 オープンソースとして公開

Core ML Stable Diffusion とは Core ML フォーマットに変換された Stable Diffusion のモデル • 従来のモデルをAppleハードウェアで動かす => CPUのみ利 用 • Core MLモデル => CPU, GPU, Neural Engineを利用 → Appleのハードウェアを最大限活かせるのがCore MLモデル 1 詳細な最適化の解説： Deploying Transformers on the Apple Neural Engine - Apple Machine Learning Research 1

• https://machinelearning.apple.com/research/neural-engine-transformers

Core ML Stable Diffusion のリポジトリ https://github.com/apple/ml-stable-diffusion • Apple謹製 • モデル変換コード、macOS/iOSで動かすためのSwiftコード等 を含む • 公式ブログ記事： Stable Diffusion with Core ML on Apple Silicon - Apple Machine Learning Research

• https://github.com/apple/ml-stable-diffusion
• https://machinelearning.apple.com/research/stable-diffusion-coreml-apple-silicon

初登場時以降の進化 2022年12月の初公開時以降あまり話題になっていないが、着々 と進化している • ControlNetサポート • SDXLサポート • 6ビット量子化（関連記事： 1, 2, 3） • Multilingual text encoder サポート • 関連記事： [iOS 17] 多言語BERT埋め込みモデルのサポー

• https://note.com/shu223/n/n346df0d6b30a
• https://note.com/shu223/n/n136945bcc210
• https://note.com/shu223/n/n54f428063564

Core ML Stable Diffusion の何が嬉しいのか？ • 手元のMacで、無料かつオフラインで使いまくれる • 自前のmacOS/iOSアプリに組み込める • For me: 様々なノウハウが詰まったCore MLの公式サンプル

CLIから動かす Swiftで書かれたCLIサンプル StableDiffusionSampleを利用 swift run StableDiffusionSample \ "a photo of an astronaut riding a horse on mars" \ --resource-path </path/to/models> --seed 93 \ --output-path </path/to/output/image> • --resource-path • コンパイル済みCore MLモデルファイル .mlmodelc が 入ったディレクトリのパスを指定

自作アプリに組み込む 1/2 Swift Package としてプロジェクトに追加する • → Package Product "StableDiffusion" をターゲットに追加

自作アプリに組み込む 2/2 import StableDiffusion ... let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL) pipeline.loadResources() let image = try pipeline.generateImages(prompt: prompt, seed: seed).first

Hugging Faceによるサンプルアプ リ • GitHub • App Store

• https://github.com/huggingface/swift-coreml-diffusers
• https://apps.apple.com/app/diffusers/id1666309574

Diffusers を iPhone 15 Pro で動かしてみる

モデルのダウンロード （初回のみ必要）

モデルのロード （起動毎に必要）

⚠ ローディングがとてつもなく長い

起動する度に 5分 2 Pipeline loaded in 293.2635090351105 2 モデルのダウンロードが完了した状態からの起動

起動状態

⚠ 画像生成時に必ずクラッシュ

使用メモリ量のスパイク （iPhone 15 Pro利用）

iOSで「まともに」動かす

iOSでまともに動かす • 使用メモリ量を小さくする • 起動時間を短くする • 画像生成時間を短くする

iOSでまともに動かす • 使用メモリ量を小さくする • 起動時間を短くする • 画像生成時間を短くする

使用メモリ量を小さくする方法 • reduceMemory オプションを有効にする • サイズの小さいモデルに変更する

使用メモリ量を小さくする方法 • ⭕ reduceMemory オプションを有効にする • サイズの小さいモデルに変更する

reduceMemory を有効にする ModelInfo.swift の reduceMemory プロパティを決定する 処理を書き換える 3 var reduceMemory: Bool { // Enable on iOS devices, except when using quantization if runningOnMac { return false } if isXL { return !deviceHas8GBOrMore } return !(quantized && deviceHas6GBOrMore) } 3 現行実装ではiPhone 15 Proの場合は必ず無効化されてしまう

reduceMemory を有効にするとどうなるか？ • 必要なときにモデルをロードし、すぐにアンロードする • 画像生成処理のオーバーヘッドを増やすことになる • ドキュメントでは解説されていないが、ソースコードを見る と起動時に prewarmResources という処理を行っている public func loadResources() throws { if reduceMemory { try prewarmResources() } ...

prewarmResources loadResources 後すぐに unloadResources を行う func prewarmResources() throws { try loadResources() unloadResources() } Both .mlpackage and .mlmodelc models are compiled upon first load. （中略） .mlmodelc files do cache this compiled asset and non-first load times are reduced to just a few seconds.

デモ

画像生成時にクラッシュしなくなった

画像生成の処理時間 Got images: [Optional(<CGImage 0x12810d510> (IP) <<CGColorSpace 0x283e04780> (kCGColorSpaceDeviceRGB)> width = 768, height = 768, bpc = 8, bpp = 24, row bytes = 2304 kCGImageAlphaNone | 0 (default byte order) | kCGImagePixelFormatPacked is mask? No, has masking color? No, has soft mask? No, has matte? No, should interpolate? Yes)] in 83.99049496650696 1枚あたり 83.99 秒 起動時間も依然として 5分 かかっている

iOSでまともに動かす • 使用メモリ量を小さくする • 起動時間を短くする • 画像生成時間を短くする

起動時間を短くする方法 起動時間のほとんどが、モデルのロード （実態としては 4 MLModel の初期化）に時間がかかっている → サイズの小さいモデルに変更する 4 ライブラリ内の loadResources メソッド

画像生成時間を短くする方法 サイズの小さいモデルに変更する The Neural Engine is capable of accelerating models with low-bit palettization: 1, 2, 4, 6 or 8 bits. （ニューラル・エンジンは、低 ビットのパレタイズでモデルを高速化できる） compressed weights are faster to fetch from memory ... （圧縮さ れたウェイトはメモリからのフェッチが高速になり、） 5 apple/ml-stable-diffusion: Stable Diffusion with Core ML on Apple Silicon のREADMEより 5

• https://github.com/apple/ml-stable-diffusion

（再掲）使用メモリ量を小さくする方法 ⭕ • reduceMemory オプションを有効にする • サイズの小さいモデルに変更する

とにかくモデルの圧縮が重要

現在の使用モデルを調べる 1/3 let loader = PipelineLoader(model: iosModel()) func iosModel() -> ModelInfo { guard deviceSupportsQuantization else { return ModelInfo.v21Base } if deviceHas6GBOrMore { return ModelInfo.xlmbpChunked } return ModelInfo.v21Palettized } → iPhone 15 Proでは ModelInfo.xlmbpChunked を利用 6 deviceSupportsQuantization が true かつ deviceHas6GBOrMore が true 6

現在の使用モデルを調べる 2/3 static let xlmbpChunked = ModelInfo( modelId: "apple/coreml-stable-diffusion-xl-base-ios", modelVersion: "SDXL base (768, iOS) [4 bit]", supportsEncoder: false, quantized: true, isXL: true )

現在の使用モデルを調べる 3/3 apple/coreml-stable-diffusion-xl-base-ios • iOS向けに 4.04 混合ビット量子化（mixed-bit palettization）したもの • ・・・ではあるが、もともとがSDXLモデルなのでクソでか い • Unetだけで 1.3 GB以上

• https://huggingface.co/apple/coreml-stable-diffusion-xl-base-ios

モデルを圧縮する or モデルの圧縮版を探す

Hugging Face で公開されている Core ML Stable Diffusion モデル → 2023年10月現在で16種類

• https://huggingface.co/models?other=core-ml&amp;search=diffusion

! どれを選べばいいのか

軽いSDモデルを選ぶポイント • SDXLではないもの・・・XLは最新の高画質版であり、デカ い • palletized とついているもの・・・圧縮されていることを示 す • 16bit -> 6-bit • SDバージョンの違いによるサイズの違いはあまりない → "apple/coreml-stable-diffusion-2-1-base-palettized"

デモ

各種パフォーマンス • 起動時間 ：100秒前後（300%高速化） 7 • 使用メモリ量のピーク値 ：約850MB 8 • 画像生成処理時間：10秒以下（800%高速化） 7 アプリの起動時間全体ではなく、パイプラインのロード時間を計測 8 サイズの小さいモデルの使用に伴い、reduceMemory は再び無効化した

さらなる起動時間の短縮 モデルのロード（MLModelの初期化）が90%以上を占める time to load unet: 57.41793203353882 time to load textEncoder: 30.462964057922363 time to load decoder: 3.1042929887771606 time to load encoder: 2.146523952484131 Pipeline resources loaded in 106.28662097454071

さらなるモデル圧縮 Mixed-Bit Palettization 公式リポジトリREADME の "Advanced Weight Compression (Lower than 6-bits)" に詳細あり

• https://github.com/apple/ml-stable-diffusion#core-ml-stable-diffusion

more • ControlNet • Multilingual text encoder • 関連記事： [iOS 17] 多言語BERT埋め込みモデルのサポー ト

• https://note.com/shu223/n/n54f428063564

まとめ Stable DiffusionのようなクソデカCore MLモデルをiOSで動かす • キャッシュ機構を利用して都度ロード & アンロード • → メモリ使用量を抑える • モデルサイズを小さくする • → 起動時間・メモリ使用量・画像生成時間すべて改善 • Hugging Face探す/coremltoolsで自前圧縮/Mixed-Bit Palettization

ご清聴ありがとうございました！