Views
No views yet
1pip install git+https://github.com/xin3he/diffusers@auto_round # A special version for auto-round
2pip install git+https://github.com/intel/auto-round.git@refs/pull/1640/head # A special version for tuning and loading1import time
2
3import torch
4from diffusers import ZImagePipeline
5
6model_path = "INCModel/Z-Image-W4A16-AutoRound"
7
8# 1. Load the pipeline
9# Use bfloat16 for optimal performance on supported GPUs
10pipe = ZImagePipeline.from_pretrained(
11 model_path,
12 torch_dtype=torch.bfloat16,
13 device_map="cuda", # set device_map="cuda" to leverage gptqmodel kernels
14)
15pipe.to("cuda")
16
17# [Optional] Attention Backend
18# Diffusers uses SDPA by default. Switch to Flash Attention for better efficiency if supported:
19# pipe.transformer.set_attention_backend("flash") # Enable Flash-Attention-2
20# pipe.transformer.set_attention_backend("_flash_3") # Enable Flash-Attention-3
21
22# [Optional] Model Compilation
23# Compiling the DiT model accelerates inference, but the first run will take longer to compile.
24# pipe.transformer.compile()
25
26# [Optional] CPU Offloading
27# Enable CPU offloading for memory-constrained devices.
28# pipe.enable_model_cpu_offload()
29
30prompt = "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."
31
32# 2. Generate Image
33print("Starting generation...")
34image = pipe(
35 prompt=prompt,
36 negative_prompt="",
37 height=1280,
38 width=720,
39 cfg_normalization=False,
40 num_inference_steps=50,
41 guidance_scale=4,
42 generator=torch.Generator("cuda").manual_seed(42),
43).images[0]
44
45image.save("example.png")CUDA_VISIBLE_DEVICES=6,7 auto-round Tongyi-MAI/Z-Image --num_inference_steps 2 --output_dir /workspace/models/Tongyi-MAI/Z-Image-W4A16-AutoRound --device_map 0,1 --low_gpu_mem_usage --enable_torch_compile