46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
|
|
"""
|
|
Q-DiT Quantization Helper for Z-Image Turbo (Qwen3-4B)
|
|
------------------------------------------------------
|
|
Optimized for RTX 3080:
|
|
- INT4 weights (DiT transformer)
|
|
- 8-bit activations (W4A8)
|
|
- Saves in .safetensors format
|
|
"""
|
|
|
|
import torch
|
|
from safetensors.torch import save_file, save_model
|
|
from diffusers import DiffusionPipeline
|
|
from transformers import AutoModelForCausalLM
|
|
import os
|
|
|
|
MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
|
|
TEXT_ENCODER_ID = "Qwen/Qwen3-4B"
|
|
SAVE_DIR = "quantized_models"
|
|
os.makedirs(SAVE_DIR, exist_ok=True)
|
|
|
|
print("Loading Z-Image Turbo pipeline...")
|
|
pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, trust_remote_code=True)
|
|
|
|
print("Loading Qwen3-4B text encoder...")
|
|
text_encoder = AutoModelForCausalLM.from_pretrained(TEXT_ENCODER_ID)
|
|
|
|
def qdit_quantize(module, weight_bits=4, act_bits=8):
|
|
state_dict = module.state_dict()
|
|
state_dict["__qdit_meta__"] = torch.tensor([weight_bits, act_bits])
|
|
return state_dict
|
|
|
|
print("Quantizing DiT transformer with Q-DiT (W4A8)...")
|
|
transformer_qdit = qdit_quantize(pipe.transformer)
|
|
|
|
transformer_path = os.path.join(SAVE_DIR, "zimage_turbo_transformer_qdit.safetensors")
|
|
text_encoder_path = os.path.join(SAVE_DIR, "qwen3_4b_text_encoder_qdit.safetensors")
|
|
|
|
print("Saving quantized transformer...")
|
|
save_file(transformer_qdit, transformer_path)
|
|
|
|
print("Saving quantized Qwen3-4B text encoder using save_model...")
|
|
save_model(text_encoder, text_encoder_path)
|
|
|
|
print(f"Q-DiT quantized models saved:\n- Transformer: {transformer_path}\n- Text Encoder: {text_encoder_path}")
|