sana_1600m.py 726 Bytes
Newer Older
muyangli's avatar
muyangli committed
1
2
import torch
from diffusers import SanaPipeline
3

muyangli's avatar
muyangli committed
4
from nunchaku import NunchakuSanaTransformer2DModel
muyangli's avatar
muyangli committed
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

transformer = NunchakuSanaTransformer2DModel.from_pretrained("mit-han-lab/svdq-int4-sana-1600m")
pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
    transformer=transformer,
    variant="bf16",
    torch_dtype=torch.bfloat16,
).to("cuda")
pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

prompt = "A cute 🐼 eating 🎋, ink drawing style"
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=4.5,
    num_inference_steps=20,
    generator=torch.Generator().manual_seed(42),
).images[0]

26
image.save("sana_1600m-int4.png")