swift-coreml-diffusers のコードを読んだメモ

この発表をする際に、自前で変換・圧縮したCore ML Stable Diffusion検証用として下記OSSアプリを使用したのだが、

その際にいろいろ中身を読んだのでメモ。主にiOS版の話。

なお reduceMemory など一部のコードリーディング結果については下記記事にも書いている。

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

ModelInfo の定義

ModelInfo の extension には static let でこういうのがズラズラと定義されている。

    static let v21Base = ModelInfo(
        modelId: "pcuenq/coreml-stable-diffusion-2-1-base",
        modelVersion: "StabilityAI SD 2.1",
        supportsEncoder: true
    )
    
    static let v21Palettized = ModelInfo(
        modelId: "apple/coreml-stable-diffusion-2-1-base-palettized",
        modelVersion: "StabilityAI SD 2.1 [6 bit]",
        supportsEncoder: true,
        supportsAttentionV2: true,
        quantized: true
    )

• modelId・・・使用するモデルのHugging Face Hubリポジトリを特定するID

• supportsEncoder・・・VAEEncoderをサポートするかどうか

  • 下記記事に書いた通り、text-to-imageでは必要ない

Stable DiffusionをCore MLモデルに変換する

• supportsAttentionV2・・・モデル変換の際にAttention Implementation として SPLIT_EINSUM_V2 を指定した場合はこちらに true を指定する

• quantized・・・圧縮したモデルかどうか。

  • Linear Quantizationに限らず、圧縮手段としてPalettizationを利用している場合も true にするようだ。（元から定義してあるv21Palettizedがそうなっているため）

iOSで利用するモデル

Loading.swiftに次のような定義がある。

func iosModel() -> ModelInfo {
    guard deviceSupportsQuantization else { return ModelInfo.v21Base }
    if deviceHas6GBOrMore { return ModelInfo.xlmbpChunked }
    return ModelInfo.v21Palettized
}

この実装に従い、iOS向けにビルドした場合は次のようにモデルが選択される

• デバイスがQuantizationをサポートしてない場合・・・v21Base

• デバイスのメモリが6GB以上ある場合・・・xlmbpChunked

• それ以外・・・v21Palettized

v21Base

    static let v21Base = ModelInfo(
        modelId: "pcuenq/coreml-stable-diffusion-2-1-base",
        modelVersion: "StabilityAI SD 2.1",
        supportsEncoder: true
    )

非圧縮なv2.1モデル。2.1baseのCore MLモデルはapple公式でも提供しているのだが、pcuenqアカウントのものを向いている。

xlmbpChunked

    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
    )

4.04混合ビットPalettizationを使用したXLモデル。

変数名としては"mbp"とされているが、リポジトリ名としては"-ios"とついており、iOSデバイスで動かせるレベルを目標として圧縮されたモデル。

This version uses 4.04 mixed-bit palettization and generates images with a resolution of 768×768. It uses SPLIT_EINSUM attention and is intended for use in iOS/iPadOS 17 or better.

apple/coreml-stable-diffusion-xl-base-ios · Hugging Face

v21Palettized

    static let v21Palettized = ModelInfo(
        modelId: "apple/coreml-stable-diffusion-2-1-base-palettized",
        modelVersion: "StabilityAI SD 2.1 [6 bit]",
        supportsEncoder: true,
        supportsAttentionV2: true,
        quantized: true
    )

6ビット量子化されたv2.1モデル。Attention V2使用。

apple/coreml-stable-diffusion-2-1-base-palettized · Hugging Face

deviceSupportsQuantization

シンプルにiOS 17以上かどうかで判定されている。

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

ダウンロードとunzip

PipelineLoaderクラスの `prepare()` メソッドにて、モデルのダウンロードと解凍（unzip）が行われる。