概要
以前FLUX.1をdiffusers + quantoでVRAM 16GB環境でも読み込む方法を記事にしました。
量子化ライブラリはほかにもあるようで、torchaoと呼ばれるライブラリでFLUX.1を量子化するやり方が最近Github上に上がっていました。
早速試したところ、quanto環境より高速に推論までできましたので、記事にしました。
以下のリポジトリを参考にしております。
目次
動作環境
- Ubuntu 22.04
- Intel i7-13700K
- RTX A4000
- RAM：48GB
環境構築
pipで以下をインストールします。
$ pip install torch torchvision transformers sentencepiece protobuf accelerate torchao diffusers huggingface_hub
9/7現在、pipのバージョンですと、diffusersとhuggingface_hubがtorchaoの量子化保存に対応しておりません。
Github上のソースから落としてくる必要があります。
$ pip install git+https://github.com/huggingface/huggingface_hub
$ pip install git+https://github.com/huggingface/diffusers
コード
前回同様FLUX.1-schnellで動かします。
量子化した後に推論する
下記が基本的なコードになります。
ただ、量子化にかなり時間がかかるため、以下のやり方だとquanto環境より遅くなってしまいます。
flux1.py
import torch
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
from torchao.quantization import int8_weight_only, quantize_
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
dtype = torch.bfloat16
bfl_repo = "black-forest-labs/FLUX.1-schnell"
# モデルの読み込み
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
bfl_repo, subfolder="scheduler"
)
text_encoder = CLIPTextModel.from_pretrained(
"openai/clip-vit-large-patch14", torch_dtype=dtype
)
tokenizer = CLIPTokenizer.from_pretrained(
"openai/clip-vit-large-patch14", torch_dtype=dtype
)
text_encoder_2 = T5EncoderModel.from_pretrained(
bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype
)
tokenizer_2 = T5TokenizerFast.from_pretrained(
bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype
)
vae = AutoencoderKL.from_pretrained(
bfl_repo, subfolder="vae", torch_dtype=dtype
)
transformer = FluxTransformer2DModel.from_pretrained(
bfl_repo, subfolder="transformer", torch_dtype=dtype
)
# 8bit量子化
quantize_(transformer, int8_weight_only())
# FluxPipelineの設定
pipe = FluxPipeline(
scheduler=scheduler,
text_encoder=text_encoder,
tokenizer=tokenizer,
text_encoder_2=text_encoder_2,
tokenizer_2=tokenizer_2,
vae=vae,
transformer=None,
)
pipe.transformer = transformer
pipe.enable_model_cpu_offload()
# シードの固定
generator = torch.Generator().manual_seed(0)
# 画像の生成
image = pipe(
prompt="A cat holding a sign that says hello world",
width=1024,
height=1024,
num_inference_steps=4,
max_sequence_length=256,
generator=generator,
guidance_scale=0.0,
).images[0]
image.save("flux-schnell.png")
量子化したモデルを保存し、ロードする
前回同様、量子化モデルを保存したほうが速いので、下記やり方をお勧めします。
量子化モデルを保存する
flux1_save.py
import torch
from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
from torchao.quantization import int8_weight_only, quantize_
dtype = torch.bfloat16
bfl_repo = "black-forest-labs/FLUX.1-schnell"
transformer = FluxTransformer2DModel.from_pretrained(
bfl_repo, subfolder="transformer", torch_dtype=dtype
)
# FluxTransformerの保存
quantize_(transformer, int8_weight_only())
transformer.save_pretrained(
"./FLUX.1-schnell-torchao/transformer",
safe_serialization=False,
)
quantoと違い、safe_serialization=Falseを指定して保存します。
保存した量子化モデルで推論する
flux1_load.py
import torch
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
from torchao.quantization import int8_weight_only, quantize_
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
dtype = torch.bfloat16
bfl_repo = "black-forest-labs/FLUX.1-schnell"
# モデルの読み込み
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
bfl_repo, subfolder="scheduler"
)
text_encoder = CLIPTextModel.from_pretrained(
"openai/clip-vit-large-patch14", torch_dtype=dtype
)
tokenizer = CLIPTokenizer.from_pretrained(
"openai/clip-vit-large-patch14", torch_dtype=dtype
)
text_encoder_2 = T5EncoderModel.from_pretrained(
bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype
)
tokenizer_2 = T5TokenizerFast.from_pretrained(
bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype
)
vae = AutoencoderKL.from_pretrained(
bfl_repo, subfolder="vae", torch_dtype=dtype
)
transformer = FluxTransformer2DModel.from_pretrained(
"./FLUX.1-schnell-torchao",
subfolder="transformer",
torch_dtype=dtype,
use_safetensors=False,
)
# FluxPipelineの設定
pipe = FluxPipeline(
scheduler=scheduler,
text_encoder=text_encoder,
tokenizer=tokenizer,
text_encoder_2=text_encoder_2,
tokenizer_2=tokenizer_2,
vae=vae,
transformer=None,
)
pipe.transformer = transformer
pipe.enable_model_cpu_offload()
# シードの固定
generator = torch.Generator().manual_seed(0)
# 画像の生成
image = pipe(
prompt="A cat holding a sign that says hello world",
width=1024,
height=1024,
num_inference_steps=4,
max_sequence_length=256,
generator=generator,
guidance_scale=0.0,
).images[0]
image.save("flux-schnell.png")
上記コードで前回と同じ画像が生成されると思います。
quantoの量子化との速度比較
前回紹介したquantoの量子化とのモデルロード～推論までの速度比較です。
| 種類 | 時間（秒） |
|---|---|
| quanto | 100.35 |
| torchao | 43.90~74.90 |
私の環境ではtorchaoの推論速度にはばらつきがありましたが、1/2～3/4までの推論時間の短縮が見られました。
まとめ
今回はtorchaoを使用しFLUX.1を量子化することで、quantoより高速に推論することができました。
diffusersにおけるFLUX.1の推論でのパフォーマンス改善にお役立ていただければ幸いです。