【SwiftUI】Core ML Stable Diffusionをアプリに実装する

今回はCore ML Stable Diffusionをアプリで実行できるようにします。

GitHub - apple/ml-stable-diffusion: Stable Diffusion with Core ML on Apple Silicon

Stable Diffusion with Core ML on Apple Silicon. Contribute to apple/ml-stable-diffusion development by creating an account on GitHub.

残念ながら私はApple Siliconを搭載したiPadを所持していない為、
macOSに対応したマルチプラットフォームiOSアプリとして作成していますが、
特にコードが変わる旨は書かれていなかったので基本的には同じコードで問題ないと思います。

また、今回はmacOS対応の為にSwiftUIを使用していますが、
Core ML Stable Diffusionのコードに関してはUIKitで使っても問題ないと思います。

下準備

モデル変換

Core ML用のStable Diffusionの学習済みモデルが必要となります。
持っていない人は前回の記事を参照して変換済みモデルを用意して下さい。
–bundle-resources-for-swift-cliオプションを使用したものが必要です。

なおMacで実行する場合はそのままのUnetファイルで問題ありませんが、
iOSデバイスで実行する場合はチャンクしたものでないといけません。
変換時には–chunk-unetオプションを忘れずにつけ下さい。

配布モデルについて

AppleがHugging Face Hubにて変換済みのStable Diffusionモデルを配布しています。
既存モデルをアプリに実装するだけの場合はそちらを利用できます。
GitHubリポジトリにREADMEにも書いてあります。

ただし、変換可能な環境を用意しておく事で、
自由にモデルを選んで使用する事ができるので、
１度環境構築しておく事をオススメします。

ml-stable-diffusionパッケージを追加する

まずはml-stable-diffusionパッケージを追加します。

左のナビゲーターの右クリックメニューまたはメニューバーのFileからAdd Packageを選択して下さい。

出てきたウィンドウの左上から検索してml-stable-diffusionを選択し追加します。

Sampleは要らないのでStable Diffusionにのみチェックを入れて追加します。

モデルファイルの追加

プロジェクトにモデルファイルを追加します。
生成方法は前置きにも書いた通り前回の記事を参考にして下さい。

必要なファイルは以下の通りです。

merges.txt
TextEncoder.mlmodelc
Unet.mlmodelc または UnetChunk1.mlmodelcとUnetChunk2.mlmodelc
VAEDecoder.mlmodelc
vocab.json

以下はオプションです

SafetyChecker.mlmodelc

Unetについてですが、macOSの場合はどちらでもよく、
iOSデバイスの場合はchunk済みのものでなければなりません。
なお両方使用できる場合はUnet.mlmodelcが優先的に読み込まれます。

今回はResourcesディレクトリごとまとめて追加します。

以下ファイル追加について

※Xcodeの基本通りで特別な事は無いので、慣れている人は読み飛ばしてOKです。

今回はCreate groupsを選択しました。
後で使用する際にアプリのルートのパスを参照するだけで済みます。

Create folder referencesを選択すると、
アプリから使用する際にResourcesを含めたパスで指定する必要があります。

コピーはお好みでどうぞ。
ファイルが大きいのでコピーすると環境次第では反映まで若干時間がかかります。

画像生成を実装する

画像生成のコード

画像生成のコードを確認していきましょう。

ml-stable-diffusionで書かれている通り以下のコードを使用します。

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

画像生成を実行するだけならpipelineを用意してgenerateImagesを実行するだけです。

resourceURLにはモデルの入っているディレクトリを指定します。

promptは画像生成に使用するテキストをStringを、seedはInt型を使用します。
パラメータの詳細についてはStableDiffusionの仕様の話になるので割愛します。

imageはCGImageで返ってきます。

実行可能な形にする

UIから実行する様にしましょう。
ひとまずパラメータは固定にしています。

struct ContentView: View {
    
    @State var dispImage:CGImage?
    
    var body: some View {
        VStack {
            if let image = dispImage {
                Image(image,scale: 1 , label: Text("Generated Image"))
            }else{
                Text("Generated Image")
                    .frame(width: 512,height: 512)
                    .border(.black)
            }
            
            Button("Generate"){
                generateImage()
            }
        }
        .padding()
    }
    
    func generateImage(){
        guard let resourceURL = Bundle.main.resourceURL else{
            return
        }
        
        let prompt = "a photo of an astronaut riding a horse on mars"
        let seed = 93
        
        do{
            print("load models")
            let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
            
            print("generate image")
            dispImage = try pipeline.generateImages(prompt: prompt, seed: seed).first ?? nil
        } catch(let error) {
            print(error.localizedDescription)
        }
    }
}

UI部分は単純にImageの表示とボタンだけです。
CGImageがnilの時は代わりにテキストを表示しています。

generateImageメソッドもパラメータを設定して実行しているだけです。
StableDiffusionの実行にエラーハンドリングが必要なのでdo-catchを追加しました。

