FLUX.1をdiffusers + torchaoでquantoよりも高速に使用する

# 概要

以前FLUX.1をdiffusers + quantoでVRAM 16GB環境でも読み込む方法を記事にしました。

https://qiita.com/yamichi77/items/1f12bac65d1584900b35#%E3%83%A2%E3%83%87%E3%83%AB%E8%AA%AD%E3%81%BF%E8%BE%BC%E3%81%BF%E9%80%9F%E5%BA%A6%E3%81%AE%E6%94%B9%E5%96%84

量子化ライブラリはほかにもあるようで、torchaoと呼ばれるライブラリでFLUX.1を量子化するやり方が最近Github上に上がっていました。
早速試したところ、quanto環境より高速に推論までできましたので、記事にしました。
以下のリポジトリを参考にしております。

https://github.com/sayakpaul/diffusers-torchao

# 目次
* [動作環境](#%E5%8B%95%E4%BD%9C%E7%92%B0%E5%A2%83)
* [環境構築](#%E7%92%B0%E5%A2%83%E6%A7%8B%E7%AF%89)
* [コード](#%E3%82%B3%E3%83%BC%E3%83%89)
* [quantoの量子化との速度比較](#quanto%E3%81%AE%E9%87%8F%E5%AD%90%E5%8C%96%E3%81%A8%E3%81%AE%E9%80%9F%E5%BA%A6%E6%AF%94%E8%BC%83)
* [まとめ](#%E3%81%BE%E3%81%A8%E3%82%81) 

# 動作環境
* Ubuntu 22.04
* Intel i7-13700K
* RTX A4000
* RAM：48GB

# 環境構築
pipで以下をインストールします。
```terminal
$ pip install torch torchvision transformers sentencepiece protobuf accelerate torchao diffusers huggingface_hub
```
:::note warn
9/7現在、pipのバージョンですと、diffusersとhuggingface_hubがtorchaoの量子化保存に対応しておりません。
Github上のソースから落としてくる必要があります。
```terminal
$ pip install git+https://github.com/huggingface/huggingface_hub
$ pip install git+https://github.com/huggingface/diffusers
```
:::

# コード
前回同様FLUX.1-schnellで動かします。

## 量子化した後に推論する
下記が基本的なコードになります。
ただ、量子化にかなり時間がかかるため、以下のやり方だとquanto環境より遅くなってしまいます。

```python: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")
```

## 量子化したモデルを保存し、ロードする
前回同様、量子化モデルを保存したほうが速いので、下記やり方をお勧めします。

### 量子化モデルを保存する
```python: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`を指定して保存します。

### 保存した量子化モデルで推論する
```python: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")
```

上記コードで前回と同じ画像が生成されると思います。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3739636/afc773eb-04c5-e005-a333-d5f73f04bd54.png)


# quantoの量子化との速度比較
前回紹介したquantoの量子化とのモデルロード～推論までの速度比較です。
| 種類 | 時間（秒） |
| ---- | ---- |
| quanto | 100.35 |
| torchao | 43.90~74.90 |

私の環境ではtorchaoの推論速度にはばらつきがありましたが、1/2～3/4までの推論時間の短縮が見られました。

# まとめ
今回はtorchaoを使用しFLUX.1を量子化することで、quantoより高速に推論することができました。
diffusersにおけるFLUX.1の推論でのパフォーマンス改善にお役立ていただければ幸いです。