# int8 weight-only (linear, per-channel) quantisation of the Core ML models; activations stay fp16.
import coremltools as ct, coremltools.optimize.coreml as cto
import os
B = os.environ.get("NEMU_OCR_WORKDIR", ".")  # holds models/ (upstream weights), coreml/ (outputs)
cfg = cto.OptimizationConfig(global_config=cto.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8", granularity="per_channel"))
for name in ["mangaocr-encoder", "mangaocr-decoder"]:
    m = ct.models.MLModel(f"{B}/coreml/{name}-fp16.mlpackage")
    q = cto.linear_quantize_weights(m, config=cfg); q.save(f"{B}/coreml/{name}-int8.mlpackage"); print(name, "ok")