resourceURLですが、今回はCreate groupsでやっているのでBundle.main.resourceURLそのままです。
Create folder referencesにしたらディレクトリ名をappendしてあげて下さい。

実行するとこのように画像が生成されます。

画像の生成にそこそこ時間がかかります。
とりあえずprintで出していますが、何も表示しないと進捗が分からないまま待つことになるので気を付けましょう。
またメインスレッドで実行しているので現状ではUI操作もブロックされる状態になっています。

非同期で実行する

メインスレッドで実行してしまっているので、
途中停止や進捗表示を実装しようにもUI操作が効きません。
サブスレッドで実行するように変更しましょう。

非同期実行にはasync/awaitを使います。
非同期実行はいくつか解説記事を書いていますが以下の記事から飛べる筈です。

非同期追加時のコード全文

struct ContentView: View {
    
    @State var dispImage:CGImage?
    @State var disableGenerateButton = false
    
    var body: some View {
        VStack {
            if let image = dispImage {
                Image(image,scale: 1 , label: Text("Generated Image"))
            }else{
                Text("Generated Image")
                    .frame(width: 512,height: 512)
                    .border(.black)
            }
            
            Button("Generate"){
                disableGenerateButton = true

                Task{
                    await generateImage()
                    disableGenerateButton = false
                }
            }.disabled(disableGenerateButton)
        }
        .padding()
    }
    
    func generateImage() async {
        guard let resourceURL = Bundle.main.resourceURL else{
            return
        }
        
        let prompt = "a photo of an astronaut riding a horse on mars"
        let seed = 93
        
        do{
            print("load models")
            let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)

            print("generate image")
            let image = try pipeline.generateImages(prompt: prompt, seed: seed).first ?? nil
            
            Task{
                await MainActor.run {
                    dispImage = image
                }
            }
            
        } catch(let error) {
            print(error.localizedDescription)
        }
    }
}

変更部分を抜粋します。

generateImageメソッド

まずはgenerateImageメソッドです。

非同期実行するメソッドなのでasyncを追加しました。

また、生成した画像をメンバ変数に入れる際はViewの更新が絡むのでメインスレッドでの実行が望ましいです。
その為、画像を一旦変数で保持し、メンバ変数の変更はMainActor.runを使ってメインスレッドへ投げています。
MainActor.runにはawaitを使うのでTaskで囲っています。

func generateImage() async {
    //省略
    
    do{
        //省略

        print("generate image")
        let image = try pipeline.generateImages(prompt: prompt, seed: seed).first ?? nil
        
        Task{
            await MainActor.run {
                dispImage = image
            }
        }
        
    } catch(let error) {
        print(error.localizedDescription)
    }
}

View

次はUI側です。

非同期実行によりUI操作が可能になるので連打防止でボタンの無効化を追加しました。

また、generateImageがasyncメソッドなのでawaitを付けて呼び出し、Taskで囲っています。
このdisableGenerateButton = falseはTask内に置かないと画像生成前に実行されるので注意して下さい。

@State var disableGenerateButton = false

///
Button("Generate"){
    disableGenerateButton = true

    Task{
        await generateImage()
        disableGenerateButton = false
    }
}.disabled(disableGenerateButton)

進捗を表示する

StableDiffusionPipelineにはProgressという構造体があります。
このProgressから画像生成の進捗を取得する事ができます。

Progressを使用する際はgenerateImagesのprogressHandler引数を使います。
なお今回はTrailing Closureで使っているので引数名は出てこないです。

@State var dispStep = 0
@State var dispStepCount = 0

///
pipeline.generateImages(prompt: prompt, seed: seed){ progress in

    Task{
        await MainActor.run {
            dispStepCount = progress.stepCount
            dispStep = progress.step
        }
    }   

    return true
}

このようにしてProgress構造体からstep数を取得しています。
stepCountが最大で、stepが現在の値です。

このprogressHandlerのクロージャはstep毎に実行されるので、
stepの進捗を随時確認する事ができます。

メンバ変数はViewの表示に関わるので先ほどと同様にメインスレッドに投げます。

Viewは適当な所に以下のTextを入れて表示します。

Text("Step:\(dispStep)/\(dispStepCount)")

変更部分のコード

struct ContentView: View {
    
    //省略
    @State var dispStep = 0
    @State var dispStepCount = 0
    
    var body: some View {
        VStack {
            //省略
            Button("Generate"){
                //省略
            }.disabled(disableGenerateButton)
            
            Text("Step:\(dispStep)/\(dispStepCount)")
        }
        .padding()
    }
    
    func generateImage() async {
        //省略
        do{
            let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)

            let image = try pipeline.generateImages(prompt: prompt, seed: seed){ progress in
                
                Task{
                    await MainActor.run {
                        dispStepCount = progress.stepCount
                        dispStep = progress.step
                    }
                }
                
                return true
            }.first ?? nil
            
            Task{
                await MainActor.run {
                    dispImage = image
                }
            }
            
        } catch(let error) {
            print(error.localizedDescription)
        }
    }
}

