flux.1-dev.py 516 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

4
from nunchaku.models.transformer_flux import NunchakuFluxTransformer2dModel
Zhekai Zhang's avatar
Zhekai Zhang committed
5

muyangli's avatar
muyangli committed
6
transformer = NunchakuFluxTransformer2dModel.from_pretrained("mit-han-lab/svdq-int4-flux.1-dev")
7
pipeline = FluxPipeline.from_pretrained(
muyangli's avatar
muyangli committed
8
    "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
Zhekai Zhang's avatar
Zhekai Zhang committed
9
).to("cuda")
muyangli's avatar
muyangli committed
10
image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
11
image.save("flux.1-dev-int4.png")