#!/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()