Add DDColor paper tiny INT8 (default colouriser), DDColor build scripts, update model card

README.md

@@ -35,14 +35,25 @@ Pin a commit in your download URL (`/resolve/<commit>/<file>`) and verify the SH
 | `u2netp.onnx` | 4,574,861 | `309c8469258dda742793dce0ebea8e6dd393174f89934733ecc8b14c76f4ddd8` | U²-Netp (small) salient-object matte (default background model) | Unmodified copy of `u2netp.onnx` from [danielgatis/rembg release v0.0.0](https://github.com/danielgatis/rembg/releases/tag/v0.0.0) (rembg is MIT), converted from the weights of [xuebinqin/U-2-Net](https://github.com/xuebinqin/U-2-Net) | Apache-2.0, © Xuebin Qin et al. |
 | `u2net.onnx` | 175,997,641 | `8d10d2f3bb75ae3b6d527c77944fc5e7dcd94b29809d47a739a7a728a912b491` | U²-Net (full) salient-object matte (HD option) | Unmodified copy of `u2net.onnx` from [danielgatis/rembg release v0.0.0](https://github.com/danielgatis/rembg/releases/tag/v0.0.0) (rembg is MIT), converted from the weights of [xuebinqin/U-2-Net](https://github.com/xuebinqin/U-2-Net) | Apache-2.0, © Xuebin Qin et al. |
 | `ddcolor-paper-tiny.onnx` | 220,552,834 | `5479b1b00d8ba2b3c312e5378e75dc93eac6f80c8497e90169801896e1eeb68d` | DDColor-T (`ddcolor_paper_tiny`), ConvNeXt-T encoder + multi-scale colour decoder, fp32, opset 17 (photo colourisation) | Built by MomoApps from the official weights [piddnad/ddcolor_paper_tiny](https://huggingface.co/piddnad/ddcolor_paper_tiny) (commit `cf9fd99c1d7472689ec7413441c1b799a51866a3`, `pytorch_model.bin` sha256 `8a1277bc90a1bfbb6d2d83933a9a6bc821931879ca93e26e4fcec12165d41fce`) and the architecture in [piddnad/DDColor](https://github.com/piddnad/DDColor) (commit `2adb63f2656ac41cbdf7b894cddd94121a3faf13`). The weights load with no missing or unexpected keys. The 8 spectral-norm layers are folded into plain weights, which is exact in eval mode. Exported with `torch.onnx.export` (torch 2.15.0). | Apache-2.0, © Xiaoyang Kang et al. (DAMO Academy, Alibaba) |
+| `ddcolor-paper-tiny-int8.onnx` | 56,619,279 | `29f3a5f5f794959faebaf034a142c34712833623f9ce672f6ef8f0a8ae8e5f5a` | DDColor-T, weight-only per-channel INT8 (default colouriser) | Built by MomoApps from `ddcolor-paper-tiny.onnx` above with `scripts/quantize_ddcolor_int8.py`. The 161 Conv/MatMul/Gemm weight tensors are stored as symmetric INT8 with one scale per output channel and expanded by `DequantizeLinear`. All compute stays float32, and Einsum, Erf, LayerNormalization, norms and biases stay in float. Against the float32 ONNX on two DDColor test photos, the mean ab difference is 0.41 / 0.32. Visually indistinguishable apart from edge artifacts the float model already has. | Apache-2.0, © Xiaoyang Kang et al. (DAMO Academy, Alibaba) |
 
 ## Input and output contracts
 
 - **GFPGAN (both files):** input `input`, float32 `[1,3,512,512]` RGB in −1..1 (one aligned face crop). Output `output`, the same shape and range.
 - **Real-ESRGAN general x4v3:** input `input`, float32 `[1,3,H,W]` RGB in 0..1, with dynamic H and W. Output `output`, `[1,3,4H,4W]`. Clamp the output to 0..1.
-- **DDColor paper tiny:** input `input`, float32 `[1,3,512,512]` RGB in 0..1. Build it as DDColor's `ColorizationPipeline` does: resize the photo to 512×512, convert to Lab (OpenCV float), keep L, set a = b = 0, and convert back to RGB. The result is the greyscale L channel replicated as RGB. Output `output`, float32 `[1,2,512,512]`: the predicted Lab a and b channels on OpenCV's float Lab scale (roughly −128..127; clamp them). Resize a and b to the original size, combine them with the original L, and convert Lab to RGB. Ops include `Einsum`, `Erf` and `LayerNormalization`, which the full ONNX Runtime Android package supports. A reduced-operator ORT Mobile build must include them.
+- **DDColor paper tiny (both files):** input `input`, float32 `[1,3,512,512]` RGB in 0..1. Build it as DDColor's `ColorizationPipeline` does: resize the photo to 512×512, convert to Lab (OpenCV float), keep L, set a = b = 0, and convert back to RGB. The result is the greyscale L channel replicated as RGB. Output `output`, float32 `[1,2,512,512]`: the predicted Lab a and b channels on OpenCV's float Lab scale (roughly −128..127; clamp them). Resize a and b to the original size, combine them with the original L, and convert Lab to RGB. Ops include `Einsum`, `Erf` and `LayerNormalization`, which the full ONNX Runtime Android package supports. A reduced-operator ORT Mobile build must include them. The INT8 file also uses `DequantizeLinear` (opset 13+, per-axis). ONNX Runtime's mobile checker recommends the CPU execution provider for both files, and both convert to `.ort` format.
 - **U²-Net / U²-Netp:** input `input.1`, float32 `[1,3,320,320]` RGB, with ImageNet normalisation applied after dividing by the peak pixel. The first output is the matte.
 
+## Build scripts (DDColor)
+
+`scripts/` holds the scripts that produced the DDColor files. They are deterministic: re-running them on the box gave byte-identical files.
+
+- `export_ddcolor_onnx.py DDCOLOR_REPO WEIGHTS_DIR OUT.onnx TEST_IMAGE [PREVIEW.jpg]` exports the float32 ONNX (opset 17) from the official weights. It folds spectral norm into plain weights and checks the result against PyTorch: max |Δab| was 0.0022 on the Audrey Hepburn test photo.
+- `quantize_ddcolor_int8.py SRC.onnx DEST.onnx` produces the weight-only per-channel INT8 file.
+- `compare_ddcolor.py REF.onnx TEST.onnx OUTDIR IMG...` reports the ab differences and writes side-by-side colourised JPGs.
+
+The tools used were torch 2.15.0 (CPU), onnx 1.23.2, onnxruntime 1.31.1 and opencv-python-headless.
+
 ## Training data note: FFHQ
 
 GFPGAN was trained on the [FFHQ](https://github.com/NVlabs/ffhq-dataset) face dataset. FFHQ's images are published under a

ddcolor-paper-tiny-int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29f3a5f5f794959faebaf034a142c34712833623f9ce672f6ef8f0a8ae8e5f5a
+size 56619279

scripts/compare_ddcolor.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+"""Compare two DDColor ONNX files on photos: ab max/mean difference plus side-by-side colourised JPGs.
+Usage: compare_ddcolor.py REF.onnx TEST.onnx OUTDIR IMG [IMG ...]"""
+import sys, time
+from pathlib import Path
+import cv2, numpy as np, onnxruntime as ort
+
+def preprocess(img_bgr):  # DDColor ColorizationPipeline preprocessing
+    img = (img_bgr / 255.0).astype(np.float32)
+    l = cv2.cvtColor(cv2.resize(img, (512, 512)), cv2.COLOR_BGR2Lab)[:, :, :1]
+    rgb = cv2.cvtColor(np.concatenate((l, 0 * l, 0 * l), -1), cv2.COLOR_LAB2RGB)
+    return rgb.transpose(2, 0, 1)[None].astype(np.float32)
+
+def colourise(img_bgr, ab):
+    h, w = img_bgr.shape[:2]
+    l = cv2.cvtColor((img_bgr / 255.0).astype(np.float32), cv2.COLOR_BGR2Lab)[:, :, :1]
+    ab = cv2.resize(ab[0].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_NEAREST)
+    return (cv2.cvtColor(np.concatenate((l, ab), -1), cv2.COLOR_LAB2BGR) * 255).round().clip(0, 255).astype(np.uint8)
+
+ref_s = ort.InferenceSession(sys.argv[1], providers=["CPUExecutionProvider"])
+tst_s = ort.InferenceSession(sys.argv[2], providers=["CPUExecutionProvider"])
+out = Path(sys.argv[3]); out.mkdir(parents=True, exist_ok=True)
+for p in sys.argv[4:]:
+    img = cv2.imread(p); x = preprocess(img)
+    a = ref_s.run(["output"], {"input": x})[0]
+    t = time.time(); b = tst_s.run(["output"], {"input": x})[0]; dt = time.time() - t
+    d = np.abs(a - b)
+    ca, cb = colourise(img, a), colourise(img, b)
+    pix = np.abs(ca.astype(int) - cb.astype(int))
+    print(f"{Path(p).name}: ab max {d.max():.3f} mean {d.mean():.4f} p99 {np.percentile(d, 99):.3f} | "
+          f"8-bit RGB max {pix.max()} mean {pix.mean():.3f} | test run {dt:.2f}s")
+    stem = Path(p).stem.split()[0]
+    gap = np.full((ca.shape[0], 16, 3), 255, np.uint8)
+    side = np.concatenate((ca, gap, cb), 1)
+    s = 1400 / side.shape[1]
+    cv2.imwrite(str(out / f"{stem}_fp32_vs_test.jpg"), cv2.resize(side, None, fx=s, fy=s, interpolation=cv2.INTER_AREA), [cv2.IMWRITE_JPEG_QUALITY, 92])

scripts/export_ddcolor_onnx.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+"""Export DDColor paper_tiny (Apache-2.0, piddnad/DDColor) to the OptiMix ONNX contract.
+
+Weights: huggingface.co/piddnad/ddcolor_paper_tiny @ cf9fd99c1d7472689ec7413441c1b799a51866a3, pytorch_model.bin
+         sha256 8a1277bc90a1bfbb6d2d83933a9a6bc821931879ca93e26e4fcec12165d41fce
+Code:    github.com/piddnad/DDColor @ 2adb63f2656ac41cbdf7b894cddd94121a3faf13 (basicsr/archs/ddcolor_arch.py)
+Contract: input  `input`  float32 [1,3,512,512] RGB 0..1, the L channel rendered as grey RGB
+          output `output` float32 [1,2,512,512] Lab a,b (OpenCV float Lab scale)
+"""
+import hashlib, json, sys
+from pathlib import Path
+import cv2, numpy as np, onnx, onnxruntime as ort, torch
+
+# Usage: export_ddcolor_onnx.py DDCOLOR_REPO_DIR WEIGHTS_DIR OUT.onnx TEST_IMAGE [PREVIEW.jpg]
+#   DDCOLOR_REPO_DIR: a checkout of github.com/piddnad/DDColor (commit above)
+#   WEIGHTS_DIR: holds config.json and pytorch_model.bin from piddnad/ddcolor_paper_tiny (commit above)
+REPO = Path(sys.argv[1]).resolve(); sys.path.insert(0, str(REPO))
+from basicsr.archs.ddcolor_arch import DDColor
+
+WEIGHTS = Path(sys.argv[2])
+OUT = Path(sys.argv[3])
+TEST_IMAGE = sys.argv[4]
+PREVIEW = sys.argv[5] if len(sys.argv) > 5 else None
+OPSET = 17
+cfg = json.loads((WEIGHTS / "config.json").read_text())
+model = DDColor(**cfg)
+state = torch.load(WEIGHTS / "pytorch_model.bin", map_location="cpu", weights_only=True)
+state = state.get("params", state)
+missing, unexpected = model.load_state_dict(state, strict=False)
+print("missing", missing, "unexpected", unexpected)
+assert not unexpected, unexpected
+assert all(k.endswith("num_batches_tracked") for k in missing), missing
+model.eval()
+
+# Reference = the official model exactly as loaded (spectral-norm hooks intact).
+img = cv2.imread(TEST_IMAGE)
+
+def preprocess(img_bgr):
+    img = (img_bgr / 255.0).astype(np.float32)
+    r = cv2.resize(img, (512, 512))
+    l = cv2.cvtColor(r, cv2.COLOR_BGR2Lab)[:, :, :1]
+    lab = np.concatenate((l, np.zeros_like(l), np.zeros_like(l)), axis=-1)
+    rgb = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
+    return rgb.transpose(2, 0, 1)[None].astype(np.float32)
+
+x = preprocess(img)
+with torch.no_grad():
+    ref = model(torch.from_numpy(x)).numpy()
+
+# Bake spectral norm into plain weights (W / sigma with the stored u, v; eval mode
+# runs no power iteration), so the graph does not carry weight_orig plus a reshaped copy.
+from torch.nn.utils import remove_spectral_norm
+baked = 0
+for mod in model.modules():
+    if hasattr(mod, "weight_orig"):
+        remove_spectral_norm(mod); baked += 1
+print("spectral norms baked", baked)
+with torch.no_grad():
+    torch.onnx.export(model, torch.rand(1, 3, 512, 512), str(OUT), opset_version=OPSET,
+                      input_names=["input"], output_names=["output"], dynamo=False,
+                      do_constant_folding=True)
+m = onnx.load(str(OUT)); onnx.checker.check_model(m)
+print("opset", [(o.domain, o.version) for o in m.opset_import])
+print("inputs", [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in m.graph.input])
+print("outputs", [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim]) for o in m.graph.output])
+
+with torch.no_grad():
+    baked_ref = model(torch.from_numpy(x)).numpy()
+print("baked vs official torch max abs diff", float(np.abs(baked_ref - ref).max()))
+sess = ort.InferenceSession(str(OUT), providers=["CPUExecutionProvider"])
+got = sess.run(["output"], {"input": x})[0]
+print("shape", got.shape, "ref range", float(ref.min()), float(ref.max()))
+print("max abs diff", float(np.abs(got - ref).max()), "mean abs diff", float(np.abs(got - ref).mean()))
+# Colourise the test image with the ONNX output, as DDColor's pipeline does.
+h, w = img.shape[:2]
+orig_l = cv2.cvtColor((img / 255.0).astype(np.float32), cv2.COLOR_BGR2Lab)[:, :, :1]
+ab = cv2.resize(got[0].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_NEAREST)
+out = cv2.cvtColor(np.concatenate((orig_l, ab), -1), cv2.COLOR_LAB2BGR)
+if PREVIEW:
+    cv2.imwrite(PREVIEW, (out * 255).round().clip(0, 255).astype(np.uint8))
+print("size", OUT.stat().st_size, "sha256", hashlib.sha256(OUT.read_bytes()).hexdigest())

scripts/quantize_ddcolor_int8.py

@@ -0,0 +1,66 @@
+#!/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)")