feat: build all split model members

Author: hi-ogawa-agent

docs/split-dsp.md

@@ -113,6 +113,14 @@ for frequency, and `5.472e-5` max absolute error and `6.695e-12` MSE for time. O
 input, the official Hugging Face model loader and the legacy fork/checkpoint loader produced
 byte-identical final output (`sha256:21d88c3dcb713451ffcf20c4c44a17306bb10f11ae0277df4ab90c3043146359`).
 
+The v2 builder accepts explicit member names or `--all`. It atomically writes the requested model
+set and a `manifest.json` containing the graph flavor, upstream version, tensor contract, source
+order, precision, specialty, artifact size and digest, and measured branch parity. A full build is:
+
+```sh
+pnpm build-model-v2 --all
+```
+
 Split artifacts must use a distinct directory, filename convention, and manifest flavor while the
 current self-contained artifacts remain supported. A split model must never be loadable as a
 waveform model by accident.
@@ -156,10 +164,10 @@ code.
 
 ### 2. Verify the split ONNX artifact
 
-- Compare frequency and time outputs independently against PyTorch.
-- Extend model-building scripts with an explicit split flavor.
-- Generate all standard and fine-tuned members only after the standard model passes.
-- Record exact model sizes and tensor metadata.
+- [x] Compare frequency and time outputs independently against PyTorch.
+- [x] Add an isolated split-model builder with explicit-member and `--all` modes.
+- [x] Generate and verify the standard model and all four fine-tuned specialists.
+- [x] Record exact model sizes, digests, parity metrics, and tensor metadata in a manifest.
 
 ### 3. Implement native runtime DSP
 

package.json

@@ -7,7 +7,7 @@
     "link-wt": "bash tools/link-wt.sh",
     "bass-cover": "uv run --no-sync tools/bass-cover.py",
     "build-model": "uv run python tools/model-export/build_models.py",
-    "build-model-v2": "uv run --project tools/model-export-v2 python tools/model-export-v2/export_split_onnx.py --out data/onnx-split/htdemucs.onnx",
+    "build-model-v2": "uv run --project tools/model-export-v2 python tools/model-export-v2/export_split_onnx.py --out data/onnx-split",
     "verify-parity": "uv run python tools/model-export/verify_parity.py",
     "model-release": "uv run --no-sync tools/model_release.py",
     "build-wasm": "wasm-pack build crates/wasm --target web --release",

tools/model-export-v2/export_split_onnx.py

@@ -5,11 +5,18 @@
 
 Usage:
     uv run --project tools/model-export-v2 python tools/model-export-v2/export_split_onnx.py \
