#!/usr/bin/env python3
"""Weight-only per-channel INT8 for ddcolor-paper-tiny.onnx (OptiMix).

Only weight initializers consumed by Conv, MatMul or Gemm are quantised (symmetric INT8, one scale per
output channel). Each becomes int8 data + float scale + DequantizeLinear(axis) feeding the original
tensor name, so compute stays float32. Einsum, Erf, LayerNormalization, norms, biases and small tensors
are left untouched in float32. DequantizeLinear with axis needs opset >= 13; the graph is opset 17.

Usage: quantize_ddcolor_int8.py SRC.onnx DEST.onnx
"""
import hashlib, sys
import numpy as np, onnx
from onnx import helper, numpy_helper

src, dest = sys.argv[1], sys.argv[2]
m = onnx.load(src)
g = m.graph
inits = {i.name: i for i in g.initializer}
consumers = {}
for n in g.node:
    for k, name in enumerate(n.input):
        consumers.setdefault(name, []).append((n, k))

def axis_for(name, arr):
    uses = consumers.get(name, [])
    if not uses:
        return None
    axes = set()
    for n, k in uses:
        if n.op_type == "Conv" and k == 1 and arr.ndim == 4:
            axes.add(0)                      # [out, in, kh, kw]
        elif n.op_type == "MatMul" and k == 1 and arr.ndim == 2:
            axes.add(1)                      # B is [K, N]; channel = N
        elif n.op_type == "Gemm" and k == 1 and arr.ndim == 2:
            tb = next((a.i for a in n.attribute if a.name == "transB"), 0)
            axes.add(0 if tb else 1)
        else:
            return None                      # any other consumer: keep float
    return axes.pop() if len(axes) == 1 else None

new_inits, new_nodes, done, kept = [], [], 0, 0
for init in list(g.initializer):
    arr = numpy_helper.to_array(init)
    ax = axis_for(init.name, arr) if arr.dtype == np.float32 and arr.size >= 4096 else None
    if ax is None:
        new_inits.append(init); kept += arr.nbytes; continue
    red = tuple(i for i in range(arr.ndim) if i != ax)
    scale = np.maximum(np.abs(arr).max(axis=red) / 127.0, 1e-12).astype(np.float32)
    shp = [1] * arr.ndim; shp[ax] = -1
    q = np.clip(np.round(arr / scale.reshape(shp)), -127, 127).astype(np.int8)
    new_inits += [numpy_helper.from_array(q, init.name + "_q"),
                  numpy_helper.from_array(scale, init.name + "_scale"),
                  numpy_helper.from_array(np.zeros_like(scale, dtype=np.int8), init.name + "_zp")]
    new_nodes.append(helper.make_node("DequantizeLinear",
        [init.name + "_q", init.name + "_scale", init.name + "_zp"], [init.name],
        name=init.name + "_dq", axis=ax))
    done += 1

del g.initializer[:]; g.initializer.extend(new_inits)
nodes = new_nodes + list(g.node)
del g.node[:]; g.node.extend(nodes)
onnx.checker.check_model(m)
onnx.save(m, dest)
data = open(dest, "rb").read()
print(f"quantised {done} weights; float32 kept {kept/1e6:.1f} MB")
print(f"{hashlib.sha256(data).hexdigest()}  {dest} ({len(data)} bytes)")