MacでStable Diffusionを動かしてみる

こんにちは。yshr10icです。

2023年12月にAppleシリコン搭載のMacに最適化されたMLライブラリである、MLXが発表されました。Pytorchライクにコードを書くことができるらしいです。
https://github.com/ml-explore/mlx

MLXの発表からまもないのですが、mlx-examplesのGitHubリポジトリには、単純なMNISTやCIFARを解くだけでなく、Transformerベースの言語モデルの学習やLLaMAによる文章生成、Stable Diffusionによる画像生成、Whisperによる会話認識といった昨今の生成AIブームに合せたサンプル実装も提供されています。
https://github.com/ml-explore/mlx-examples

Macでそんなに簡単にStable Diffusion動かせるの？と衝撃を受けたので、試しに実行してみました。

まずは、mlx-examplesのGitHubリポジトリをクローンします。

$ git clone [e-mail]:ml-explore/mlx-examples.git

次に必要なライブラリをインストールします。

$ cd mlx-examples/stable_diffusion
$ pip install -r requirements.txt

requirements.txtに定義されているのは、7つのライブラリで、mlxは2024年1月27日時点だと、0.0.11がインストールされました。

requirements.txt

mlx
safetensors
huggingface-hub
regex
numpy
tqdm
Pillow

たったこれだけで準備完了なので、実際に動かしてみます。

$ python txt2image.py "A cat sitting on the desk in front of the macbook." --n_images 1 --n_rows 1 --output sample.png

出来ました！初回はHugging Faceからモデルをダウンロードするので、10分くらいかかりました。また、モデルの重みなどを含めてだいたい5,6GBくらいストレージ容量は必要となります。
2回目以降の実行に関しては、モデルの読み込みと合わせて1分かからないくらいでした。なお、自分の環境はM2 Macbook Proのメモリ32GBです。

コマンドライン引数は、以下が用意されています。

引数	デフォルト値	説明
--n_images	4	生成画像数
--steps	50	ステップ数
--cfg	0.75	CFG Scale。プロンプトに対してどれだけ忠実な画像を生成するかを指定する数値
--negative_prompt	""	ネガティブプロンプト
--n_rows	2	生成画像を何行で出力するか
--decoding_batch_size	1	バッチサイズ
--output	out.png	出力ファイル名

先ほどの生成された画像から、キーボードとマウスを出さないようにネガティブプロンプトを指定して試したのがこちらです。

$ python txt2image.py "A cat sitting on the desk in front of the macbook." --n_images 1 --n_rows 1 --negative_prompt "keyboard, mouse"  --output sample2.png

今回は簡単な検証でしたが、MacでStable Diffusionが試せるのはすごいですね！次はLLaMAやWhisperを試してみたいと思います。