Zenn
🖼️

Core ML Stable Diffusionを自分のiOSアプリに組み込む手順

2024/07/22に公開
iOS
Swift
Core ML
Stable Diffusion
visionOS
tech

はじめに

Core ML Stable Diffusion（を用いた画像生成機能）を自分のiOSアプリに追加する手順をまとめる。

以下のようにiOSオンデバイスでText-to-Imageの画像生成ができるようになる。


iOS-GenAI-Samplerに画像生成サンプルを追加した

なおiOSアプリ以外にも、Core MLが利用可能なmacOS、visionOS、watchOSアプリにも同様の手順でCore ML Stable Diffusionを組み込むことが可能 [1]。

サンプルコード

本記事に書いてある実装は生成AI × iOSのサンプルコード集である下記OSSで公開している。

https://github.com/shu223/iOS-GenAI-Sampler

アプリとして実際に動かしてみつつ、コード全体を確認できる。

関連記事

ml-stable-diffusion を用いたCore ML形式への変換方法は下記記事に、

https://zenn.dev/shu223/articles/ml-stable-diffusion

swift-coreml-diffusers の実装解説は下記記事にまとめてある。

https://note.com/shu223/n/n8529c602a78c

またそもそものCore ML Stable Diffusionの入門としては下記記事が参考になると思う。

https://zenn.dev/shu223/articles/coreml-stable-diffusion

ml-stable-diffusion のREADME

Core ML Stable DiffusionをiOSアプリで動かすには、ml-stable-diffusion リポジトリにあるSwift Packageを使う。

同リポジトリのREADMEの「Image Generation with Swift」配下に「ライブラリ使用例（Example Library Usage）」として以下のコードが書かれており、

import StableDiffusion
...
let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
pipeline.loadResources()
let image = try pipeline.generateImages(prompt: prompt, seed: seed).first

非常に簡単に実装できそうだが、残念ながらこのコードは概念コードのようなもので、現実的にはこうシンプルにはいかない。

コンパイル済みモデル群のある resourceURL を与えるためにモデルをダウンロードしてきて展開する実装が必要 [2] だし、 StableDiffusionPipeline の初期化、loadResources 、generateImages メソッドのコールが直列に並んでいるのも実際のアプリでこんな実装をするわけにはいかない [3]。

あと、そもそも普通にコード自体がいろいろ間違っていて、StableDiffusionPipeline のイニシャライザの引数が足りないし [4]、loadResources を呼ぶときに try が必要だし、generateImages に seed なんて引数はない。

Swift Package Details

本記事では解説しないが、同READMEにある、「Swift Package Details」という項目を転載しておく

Swift Package Details

This Swift package contains two products:

  • StableDiffusion library
  • StableDiffusionSample command-line tool

