FLUX.1で画像生成

FLUX.1のSchnellモデルを使った画像生成をしてみます。

FLUXのオリジナルの実装

オリジナルの実装はこちらです。

モデル: https://huggingface.co/black-forest-labs/FLUX.1-schnell
コード: https://github.com/black-forest-labs/flux

指示にしたがって環境を設定し、実行すると、44.5GBのダウンロードがはじまります。

これは src/flux/util.py を読むと google/t5-v1_1-xxl をロードするようになっており、 https://huggingface.co/google/t5-v1_1-xxl/tree/main を見るとちょうどモデルサイズが44.5GBなので、おそらくこれをダウンロードしようとしているのだと思われます。

https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main に配置されているテキストエンコーダが使われない理由は謎です。

ダウンロードサイズは大きいですが https://github.com/black-forest-labs/flux/issues/7 によると、16GBくらいVRAMがあればよさそうです。sequential offloadをすれば2GBくらいらしいのですが、cli.pyのオプションを見てもそれらしきものを見つけることができませんでした。

量子化版

とはいえもう少し小さいモデルでまずは動かしてみたいので、量子化済みのモデルで試してみます。いくつかHuggingFaceにアップロードされていますが、今回はhttps://huggingface.co/aifoundry-org/FLUX.1-schnell-Quantizedを使ってみます。

現時点では[WIP]と書かれているものの動かし方がいちおうは記載されているので、そのとおりに実行してみます。

最初に table-diffusion.cpp をダウンロードしてビルドします。

次に、 Schnell-Q2_K と Autoencoder と CLIP_L と T5XXL をダウンロードして、models/以下に配置します。

これで準備は完了です。先ほどダウンロードした一番小さいモデル(Q2_K)で試してみます。

./build/bin/sd --diffusion-model models/flux1-schnell-Q2_K.gguf --vae models/ae.safetensors --clip_l models/clip_l.safetensors --t5xxl models/t5xxl_fp16.safetensors -p "a frog holding a sign saying 'hi' " -o ../frog.png -v --cfg-scale 1.0 --sampling-method euler -v --seed 42 --steps 4

を実行します。

以下の様な蛙の画像が出力されます。

他にも

./build/bin/sd --diffusion-model models/flux1-schnell-Q2_K.gguf --vae models/ae.safetensors --clip_l models/clip_l.safetensors --t5xxl models/t5xxl_fp16.safetensors -p "Photo-realistic Rabbit running in meadow" -o ../frog.png -v --cfg-scale 1.0 --sampling-method euler -v --seed 10 --steps 20

を実行すると

となります。

T5XXLはモデルサイズは大きいもののGPUメモリはあまり必要ないようで、デフォルトの512x512のサイズの画像生成時に観測できた範囲ではGPUメモリの使用量は4.7GBくらいでした。

./build/bin/sd --diffusion-model models/flux1-schnell-Q4_K.gguf --vae models/ae.safetensors --clip_l models/clip_l.safetensors --t5xxl models/t5xxl_fp16.safetensors -p "Photo-realistic Rabbit running in meadow" -o ../frog.png -v --cfg-scale 1.0 --sampling-method euler -v --seed 13 --steps 20

のようにQ2_KではなくQ4_Kを使ってみると

となります。

Q2_KではなくQ4_Kのモデルの方を使った場合、デフォルトの512x512のサイズの画像生成時に観測できた範囲ではGPUメモリの使用量は7.4GBくらいでした。

試した範囲では実行結果に再現性があり、コマンドの引数が同じであれば同じ結果が出力されるようです。