#!/usr/bin/env python3
"""
Quantize ISNet ONNX model from fp32 to INT8.

Usage:
    pip install -r requirements.txt
    python quantize.py

Output: onnx/model_int8.onnx (~43MB, down from ~168MB fp32)
"""

import os
import sys

try:
    from onnxruntime.quantization import quantize_dynamic, QuantType
except ImportError:
    print("Error: onnxruntime not installed. Run: pip install -r requirements.txt")
    sys.exit(1)


def main():
    input_path = os.path.join("onnx", "model.onnx")
    output_path = os.path.join("onnx", "model_int8.onnx")

    if not os.path.exists(input_path):
        print(f"Error: {input_path} not found.")
        print("Download the fp32 model first:")
        print("  huggingface-cli download imgly/isnet-general-onnx onnx/model.onnx --local-dir .")
        sys.exit(1)

    input_size = os.path.getsize(input_path) / 1024 / 1024
    print(f"Input:  {input_path} ({input_size:.1f} MB)")
    print(f"Output: {output_path}")
    print("Quantizing to INT8...")

    quantize_dynamic(
        model_input=input_path,
        model_output=output_path,
        weight_type=QuantType.QInt8,
    )

    output_size = os.path.getsize(output_path) / 1024 / 1024
    reduction = (1 - output_size / input_size) * 100
    print(f"Done! {output_size:.1f} MB ({reduction:.0f}% smaller)")


if __name__ == "__main__":
    main()