File size: 2,018 Bytes
926d100 ae3512a 926d100 ae3512a 926d100 ae3512a 926d100 b1b85b3 926d100 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | #!/usr/bin/env python3
"""Rebuild the OrbitQuant components of Turbo-Image-2.1 from WaveCut/Turbo-Image-2.1.
python quantize.py --out ./turbo-image-2.1-orbitquant
"""
import argparse
import os
from pathlib import Path
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import orbitquant
import torch
from diffusers import QwenImage21Transformer2DModel
from orbitquant import recipe
from transformers import Qwen3VLForConditionalGeneration
SOURCE = "WaveCut/Turbo-Image-2.1"
DIT_KEEP = ["img_in", "txt_in", "modulation.1"] + [
f"transformer_blocks.{i}.img_mlp.out" for i in (0, 4, 5, 6, 7, 8, 9, 10, 11, 13)
]
TE_KEEP = [f"model.language_model.layers.{i}." for i in (6, 16, 34, 35)]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--source", default=SOURCE)
ap.add_argument("--revision")
ap.add_argument("--out", type=Path, required=True)
args = ap.parse_args()
dit_cfg = recipe("w4a4", target_policy="universal", runtime_mode="auto_fused",
modules_to_not_convert=DIT_KEEP)
dit = QwenImage21Transformer2DModel.from_pretrained(
args.source, subfolder="transformer", revision=args.revision, quantization_config=dit_cfg,
dtype=torch.float16, low_cpu_mem_usage=True, quantization_device="cuda")
dit.save_pretrained(args.out / "transformer", safe_serialization=True, max_shard_size="4GB")
del dit
torch.cuda.empty_cache()
te_cfg = recipe("w4a4", weight_bits=6, activation_bits=6, target_policy="universal", runtime_mode="auto_fused",
modules_to_not_convert=TE_KEEP, modules_to_convert=["*"])
te = Qwen3VLForConditionalGeneration.from_pretrained(
args.source, subfolder="text_encoder", revision=args.revision, quantization_config=te_cfg,
dtype=torch.float16, low_cpu_mem_usage=True)
te.save_pretrained(args.out / "text_encoder", safe_serialization=True, max_shard_size="4GB")
print(f"saved to {args.out}")
if __name__ == "__main__":
main()
|