sana1.6b.py 758 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
transformer = NunchakuSanaTransformer2DModel.from_pretrained(
    "mit-han-lab/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors"
)
muyangli's avatar
muyangli committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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]

28
image.save("sana1.6b-int4.png")