|
|
| """
|
| inline_onnx.py
|
| ==============
|
|
|
| 把 `.model4` 里的 ONNX 图**内联成独立可跑的 .onnx**（不依赖 iZotope 运行时）。
|
|
|
| 背景
|
| ----
|
| ONNX 图里大部分权重是 `external_data`，`location` 写成 `model4_<hex(张量名)>.bin`——
|
| 这些文件在磁盘上根本不存在，是 iZotope 的 Photon 在运行时从 safetensors 载荷里喂进去的。
|
| 所以抽出来的 .onnx 用 onnxruntime 打不开（`should be stored in model4_xxx.bin, but it is not regular file`）。
|
|
|
| 本脚本用容器自带的 `graphs[backend=onnx].bindings.tensors` 表把
|
| `backend_name -> 容器张量名` 解析出来，把字节写回 initializer，输出标准 .onnx。
|
|
|
| 张量来源（三级回退）
|
| --------------------
|
| 1. 容器里同名张量（规范名）
|
| 2. `bindings.tensors` 里的 `backend_name` 反查
|
| 3. `onnx.source.<sha>.<name>` 去掉前缀后的同名张量
|
|
|
| 用法
|
| ----
|
| python inline_onnx.py # 处理非 pymss 的 10 个模型
|
| python inline_onnx.py --all # 连 pymss_models 里的 4 个也一起导
|
| python inline_onnx.py --only indoor
|
| """
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import json
|
| import os
|
| import re
|
| import struct
|
| import sys
|
|
|
| import numpy as np
|
| import onnx
|
| import torch
|
|
|
| import model_io
|
| from onnx import numpy_helper
|
|
|
| HERE = os.path.dirname(os.path.abspath(__file__))
|
| SF_DIR = os.path.join(HERE, "..", "raw", "safetensors")
|
| BLOBS_DIR = os.path.join(HERE, "..", "raw", "blobs")
|
| CKPT_DIR = os.path.join(HERE, "..", "raw", "ckpt")
|
| OUT_DIR = os.path.join(HERE, "..", "onnx")
|
|
|
|
|
| PYMSS_OK = {
|
| "voice2stem_v03_fp16c_t2", "voicesep2stem_fp16c_t2",
|
| "drumsep_1.3", "music4stemhq_fp16c_t2",
|
| }
|
|
|
|
|
| def read_header(model: str) -> dict:
|
| p = model_io.safetensors_path(model)
|
| with open(p, "rb") as f:
|
| n = struct.unpack("<Q", f.read(8))[0]
|
| hdr = json.loads(f.read(n))
|
| return json.loads(hdr["__metadata__"]["model4"])
|
|
|
|
|
| def container_tensors(model: str) -> dict:
|
| sd = model_io.state_dict_for(model)
|
| out = {}
|
| for k, v in sd.items():
|
| if k.startswith("__blob__"):
|
| continue
|
| out[k] = v
|
|
|
| m = re.match(r"^(?:onnx|coreml)\.source\.[0-9a-f]+\.(.*)$", k)
|
| if m:
|
| out.setdefault(m.group(1), v)
|
| m = re.match(r"^coreml\.fp16\.(.*)_to_fp16$", k)
|
| if m:
|
| out.setdefault(m.group(1), v)
|
| return out
|
|
|
|
|
| def norm_name(k: str) -> str:
|
| n = k.replace("/", ".").replace("..", ".").strip(".")
|
| n = re.sub(r"^model\.", "", n)
|
| n = re.sub(r"^__onnx__\.", "", n)
|
| return n
|
|
|
|
|
| def build_lookup(tensors: dict) -> dict:
|
| """归一化名 -> tensor，供名字只差分隔符/前缀时兜底。"""
|
| out = {}
|
| for k, v in tensors.items():
|
| out.setdefault(norm_name(k), v)
|
| return out
|
|
|
|
|
| def binding_map(meta: dict) -> dict:
|
| """backend_name -> 容器张量名"""
|
| g = next((x for x in meta.get("graphs", []) if x.get("backend") == "onnx"), None)
|
| if not g:
|
| return {}
|
| out = {}
|
| for b in g.get("bindings") or []:
|
| for tname, v in (b.get("tensors") or {}).items():
|
| bn = v.get("backend_name")
|
| if bn and bn not in out:
|
| out[bn] = tname
|
| return out
|
|
|
|
|
| def pick_onnx_blob(model: str):
|
| bd = model_io.blobs_dir(model)
|
| if not os.path.isdir(bd):
|
| return None
|
| cands = sorted((p for p in os.listdir(bd) if p.endswith(".onnx")),
|
| key=lambda n: os.path.getsize(os.path.join(bd, n)), reverse=True)
|
| for n in cands:
|
| p = os.path.join(bd, n)
|
| try:
|
| m = onnx.load(p, load_external_data=False)
|
| except Exception:
|
| continue
|
| if len(m.graph.initializer) > 0:
|
| return p, m
|
| return None, None
|
|
|
|
|
| def inline(model: str, tensors: dict, meta: dict):
|
| path, m = pick_onnx_blob(model)
|
| if m is None:
|
| return None, {"error": "没有可解析的 onnx blob"}
|
|
|
| bmap = binding_map(meta)
|
| look = build_lookup(tensors)
|
|
|
|
|
| try:
|
| sys.path.insert(0, HERE)
|
| from to_pymss import load_canon, map_bs_roformer
|
| ck = model_io.model_path_for(model)
|
| if ck and os.path.exists(ck):
|
| canon, al = load_canon(ck)
|
| if any(k.startswith("band_split") for k in canon):
|
| sd, _p, _u = map_bs_roformer(canon, al)
|
| for k, v in sd.items():
|
| look.setdefault(norm_name("model." + k), v)
|
|
|
| m2 = re.match(r"^(layers\.\d+\.\d+\.layers\.\d+\.\d+)\.(.+)$", k)
|
| if m2:
|
| look.setdefault(norm_name(f"model.{m2.group(1)}.branch.{m2.group(2)}"), v)
|
| except Exception as e:
|
| print(" (roformer 还原跳过:", e, ")")
|
| aliases = meta.get("tensor_aliases") or {}
|
|
|
| def deref(n):
|
| """沿 tensor_aliases 追到实际存储的张量名。"""
|
| seen = set()
|
| while isinstance(n, str) and n in aliases and n not in seen:
|
| seen.add(n)
|
| n = (aliases[n] or {}).get("tensor", n)
|
| return n
|
| need = [i for i in m.graph.initializer if any(e.key == "location" for e in i.external_data)]
|
| resolved = missing = 0
|
| for init in m.graph.initializer:
|
| loc = next((e.value for e in init.external_data if e.key == "location"), None)
|
| if loc is None:
|
| continue
|
| name = init.name
|
| t = tensors.get(name)
|
| if t is None and name in bmap:
|
| t = tensors.get(bmap[name])
|
| if t is None:
|
| try:
|
| dec = bytes.fromhex(loc.split("model4_")[-1].replace(".bin", "")).decode()
|
| t = tensors.get(dec) or tensors.get(bmap.get(dec, ""))
|
| except Exception:
|
| t = None
|
| if t is None:
|
|
|
|
|
| m2 = re.match(r"^model\.(band_split\.norm_fc_modules\.\d+)\.combined\.([01])\.(\w+)$", name)
|
| if m2:
|
| sub = "norm" if m2.group(2) == "0" else "fc"
|
| t = look.get(norm_name(f"{m2.group(1)}.{sub}.{m2.group(3)}"))
|
| if t is None:
|
| t = look.get(norm_name(deref(name)))
|
| if t is None:
|
| t = look.get(norm_name(name))
|
| if t is None and name in bmap:
|
|
|
|
|
| t = look.get(norm_name(deref(bmap[name])))
|
| if t is None and name in bmap:
|
| t = look.get(norm_name(bmap[name]))
|
| if t is None:
|
| missing += 1
|
| continue
|
| want = onnx.helper.tensor_dtype_to_np_dtype(init.data_type)
|
| arr = t.detach().cpu().numpy()
|
| if arr.dtype != want:
|
| arr = arr.astype(want, copy=False)
|
| if tuple(init.dims) and tuple(arr.shape) != tuple(init.dims):
|
| if arr.size != int(np.prod(init.dims)):
|
| missing += 1
|
| continue
|
| arr = arr.reshape(init.dims)
|
| init.CopyFrom(numpy_helper.from_array(arr, name=name))
|
| resolved += 1
|
|
|
|
|
| for init in m.graph.initializer:
|
| if not any(e.key == "location" for e in init.external_data):
|
| continue
|
| t2 = look.get(norm_name(init.name))
|
| if t2 is None:
|
| continue
|
| want2 = onnx.helper.tensor_dtype_to_np_dtype(init.data_type)
|
| arr2 = t2.detach().cpu().numpy()
|
| if arr2.dtype != want2:
|
| arr2 = arr2.astype(want2, copy=False)
|
| if tuple(init.dims) and tuple(arr2.shape) != tuple(init.dims):
|
| if arr2.size != int(np.prod(init.dims)):
|
| continue
|
| arr2 = arr2.reshape(init.dims)
|
| init.CopyFrom(numpy_helper.from_array(arr2, name=init.name))
|
| resolved += 1
|
| missing = max(0, missing - 1)
|
|
|
| os.makedirs(OUT_DIR, exist_ok=True)
|
| dst = os.path.join(OUT_DIR, model + ".onnx")
|
| onnx.save(m, dst)
|
| info = {"model": model, "source_blob": os.path.basename(path),
|
| "nodes": len(m.graph.node), "initializers": len(m.graph.initializer),
|
| "external_refs": len(need), "resolved": resolved, "missing": missing,
|
| "bytes": os.path.getsize(dst)}
|
|
|
| try:
|
| onnx.checker.check_model(dst)
|
| info["checker"] = "OK"
|
| except Exception as e:
|
| info["checker"] = str(e)[:160]
|
| return dst, info
|
|
|
|
|
| def main() -> int:
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--all", action="store_true", help="包含已塞进 pymss 的 4 个")
|
| ap.add_argument("--only", default=None)
|
| args = ap.parse_args()
|
| os.makedirs(OUT_DIR, exist_ok=True)
|
|
|
| models = model_io.list_models()
|
| if not args.all:
|
| models = [n for n in models if n not in PYMSS_OK]
|
| if args.only:
|
| models = [n for n in models if args.only.lower() in n.lower()]
|
|
|
| report = []
|
| for model in models:
|
| try:
|
| meta = read_header(model)
|
| tensors = container_tensors(model)
|
| dst, info = inline(model, tensors, meta)
|
| except Exception as e:
|
| info = {"model": model, "error": f"{type(e).__name__}: {e}"}
|
| dst = None
|
| tag = "OK" if info.get("missing") == 0 and info.get("checker") == "OK" else "!!"
|
| print(f"[{tag}] {model:<32} " + (
|
| f"nodes={info.get('nodes')} inits={info.get('initializers')} "
|
| f"外部={info.get('external_refs')} 解析={info.get('resolved')} "
|
| f"缺={info.get('missing')} checker={info.get('checker')}"
|
| if "error" not in info else info["error"]))
|
| report.append(info)
|
|
|
| with open(os.path.join(OUT_DIR, "_inline_report.json"), "w", encoding="utf-8") as f:
|
| json.dump(report, f, indent=2, ensure_ascii=False)
|
| ok = sum(1 for r in report if r.get("missing") == 0 and r.get("checker") == "OK")
|
| print(f"\n完整内联 {ok}/{len(report)}")
|
| return 0
|
|
|
|
|
| if __name__ == "__main__":
|
| sys.exit(main())
|
|
|