-        --out data/onnx-split/htdemucs.onnx
+        htdemucs --out data/onnx-split
+    uv run --project tools/model-export-v2 python tools/model-export-v2/export_split_onnx.py \
+        --all --out data/onnx-split
 """
 
 import argparse
+import hashlib
+import json
+import os
 from pathlib import Path
+import shutil
+import tempfile
 
 import numpy as np
 import onnx
@@ -20,6 +27,15 @@
 from demucs.pretrained import get_model
 from einops import rearrange
 
+MEMBERS = (
+    "htdemucs",
+    "htdemucs_ft_drums",
+    "htdemucs_ft_bass",
+    "htdemucs_ft_other",
+    "htdemucs_ft_vocals",
+)
+FT_PREFIX = "htdemucs_ft_"
+
 
 class SplitHTDemucs(torch.nn.Module):
     """HTDemucs from packed spectrogram/waveform inputs to decoded branch outputs."""
@@ -106,18 +122,45 @@ def forward(
         return frequency, time
 
 
-def get_core() -> HTDemucs:
-    model = get_model("htdemucs")
+def unwrap_model(model: HTDemucs | BagOfModels, name: str) -> HTDemucs:
     if isinstance(model, BagOfModels):
         if len(model.models) != 1:
-            raise ValueError(f"expected one htdemucs member, got {len(model.models)}")
+            raise ValueError(f"expected one {name} member, got {len(model.models)}")
         model = model.models[0]
     if not isinstance(model, HTDemucs):
         raise TypeError(f"expected HTDemucs, got {type(model)}")
     model.eval()
     return model
 
 
+def load_members(requested: list[str]) -> dict[str, HTDemucs]:
+    members = {}
+    if "htdemucs" in requested:
+        members["htdemucs"] = unwrap_model(get_model("htdemucs"), "htdemucs")
+
+    requested_ft = {name.removeprefix(FT_PREFIX) for name in requested if name.startswith(FT_PREFIX)}
+    if requested_ft:
+        bag = get_model("htdemucs_ft")
+        if not isinstance(bag, BagOfModels):
+            raise TypeError(f"expected htdemucs_ft bag, got {type(bag)}")
+        for core, weights in zip(bag.models, bag.weights):
+            selected = [source for source, weight in zip(bag.sources, weights) if weight > 0]
+            if len(selected) != 1:
+                raise ValueError(f"expected one-hot htdemucs_ft weights, got {weights}")
+            source = selected[0]
+            if source not in requested_ft:
+                continue
+            if not isinstance(core, HTDemucs):
+                raise TypeError(f"expected HTDemucs, got {type(core)}")
+            core.eval()
+            members[f"{FT_PREFIX}{source}"] = core
+
+    missing = [name for name in requested if name not in members]
+    if missing:
+        raise ValueError(f"failed to load member(s): {', '.join(missing)}")
+    return members
+
+
 def pack_spectrogram(core: HTDemucs, mix: torch.Tensor) -> torch.Tensor:
     z = core._spec(mix)
     batch, channels, frequencies, frames = z.shape
@@ -161,7 +204,7 @@ def validate_onnx(
     spectrogram: torch.Tensor,
     expected_frequency: torch.Tensor,
     expected_time: torch.Tensor,
-) -> None:
+) -> dict[str, dict[str, float]]:
     model = onnx.load(str(path), load_external_data=False)
     onnx.checker.check_model(model)
     inputs = [(value.name, [dim.dim_value for dim in value.type.tensor_type.shape.dim])
@@ -200,6 +243,7 @@ def validate_onnx(
         ["frequency", "time"],
         {"waveform": mix.numpy(), "spectrogram": spectrogram.numpy()},
     )
+    parity = {}
     for name, actual, expected in (
         ("frequency", frequency, expected_frequency.numpy()),
         ("time", time, expected_time.numpy()),
@@ -210,34 +254,115 @@ def validate_onnx(
         print(f"ONNX {name}: max abs {max_abs:.3e}, mse {mse:.3e}")
         if max_abs > 2e-3 or mse > 1e-7:
             raise SystemExit(f"ONNX {name} parity failed: max abs {max_abs:.3e}, mse {mse:.3e}")
+        parity[name] = {"max_abs": max_abs, "mse": mse}
     print(f"split graph: {path.stat().st_size / 1e6:.1f} MB, no DFT payload")
+    return parity
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--out", type=Path, required=True)
-    args = parser.parse_args()
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as file:
+        while chunk := file.read(1024 * 1024):
+            digest.update(chunk)
+    return digest.hexdigest()
+
 
-    core = get_core()
+def export_member(name: str, core: HTDemucs, out: Path) -> dict:
+    print(f"exporting {name} ...")
     split = SplitHTDemucs(core).eval()
     generator = torch.Generator().manual_seed(42)
     mix = torch.randn(1, 2, mix_length(core), generator=generator)
     with torch.no_grad():
         spectrogram = pack_spectrogram(core, mix)
     frequency, time = verify_python_seam(core, split, mix, spectrogram)
 
-    args.out.parent.mkdir(parents=True, exist_ok=True)
+    path = out / f"{name}.onnx"
     torch.onnx.export(
         split,
         (mix, spectrogram),
-        args.out,
+        path,
         export_params=True,
         opset_version=17,
         do_constant_folding=True,
         input_names=["waveform", "spectrogram"],
         output_names=["frequency", "time"],
     )
-    validate_onnx(args.out, mix, spectrogram, frequency, time)
+    parity = validate_onnx(path, mix, spectrogram, frequency, time)
+    return {
+        "file": path.name,
+        "member": name,
+        "precision": "fp32",
+        "sha256": sha256(path),
+        "size": path.stat().st_size,
+        "specialty": name.removeprefix(FT_PREFIX) if name.startswith(FT_PREFIX) else None,
+        "parity": parity,
+    }
+
+
+def replace_output(staging: Path, output: Path) -> None:
+    backup = None
+    if output.exists():
+        backup = Path(tempfile.mkdtemp(prefix=f".{output.name}.old-", dir=output.parent))
+        backup.rmdir()
+        os.replace(output, backup)
+    try:
+        os.replace(staging, output)
+    except BaseException:
+        if backup is not None:
+            os.replace(backup, output)
+        raise
+    if backup is not None:
+        shutil.rmtree(backup)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("members", nargs="*", metavar="MEMBER")
+    parser.add_argument("--all", action="store_true", help="build all model members")
+    parser.add_argument("--out", type=Path, required=True)
+    args = parser.parse_args()
+
+    if args.all and args.members:
+        parser.error("pass either --all or explicit members, not both")
+    if not args.all and not args.members:
+        parser.error("pass --all or at least one member")
+    requested = list(MEMBERS) if args.all else list(dict.fromkeys(args.members))
+    unknown = [name for name in requested if name not in MEMBERS]
+    if unknown:
+        parser.error(f"unknown member(s): {', '.join(unknown)}; expected: {', '.join(MEMBERS)}")
+
+    output = args.out.resolve()
+    output.parent.mkdir(parents=True, exist_ok=True)
+    staging = Path(tempfile.mkdtemp(prefix=f".{output.name}.new-", dir=output.parent))
+    try:
+        cores = load_members(requested)
+        models = [export_member(name, cores[name], staging) for name in requested]
+        manifest = {
+            "format": 1,
+            "graph_flavor": "split-dsp",
+            "models": models,
+            "onnx_contract": {
+                "inputs": {
+                    "waveform": [1, 2, 343980],
+                    "spectrogram": [1, 4, 2048, 336],
+                },
+                "outputs": {
+                    "frequency": [1, 4, 4, 2048, 336],
+                    "time": [1, 4, 2, 343980],
+                },
+                "sample_rate": 44100,
+                "sources": ["drums", "bass", "other", "vocals"],
+            },
+            "upstream": {"demucs": "4.1.0"},
+        }
+        (staging / "manifest.json").write_text(
+            json.dumps(manifest, indent=2, sort_keys=True) + "\n"
+        )
+        replace_output(staging, output)
+    except BaseException:
+        shutil.rmtree(staging, ignore_errors=True)
+        raise
+    print(f"wrote {len(requested)} model(s) and manifest to {output}")
 
 
 if __name__ == "__main__":