74 lines
3.2 KiB
Python
74 lines
3.2 KiB
Python
import torch
|
|
import os
|
|
import quanto
|
|
from diffusers import DiffusionPipeline, Transformer2DModel
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from safetensors.torch import save_file
|
|
|
|
def quantize_and_save():
|
|
MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
|
|
SAVE_DIR = "models/quantized_models"
|
|
os.makedirs(SAVE_DIR, exist_ok=True)
|
|
|
|
print(f"Starting quantization for {MODEL_ID}...")
|
|
|
|
# 1. Quantize Text Encoder (Qwen3-4B)
|
|
print("Loading Text Encoder...")
|
|
try:
|
|
text_encoder = AutoModelForCausalLM.from_pretrained(MODEL_ID, subfolder="text_encoder", trust_remote_code=True, dtype=torch.bfloat16)
|
|
except Exception:
|
|
print("Could not load from subfolder, trying direct Qwen/Qwen3-4B...")
|
|
text_encoder = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B", trust_remote_code=True, dtype=torch.bfloat16)
|
|
|
|
print("Quantizing Text Encoder (Weights: int4, Activations: int8)...")
|
|
# Collect modules first to avoid modifying graph while iterating (prevents RecursionError)
|
|
modules_to_quantize = []
|
|
for name, module in text_encoder.named_modules():
|
|
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
|
|
modules_to_quantize.append(module)
|
|
|
|
for module in modules_to_quantize:
|
|
quanto.quantize(module, weights=quanto.qint4, activations=quanto.qint8)
|
|
quanto.freeze(text_encoder)
|
|
|
|
te_path = os.path.join(SAVE_DIR, "qwen_text_encoder_qdit.safetensors")
|
|
print(f"Saving Text Encoder to {te_path}...")
|
|
# Save state dict, filtering out non-tensor values (quanto metadata)
|
|
te_state_dict = {k: v for k, v in text_encoder.state_dict().items() if isinstance(v, torch.Tensor)}
|
|
save_file(te_state_dict, te_path)
|
|
del text_encoder
|
|
del te_state_dict
|
|
|
|
# 2. Quantize Transformer
|
|
print("Loading Transformer...")
|
|
# Try to load ZImageTransformer2DModel
|
|
try:
|
|
from diffusers import ZImageTransformer2DModel
|
|
transformer = ZImageTransformer2DModel.from_pretrained(MODEL_ID, subfolder="transformer", torch_dtype=torch.bfloat16)
|
|
except ImportError:
|
|
print("ZImageTransformer2DModel not found. Using Transformer2DModel (might fail if architecture is custom)...")
|
|
transformer = Transformer2DModel.from_pretrained(MODEL_ID, subfolder="transformer", torch_dtype=torch.bfloat16)
|
|
|
|
print("Quantizing Transformer (Weights: int4, Activations: int8)...")
|
|
# Selective quantization to avoid LayerNorm crash and improve quality
|
|
modules_to_quantize = []
|
|
for name, module in transformer.named_modules():
|
|
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
|
|
modules_to_quantize.append(module)
|
|
|
|
for module in modules_to_quantize:
|
|
quanto.quantize(module, weights=quanto.qint4, activations=quanto.qint8)
|
|
quanto.freeze(transformer)
|
|
|
|
tr_path = os.path.join(SAVE_DIR, "zimage_turbo_transformer_qdit.safetensors")
|
|
print(f"Saving Transformer to {tr_path}...")
|
|
tr_state_dict = {k: v for k, v in transformer.state_dict().items() if isinstance(v, torch.Tensor)}
|
|
save_file(tr_state_dict, tr_path)
|
|
del transformer
|
|
del tr_state_dict
|
|
|
|
print("Quantization complete!")
|
|
|
|
if __name__ == "__main__":
|
|
quantize_and_save()
|