Cursor Agent
6da116518f
AppImage release / appimage (push) Successful in 1m42s
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.
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
#!/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())
|