Rewrite oemer U-Net ConvTranspose pads before packaging.

ONNX Runtime 1.30+ rejects pads [0, 0, -1, -1] on the unet checkpoint. The build now rewrites those nodes to pads [0, 0, 0, 0] and output_padding [1, 1], which keeps the same output size. Version 0.1.1.

scripts/patch_oemer_unet.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Rewrite oemer U-Net ConvTranspose pads so ONNX Runtime can load the model.
+
+oemer 0.1.8's 1st_model.onnx ships ConvTranspose nodes with pads
+[0, 0, -1, -1]. ONNX Runtime 1.30 and newer reject those during shape
+inference. pads [0, 0, 0, 0] with output_padding [1, 1] is the same
+output size for these stride-2, kernel-1 nodes. Weights are left as they are.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import onnx
+from onnx import helper
+
+BAD_PADS = [0, 0, -1, -1]
+FIXED_PADS = [0, 0, 0, 0]
+FIXED_OUTPUT_PADDING = [1, 1]
+CRASH_NODE = "model/conv2d_transpose_1/conv2d_transpose"
+
+
+def _attr(node: onnx.NodeProto, name: str) -> onnx.AttributeProto | None:
+    for attr in node.attribute:
+        if attr.name == name:
+            return attr
+    return None
+
+
+def _ints(node: onnx.NodeProto, name: str) -> list[int] | None:
+    attr = _attr(node, name)
+    if attr is None:
+        return None
+    return list(attr.ints)
+
+
+def _set_ints(attr: onnx.AttributeProto, values: list[int]) -> None:
+    del attr.ints[:]
+    attr.ints.extend(values)
+
+
+def rewrite(model_path: Path) -> int:
+    model = onnx.load(model_path)
+    changed = 0
+    for node in model.graph.node:
+        if node.op_type != "ConvTranspose":
+            continue
+        pads = _ints(node, "pads")
+        if pads is None or pads == FIXED_PADS:
+            continue
+        if pads != BAD_PADS:
+            if any(value < 0 for value in pads):
+                raise SystemExit(
+                    f"{model_path}: {node.name} has unexpected negative pads {pads}"
+                )
+            continue
+        pads_attr = _attr(node, "pads")
+        if pads_attr is None:
+            raise SystemExit(f"{model_path}: {node.name} is missing pads")
+        _set_ints(pads_attr, FIXED_PADS)
+        output_padding = _attr(node, "output_padding")
+        if output_padding is None:
+            node.attribute.append(
+                helper.make_attribute("output_padding", FIXED_OUTPUT_PADDING)
+            )
+        else:
+            _set_ints(output_padding, FIXED_OUTPUT_PADDING)
+        changed += 1
+        print(
+            f"rewrote {node.name}: pads {FIXED_PADS} output_padding {FIXED_OUTPUT_PADDING}"
+        )
+    if changed:
+        onnx.save(model, model_path)
+    verify(model_path)
+    print(f"rewrote {changed} ConvTranspose node(s) in {model_path}")
+    return changed
+
+
+def verify(model_path: Path) -> None:
+    model = onnx.load(model_path)
+    found = False
+    for node in model.graph.node:
+        if node.op_type != "ConvTranspose":
+            continue
+        pads = _ints(node, "pads") or []
+        if any(value < 0 for value in pads):
+            raise SystemExit(f"{model_path}: {node.name} still has negative pads {pads}")
+        if node.name != CRASH_NODE:
+            continue
+        found = True
+        output_padding = _ints(node, "output_padding")
+        if pads != FIXED_PADS or output_padding != FIXED_OUTPUT_PADDING:
+            raise SystemExit(
+                f"{model_path}: {node.name} is {pads} / {output_padding}, "
+                f"expected {FIXED_PADS} / {FIXED_OUTPUT_PADDING}"
+            )
+    if not found:
+        raise SystemExit(f"{model_path}: missing {CRASH_NODE}")
+
+
+def main() -> int:
+    if len(sys.argv) != 2:
+        print("usage: patch_oemer_unet.py MODEL.onnx", file=sys.stderr)
+        return 2
+    path = Path(sys.argv[1])
+    if not path.is_file():
+        print(f"missing model: {path}", file=sys.stderr)
+        return 1
+    rewrite(path)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())