flux.1-dev.py 686 Bytes
Newer Older
Zhekai Zhang's avatar
Zhekai Zhang committed
1
import torch
2
from diffusers import FluxPipeline
Zhekai Zhang's avatar
Zhekai Zhang committed
3

muyangli's avatar
muyangli committed
4
from nunchaku import NunchakuFluxTransformer2dModel
5
from nunchaku.utils import get_precision
Zhekai Zhang's avatar
Zhekai Zhang committed
6

7
precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
8
9
10
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
    f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors"
)
11
pipeline = FluxPipeline.from_pretrained(
muyangli's avatar
muyangli committed
12
    "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
Zhekai Zhang's avatar
Zhekai Zhang committed
13
).to("cuda")
muyangli's avatar
muyangli committed
14
image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
15
image.save(f"flux.1-dev-{precision}.png")