【Diffusers】FLUX.1-dev の画像生成時間を短縮させる

はじめに

量子化を駆使して画像生成時間の短縮にトライしました。

トランスフォーマーとテキストエンコーダー2を4bitに変換しました。

もちろん画質は悪くなっていると思われます。

結果

RTX 4090を使っています。

VRAMは15GB程度使用していました。16GBのVRAMがあればかなり高速化できると思います。

結果①

プロンプト

A cat holding a sign that says hello world

生成画像

時間

time: 43.34sec

結果②

プロンプト

an insect robot preparing a delicious meal, anime style

生成画像

時間

time: 44.59sec

結果③

プロンプト

A photorealistic portrait of a young Japanese woman with long black hair and natural makeup, wearing a casual white blouse, sitting in a modern Tokyo cafe with soft window light

生成画像

時間

time: 45.07sec

Pythonスクリプト

モデルのロードに時間がかかるのであらかじめ量子化したものを保存しておきました。

import torch
from diffusers import FluxTransformer2DModel
from transformers import T5EncoderModel
from transformers import BitsAndBytesConfig as transformers_config
from diffusers import BitsAndBytesConfig as diffusers_config

model_id = "black-forest-labs/FLUX.1-dev"

# transformer
nf4_config = diffusers_config(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model_nf4 = FluxTransformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=nf4_config,
    torch_dtype=torch.bfloat16
)

model_nf4.save_pretrained("FLUX.1-transformer-nf4")

# text_encoder_2
text_encoder_2_config = transformers_config(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
text_encoder_2 = T5EncoderModel.from_pretrained(
    model_id,
    subfolder="text_encoder_2",
    quantization_config=text_encoder_2_config,
    torch_dtype=torch.bfloat16
)

text_encoder_2.save_pretrained("FLUX.1-text_encoder_2-nf4")

以下が画像生成の実行スクリプトです。

import time
import torch
from diffusers import FluxTransformer2DModel, FluxPipeline
from transformers import T5EncoderModel

start = time.perf_counter()

transformer = FluxTransformer2DModel.from_pretrained(
    "FLUX.1-transformer-nf4",
    torch_dtype=torch.bfloat16
)
text_encoder_2 = T5EncoderModel.from_pretrained(
    "FLUX.1-text_encoder_2-nf4",
    torch_dtype=torch.bfloat16
)

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer,
    text_encoder_2=text_encoder_2,
    torch_dtype=torch.bfloat16
)

pipe.enable_model_cpu_offload()

generator = torch.Generator().manual_seed(123)
prompt = "A photorealistic portrait of a young Japanese woman with long black hair and natural makeup, wearing a casual white blouse, sitting in a modern Tokyo cafe with soft window light"
out = pipe(
    prompt=prompt,
    height=768,
    width=1360,
    num_inference_steps=50,
    generator=generator,
    guidance_scale=3.5
).images[0]
out.save("girl.jpg")

end = time.perf_counter()
print(f"time: {(end - start):.2f}sec")