# quantize_onnx.py
from pathlib import Path
from onnxruntime.quantization import quantize_dynamic, QuantType
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--in_dir", default="onnx_models")
parser.add_argument("--out_dir", default="onnx_models_quant")
args = parser.parse_args()

Path(args.out_dir).mkdir(parents=True, exist_ok=True)

for fname in ["text_encoder.onnx", "unet.onnx", "vae_decoder.onnx"]:
    in_path = Path(args.in_dir) / fname
    out_path = Path(args.out_dir) / fname
    if not in_path.exists():
        print("Missing:", in_path)
        continue
    print("Quantizing", in_path, "->", out_path)
    # Dynamic quantization (weights -> INT8 where supported)
    quantize_dynamic(str(in_path), str(out_path), weight_type=QuantType.QUInt8)
    print("Wrote", out_path)