はじめに
Swift・Core MLで音源分離するライブラリ「swift-spleeter」を作りました。本記事では、このライブラリの使い方と制作過程を紹介します。
Core MLを扱ったネット上の記事は意外と少なく、苦労した点も多かったので、この記事がこれからCore MLを触ってみる人の助けになれば幸いです。
swift-spleeterとは
swift-spleeterとは、Spleeterという音源分離モデルをiOS等Appleプラットフォーム上で使えるようにしたものです。
Spleeterでは
- 2stems (ボーカル・伴奏)
- 4stems (ボーカル・ドラム・ベース・その他)
- 5stems (ボーカル・ドラム・ベース・ピアノ・その他)
という3種類の学習済みモデルが使用可能で、swift-spleeterでも全て対応しています。
使い方
音源ファイルが1つあれば、以下のコードで簡単に分離を実行できます。
import SwiftSpleeter
// コンパイル済みモデルを用意 ※
let modelURL = URL(fileURLWithPath: "path/to/Spleeter2.mlmodelc")
let separator = try AudioSeparator2(modelURL: modelURL)
// 分離したい音源ファイルのURL
let inputURL = URL(fileURLWithPath: "path/to/mixture.wav")
// 分離結果の保存先URL
let outputURLs = Stems2(
vocals: URL(fileURLWithPath: "path/to/vocals.wav"),
instruments: URL(fileURLWithPath: "path/to/instruments.wav")
)
// 分離実行 (for awaitで進捗を監視)
for try await progress in separator.separate(from: inputURL, to: outputURLs) {
print("Progress: \(progress.current)/\(progress.total)")
}
※
コンパイル済みモデルは、GitHub Releasesからダウンロードするか、jiyimeta/spleeter-pytorchを使って自分で作成してください (簡単にできます)。
制作過程
モデル選び
Pythonで音源分離を行うモデルはいくつか公開されていますが、以下のような条件があり、何でも良いわけではありません。
- iOSアプリで音源分離を実行するため、それなりに軽量である
- 汎用性を損なわないため、ExecuTorchやLiteRTのようなサードパーティーツールへの依存がない
- Core ML対応のため、モデル内部で未対応の処理を行っていない
- 公開アプリに組み込むため、商用利用可の学習済みモデルが公開されている
例えば、Demucsという高精度な音源分離モデルがありますが、これは処理が複雑、かつ波形ベースのため処理が重く、条件1を満たしません。また、内部にはCore ML未対応の処理が多く含まれており、条件3も満たしません。
Open-Unmixは比較的軽量ですが、内部にLSTMのようなCore ML対応不可のレイヤが含まれており、条件3を満たしません。こちらは変換可能なサードパーティーツールもあり、条件2を無視すれば使えそうにも見えましたが、学習済みモデルに使用されているデータセットMUSDB18には「学術目的のみ使用可」と明記されており条件4がグレーでした。
SpleeterはCore ML未対応のSTFTを含んでいるものの、前処理・後処理の段階で使われているだけだったので、これはSwift側で実装すれば解決可能でした。また、学習に使われているデータセットはDeezer社内のデータということで、リポジトリがMITライセンスなのも踏まえ、商用利用可能のようでした (実際様々なサービスで商用利用されている模様)。さらに、SpleeterをPyTorchで書き直した上でCore ML変換するところまで実装した天才的なリポジトリが既にあったので、ありがたくこれを使わせていただきました。
PyTorch実装の修正
元のSpleeterでは、2, 4, 5stems各モデルに対して、UNetに含まれるアクティベーション関数が設定ファイルに記述されており、それぞれ異なるパラメータが指定されています。
しかし、既存のPyTorch実装のSpleeterでは2stemsのみ対応していました。
そこで、元のSpleeterのTensorFlow実装と比較し、モデルを修正しました。
具体的な差分はこちらのコミットをご覧ください。
Core MLモデルの用意
coremltoolsというフレームワークを使うことで、Pythonで実装された学習済みモデルをCore ML形式に変換することができます。
この際、入出力の変数名は内部で利用されている文字列に自動設定されているため、人間が見て分かりやすい名前に置き換えました。
ct.utils.rename_feature(spec, "stft_mag_1", "magnitude")
ct.utils.rename_feature(spec, "var_614", "vocalsMask")
ct.utils.rename_feature(spec, "var_648", "instrumentsMask")
ここまでで出力されたモデルファイル (.mlprogram形式) のままでも、Xcodeに入れてiOS等のアプリで使用できるのですが、アプリ本体に組み込むことになるためアプリサイズが大きくなってしまうという問題があります。実際今回作ったモデルは35MB程度あり、無視できるか微妙なラインでした。そこで、公式記事に従ってスクリプトを書き、モデルをコンパイルし後からダウンロードして使えるようにしました。
ここで生成されたモデルは、固定サイズのデータしか受け取れないため、入力データを自分で細かく区切って与える必要があります。可変サイズ入力を受け付けるよう変更する方法もあるようですが、大規模データを一気に処理するのはデバイスのメモリが心配になるので、あえてこのままにしました。
STFTの実装
STFT (短時間フーリエ変換) は、簡単に言えば、音声波形という1次元データ (時間軸のみ) をスペクトログラムという2次元データ (時間軸×周波数軸) に変換する処理です。
Spleeterはスペクトログラムを用いて推論するモデルです。
先述の通り、Core ML対応の都合で、本来モデルに含まれていたSTFTをモデルから除いてしまったので、Swift側で実装する必要があります。
AccelerateフレームワークのvDSP (大規模配列を高速に処理するApple純正API群) をふんだんに用いて、Python実装をなるべく再現することを意識しながら、おおまかに以下のような処理をするようにしました。
- 入力波形の両端のデータ整合性を保つため、ゼロ埋めを両端に追加
- 周波数領域に変換するため、4096サンプルだけ取り出してhann窓をかけFFT (高速フーリエ変換) する
- この後の推論で冗長な (結果の良し悪しにあまり影響しない) 処理を削るため、適当な値以上の周波数データをまるごと消す
- strideを用いてずらしながら2~3を繰り返し、各時間ごとの周波数フレームを得る
詳細は実装を参照してください。
音源ファイルの入出力
AVFAudioを用いて音源ファイルの入出力を行います。ただし、長いファイルを読み込むとデバイスのメモリを圧迫し、最悪の場合アプリのクラッシュを引き起こす可能性があるため、指定サンプルだけを読み込み、指定サンプルだけを追加で書き出すように実装しました。
実際使うときは、先述のモデル入力の固定サイズに合わせた数のサンプルを取得するようにします。
推論の実装
Core MLにはMLModelというclassがあり、ここにモデルファイルを読み込んで使うのですが、このままではなかなか使いづらいです。
そこで、コンパイル前のモデル (.mlprogram) をXcodeに突っ込んだ時に自動生成されるコードをコピペ・改変し、MLModelを内包する扱いやすいインターフェースを用意しました。
これにより、
let model = Spleeter2Model(contentsOf: modelURL)
// MLTensor型の`magnitude` (絶対値スペクトログラム) を入力
let result = try await model.prediction(magnitude: magnitude)
// MLTensor型の出力を取得
let vocalsMask = result.vocalsMaskTensor
let instrumentsMask = result.instrumentsMaskTensor
のように、簡単に推論を行うことができるようになっています。
ドキュメント系の執筆
各APIのDocCコメント、README、およびLICENSEは、(人の手で修正しつつも大部分は) ChatGPTが書いてくれました。英語が苦手なので大変助かりました。
あとがき
Swiftで音源分離するswift-spleeterの使い方や制作過程を紹介しました。この記事を読まれた方は、ぜひswift-spleeterを使ってみてください & Core MLも触ってみてください。
iOSアプリ内での音源分離は数年前から構想していたのですが、多忙と怠惰により放置してしまっていました。今回開発したこのライブラリを使って、自分のアプリにも積極的に音源分離を導入していきたいと思います。