flux.1-fill-dev.py 1.04 KB
Newer Older
muyangli's avatar
muyangli committed
1
2
3
4
import torch
from diffusers import FluxFillPipeline
from diffusers.utils import load_image

muyangli's avatar
muyangli committed
5
from nunchaku import NunchakuFluxTransformer2dModel
6
from nunchaku.utils import get_precision
muyangli's avatar
muyangli committed
7

Muyang Li's avatar
Muyang Li committed
8
9
image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")
muyangli's avatar
muyangli committed
10

11
precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
12
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
13
    f"nunchaku-tech/nunchaku-flux.1-fill-dev/svdq-{precision}_r32-flux.1-fill-dev.safetensors"
14
)
muyangli's avatar
muyangli committed
15
16
17
18
pipe = FluxFillPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Fill-dev", transformer=transformer, torch_dtype=torch.bfloat16
).to("cuda")
image = pipe(
Muyang Li's avatar
Muyang Li committed
19
    prompt="a white paper cup",
muyangli's avatar
muyangli committed
20
21
22
23
24
25
26
27
    image=image,
    mask_image=mask,
    height=1024,
    width=1024,
    guidance_scale=30,
    num_inference_steps=50,
    max_sequence_length=512,
).images[0]
28
image.save(f"flux.1-fill-dev-{precision}.png")