Both of these products require the Core ML models and tokenization resources to be supplied. When specifying resources via a directory path that directory must contain the following:

  • TextEncoder.mlmodelc or `TextEncoder2.mlmodelc (text embedding model)
  • Unet.mlmodelc or UnetChunk1.mlmodelc & UnetChunk2.mlmodelc (denoising autoencoder model)
  • VAEDecoder.mlmodelc (image decoder model)
  • vocab.json (tokenizer vocabulary file)
  • merges.text (merges for byte pair encoding file)

Optionally, for image2image, in-painting, or similar:

  • VAEEncoder.mlmodelc (image encoder model)

Optionally, it may also include the safety checker model that some versions of Stable Diffusion include:

  • SafetyChecker.mlmodelc

Optionally, for the SDXL refiner:

  • UnetRefiner.mlmodelc (refiner unet model)

Optionally, for ControlNet:

  • ControlledUNet.mlmodelc or ControlledUnetChunk1.mlmodelc & ControlledUnetChunk2.mlmodelc (enabled to receive ControlNet values)

  • controlnet/ (directory containing ControlNet models)

    • LllyasvielSdControlnetMlsd.mlmodelc (for example, from lllyasviel/sd-controlnet-mlsd)
    • LllyasvielSdControlnetDepth.mlmodelc (for example, from lllyasviel/sd-controlnet-depth)
    • Other models you converted

Note that the chunked version of Unet is checked for first. Only if it is not present will the full Unet.mlmodelc be loaded. Chunking is required for iOS and iPadOS and not necessary for macOS.

手順

まずはHugging Faceが出しているサンプルアプリ、swift-coreml-diffusers の実装を借りて手軽に済ませる手順を示す。

1. Swift Package追加

https://github.com/apple/ml-stable-diffusion から追加

Package Productが "Stable Diffusion"、Kindが ”Library" の方をアプリのターゲットに追加する。

2. swift-coreml-diffusers から実装をコピー

Diffusion/Common 配下

.
├── DiffusionImage.swift
├── Downloader.swift
├── ModelInfo.swift
├── Pipeline
│   ├── Pipeline.swift
│   └── PipelineLoader.swift
├── State.swift
├── Utils.swift
└── Views
    └── PromptTextField.swift

Diffusion/Views 配下

  • Loading.swift
  • TextToImage.swift

Diffusion 配下

  • DiffusionImage+iOS.swift
  • Utils_iOS.swift

3. ビルドが通るよう修正する

ZIPFoundation パッケージを追加

https://github.com/weichsel/ZIPFoundation

定義をコピーしてくる

DiffusionApp.swift にある以下の定義をどこかにコピー

let runningOnMac = ProcessInfo.processInfo.isMacCatalystApp
let deviceHas6GBOrMore = ProcessInfo.processInfo.physicalMemory > 5910000000   // Reported by iOS 17 beta (21A5319a) on iPhone 13 Pro: 5917753344
let deviceHas8GBOrMore = ProcessInfo.processInfo.physicalMemory > 7900000000   // Reported by iOS 17.0.2 on iPhone 15 Pro Max: 8021032960

let deviceSupportsQuantization = {
    if #available(iOS 17, *) {
        true
    } else {
        false
    }
}()

一部コード修正

StableDiffusion パッケージの最新リリース版（ 2024年7月現在は 1.1.1）を取り込みつつ、ソースコードとしては swift-coreml-diffusers リポジトリのものをコピペしてくると、いろいろビルドエラーが発生する。swift-coreml-diffusers では main ブランチを取り込んで実装しているため。

たとえば StableDiffusion3Pipeline という定義はStableDiffusion パッケージの 1.1.1 にはない。

では StableDiffusion の 1.1.1 に対応した swift-coreml-diffusers のバージョンをフェッチしてくればよいかというと、そういうタグは打たれていない [5]。

というわけで、コピーしてきた実装を適当に修正してビルドエラーを直していく。

  • isSD3 の定義と、それを参照している部分を削除 → StableDiffusion3Pipeline を参照している部分もなくなる
  • discreteFlowScheduler 関連の実装を削除

    • Scheduler for rectified flow based multimodal diffusion transformer models

    • らしい

いちおうこれだけの修正でビルドは通るようになる。

4. 動作確認

LoadingView を表示すればダウンロードが始まり、

@main
struct DiffusionApp: App {
    var body: some Scene {
        WindowGroup {
            LoadingView()
        }
    }
}

ダウンロードが終わればパイプラインの初期化とモデルの読み込みが始まり、
それらが終われば TextToImage ビューが表示され、画像生成が行える。

ただし、コピーしてきた実装のままではXLモデルが利用される（iPhone 15 Proを使用した場合） [6]。

カスタマイズ

iOS-GenAI-Sampler では簡単に以下のようなカスタマイズを行っている。

  • モデルを選べるようにする

    • iosModel の定義をなくし、ビューのイニシャライザでモデルを指定できるようにした
    • v2.1の圧縮版とXLの圧縮版を試せるようにした

  • 不要な実装を削除

    • runningOnMac・・・iOS向けは不要
    • deviceSupportsQuantization・・・iOS 17以上であれば不要

実プロダクトに組み込むにはコピペベースで済ませるわけにはいかないと思うが、ここから不要なものを削りつつ足りないものを追加しつつUIは全部書き直しつつ、という感じで自分なら進めると思う。

脚注

  1. 現実的にはwatchOSデバイスでStable Diffusionを動かすにはメモリや処理能力が足りないと思われる。 ↩︎

  2. アプリにバンドルしても良いが、合計数GBのモデルを同梱するのは非現実的。 ↩︎

  3. 前述の入門記事に書いたが、loadResources はCore MLモデルの読み込みを伴う処理で、非常に時間がかかる。画像生成のたびに行う必要はない。 ↩︎

  4. controlNet 引数にはデフォルト値が設定されていない。 ↩︎

  5. サンプルアプリだからか、あまり本家パッケージのリリースバージョンと対応したバージョニングや開発は行われておらず、「常にmainを向いたまま随時更新」という感じで運用されているようだ。 ↩︎

  6. 再掲になるが、このあたりの実装の詳細はこちらの記事を参照： swift-coreml-diffusers のコードを読んだメモ｜shu223 ↩︎

Discussion