画像生成を中止する

先ほどのprogressHandlerでtrueを返していました。
ここをfalseにすると画像生成が中止されます。

以下のような形でStopボタンを作って、returnの値を変えるようにすることで中止できます。
Generateボタンを押して辺りでtrueになるようにしておきましょう。

@State var continueGenerate = false

///
Button("Generate"){
    continueGenerate = true
    //略
}

Button("Stop"){
    continueGenerate = false
}.disabled(!continueGenerate)

///
pipeline.generateImages(prompt: prompt, seed: seed){ progress in
    //略

    return continueGenerate
}

generateImagesで使用できるパラーメータについて

generateImagesの宣言部を抜粋しました。

    /// Text to image generation using stable diffusion
    ///
    /// - Parameters:
    ///   - prompt: Text prompt to guide sampling
    ///   - stepCount: Number of inference steps to perform
    ///   - imageCount: Number of samples/images to generate for the input prompt
    ///   - seed: Random seed which
    ///   - disableSafety: Safety checks are only performed if `self.canSafetyCheck && !disableSafety`
    ///   - progressHandler: Callback to perform after each step, stops on receiving false response
    /// - Returns: An array of `imageCount` optional images.
    ///            The images will be nil if safety checks were performed and found the result to be un-safe
    public func generateImages(
        prompt: String,
        imageCount: Int = 1,
        stepCount: Int = 50,
        seed: Int = 0,
        disableSafety: Bool = false,
        progressHandler: (Progress) -> Bool = { _ in true }
    ) throws -> [CGImage?]

以下の引数が使用できます。

引数	説明	備考
prompt	画像生成に使用するテキスト
imageCount	生成する画像の枚数
stepCount	ステップ数
seed	画像生成に使用するseed値	Uint32までマイナス値はダメ
disableSafety	セーフティを無効化するいわゆるNSFW Filterを無効化する	適当にnudeとか入れるとすぐに分かる trueにしないと生成されない
progressHandler	Progressを扱う為のクロージャを渡す

StableDiffusionを触っている方だと色々思う所はあると思いますが、それはさておき、
なんと解像度の指定がありません。
GitHubを見ると現在は対応していないのでforkして自分で実装してくれと書いてあります。

自力で頑張るか今後の実装に期待といった所ですね・・・

最後に

ざっくりとCore ML Stable Diffusionのアプリへの実装を紹介しました。

現在でもmacOSで動作しましたが、macOS 13.1、iPadOS16.2以降が対象のようです。
OSのアップデートまで時間があるので、Core ML Stable Diffusionも何らかの更新があるかもしれません。

まだ発表されたばかりなので、簡単にそのまま実用とはいきませんが今後の更新に期待したい所です。

おまけ

おまけ①

今回の記事で作成したコードに若干手を加えたものがあるので、せっかくなので公開します。
ざっくり作っただけなので粗がありますが参考までに

おまけ②

画像生成の中止まで加えたコード全文

import SwiftUI
import StableDiffusion

struct ContentView: View {
    
    
    @State var dispImage:CGImage?
    @State var disableGenerateButton = false
    
    @State var dispStep = 0
    @State var dispStepCount = 0
    
    var body: some View {
        VStack {
            
            if let image = dispImage {
                Image(image,scale: 1 , label: Text("Generated Image"))
            }else{
                Text("Generated Image")
                    .frame(width: 512,height: 512)
                    .border(.black)
            }
            
            
            Button("Generate"){
                disableGenerateButton = true
                
                Task{
                    await generateImage()
                    
                    disableGenerateButton = false
                }
                
            }
            
            Text("Step:\(Int(dispStep))/\(Int(dispStepCount))")
        }
    }
    
    
    func generateImage() async {
        
        guard let resourceURL = Bundle.main.resourceURL else{
            return
        }
        
        let prompt = "a photo of an astronaut riding a horse on mars"
        let seed = 93
        
        do{
            print("Models loding")
            let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
            
            print("Generate Image")
            let image = try pipeline.generateImages(prompt: prompt, seed: seed ){ progress in
                
                Task{
                    await MainActor.run {
                        dispStepCount = progress.stepCount
                        dispStep = progress.step
                    }
                }
                
                return true
                
            }.first ?? nil
            
            dispImage = image
            
            Task{
                await MainActor.run {
                    dispImage = image
                }
            }
            
        } catch(let error) {
            print(error.localizedDescription)
        }
    }
}

【Swift】Core ML Stable Diffusionを使ってみた（Apple Silicon） | thwork より:
2022年12月5日 9:35 PM
[…] 【SwiftUI】Core ML Stable Diffusionをアプリに組み込む今回はCore ML Stable Diffusionを… […]
返信
【SwiftUI】Core ML Stable Diffusionをアプリに実装するサンプルコード | thwork より:
2022年12月7日 4:52 PM
[…] 【SwiftUI】Core ML Stable Diffusionをアプリに実装する今回はCore ML Stable Diffusionを… […]
返信