Blame for scripts/classify_models.py

#!/usr/bin/env python3
"""

classify_models.py

==================



根据 checkpoint 里的参数名 / 形状，判断 .ckpt 模型属于哪种分离模型架构

（bs_roformer / mel_band_roformer / scnet / bandit / … ）。



判据来自 E:\\vs\\pymss-core\\pymss_core\\modules 下的参考实现：



  bs_roformer.py / mel_band_roformer.py   ->  band_split(BandSplit) + mask_estimators(MaskEstimator) + layers(Transformer)

  bands.py  BandSplit       : self.to_features = [Sequential(RMSNorm(gamma), Linear)]

  bands.py  MaskEstimator   : self.to_freqs    = [Sequential(MLP, GLU)]

  transformer.py Attention  : self.norm(gamma) / to_qkv / to_gates / to_out

  scnet.py  SCNet           : self.separation_net / SDlayer / dp_modules / encoder / decoder

  bandit_v2/bandit.py       : self.band_split(norm_fc_modules) / tf_model(seqband) / mask_estim[stem]



bs_roformer 与 mel_band_roformer 的判别点（两者 band_split / mask_estimators 结构相同）：

  * BSRoformer.final_norm      = RMSNorm(dim)  -> state_dict 里有 "final_norm.gamma"

  * MelBandRoformer.final_norm  = nn.Identity() -> 没有任何 final_norm 参数

  * BSRoformer 常用 use_shared_bias -> 有 "linear_62_bias_0" / "linear_64_bias_0"

  * freq 带宽分布：BS 默认 DEFAULT_FREQS_PER_BANDS，mel 版本是按 mel 滤波器分带



用法

----

    python classify_models.py                    # 扫描 ./ckpt

    python classify_models.py --dir <ckpt目录>

    python classify_models.py --json out.json

"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from collections import defaultdict

import torch

import model_io

HERE = os.path.dirname(os.path.abspath(__file__))

# pymss-core 里 BS-RoFormer 的默认频带划分（单声道）
BS_DEFAULT_FREQS = (2,) * 24 + (4,) * 12 + (12,) * 8 + (24,) * 8 + (48,) * 8 + (128, 129)


def idx_set(pattern: str, keys) -> set[str]:
    out = set()
    for k in keys:
        m = re.search(pattern, k)
        if m:
            out.add(m.group(1))
    return out


def shape_of(sd, pattern: str):
    for k, v in sd.items():
        if re.search(pattern, k):
            return tuple(v.shape)
    return None


def classify(sd: dict) -> dict:
    all_keys = list(sd)
    # 别名键（coreml.source.* / onnx.source.*）保留了原始 PyTorch 名。有些模型
    # （例如 drumsep）的 mask estimator 只以别名形式存在，因此判定时使用全量键。
    keys = all_keys
    info: dict = {}

    # ---- Bandit / Bandit v2 ：band_split(norm_fc_modules) + tf_model(seqband) ----
    if any("norm_fc_modules" in k for k in keys) or any(k.startswith("tf_model") for k in keys):
        stems = sorted({m.group(1) for k in keys
                        if (m := re.match(r"mask_estim\.([A-Za-z0-9_]+)\.", k))})
        info.update(
            family="bandit",
            subtype="bandit_v2" if "band_split.norm_fc" in " ".join(keys) else "bandit",
            bands=len(idx_set(r"band_split\.norm_fc_modules\.(\d+)\.", keys)),
            stems=stems,
            stft_nfft=2048,
        )
        return info

    # ---- SCNet ----
    if any(k.startswith(("separation_net", "model.separation_net")) for k in keys) or \
       any("SDlayer" in k for k in keys):
        info.update(family="scnet",
                    encoders=len(idx_set(r"(?:^|[./])encoder\.(\d+)\.", keys)),
                    decoders=len(idx_set(r"(?:^|[./])decoder\.(\d+)\.", keys)))
        return info

    # ---- ConvNeXt U-Net generator（GAN 类降噪器）----
    if any("generator.trainable" in k for k in keys):
        info.update(family="convnext_gan_generator",
                    convnext_blocks=len(idx_set(r"convnext\.(\d+)\.", keys)),
                    u_blocks=len(idx_set(r"u_blocks\.(\d+)\.", keys)),
                    d_blocks=len(idx_set(r"d_blocks\.(\d+)\.", keys)),
                    snake=any(k.endswith(".alpha") for k in keys))
        return info

    # ---- ReDimNet 说话人编码器 ----
    if any(k.startswith("backbone.stage") for k in keys):
        info.update(family="redimnet",
                    stages=len(idx_set(r"backbone\.stage(\d+)\.", keys)))
        return info

    # ---- 小型打分头 ----
    if any(k.startswith(("scorer.", "temporal.")) for k in keys):
        info.update(family="scoring_head",
                    layers=len(idx_set(r"temporal\.layers\.(\d+)\.", keys)))
        return info

    # ---- Roformer 家族 ----
    if any("band_split" in k for k in keys) and \
       any(k.startswith("mask_est") or "mask_estimator" in k for k in keys):
        bands = (idx_set(r"band_split\.(\d+)\.linear\.weight", keys)
                 or idx_set(r"band_split\.to_features\.(\d+)\.", keys))
        n_bands = len(bands)
        # 每带 linear 的 (out_features=dim, in_features=复数频点×声道)
        widths = []
        for k in keys:
            m = re.fullmatch(r"band_split\.(\d+)\.linear\.weight", k)
            if m:
                widths.append((int(m.group(1)), sd[k].shape[0], sd[k].shape[1]))
        if not widths:
            for k in keys:
                m = re.fullmatch(r"band_split\.to_features\.(\d+)\.1\.weight", k)
                if m:
                    widths.append((int(m.group(1)), sd[k].shape[0], sd[k].shape[1]))
        widths.sort()
        per_band = [w // 2 for _, _, w in widths]   # 立体声 → /2 得单声道频点数

        final_norm = any("final_norm" in k for k in keys)
        shared_bias = any("linear_62_bias_0" in k or "linear_64_bias_0" in k for k in keys)
        branch = any(".branch." in k for k in keys)
        depth = len(idx_set(r"^(?:layers|blk)\.(\d+)\.", keys))
        stems_raw = (idx_set(r"(?:mask_estimators?|mask_est)\.(\d+)\.", keys)
                     | idx_set(r"mask_estimators?_(\d+)_", keys))
        heads = shape_of(sd, r"to_gates\.weight$")

        info.update(
            family="bs_roformer" if final_norm else "mel_band_roformer",
            bands=n_bands,
            depth=depth,
            dim=widths[0][1] if widths else None,
            heads=heads[0] if heads else None,
            stems=len(stems_raw),
            final_norm=final_norm,
            shared_bias=shared_bias,
            attention_branch=branch,
            freqs_per_band=per_band,
            freq_layout="BS-DEFAULT(stereo)" if per_band == [2 * f for f in BS_DEFAULT_FREQS]
                        else "BS-DEFAULT(mono)" if per_band == list(BS_DEFAULT_FREQS)
                        else "custom",
        )
        return info

    info.update(family="unknown")
    return info


def describe(info: dict) -> str:
    fam = info.get("family", "?")
    if fam in ("bs_roformer", "mel_band_roformer"):
        bits = [f"{info['bands']} bands", f"dim={info['dim']}", f"depth={info['depth']}",
                f"{info['stems']} stem(s)", f"freq={info.get('freq_layout')}"]
        if info.get("shared_bias"):
            bits.append("shared_bias")
        if info.get("attention_branch"):
            bits.append("attn-branch")
        return "  ".join(bits)
    if fam == "bandit":
        return f"n_bands={info.get('bands')}  stems={info.get('stems')}"
    if fam == "scnet":
        enc, dec = info.get("encoders") or 0, info.get("decoders") or 0
        return f"{enc} enc / {dec} dec" if enc else f"{dec} decoder(s)"
    if fam == "convnext_gan_generator":
        return (f"convnext={info.get('convnext_blocks')} u_blocks={info.get('u_blocks')} "
                f"d_blocks={info.get('d_blocks')} snake={info.get('snake')}")
    if fam == "redimnet":
        return f"stages={info.get('stages')}"
    if fam == "scoring_head":
        return f"temporal_layers={info.get('layers')}"
    return ""


def main() -> int:
    ap = argparse.ArgumentParser(description="识别 ckpt 的分离模型类型")
    ap.add_argument("--dir", default=os.path.join(HERE, "..", "raw"), help="模型目录（safetensors / ckpt）")
    ap.add_argument("--json", default=None, help="把结果另存为 JSON")
    args = ap.parse_args()

    paths = []
    for root, _d, files in os.walk(args.dir):
        for n in files:
            if n.endswith((".ckpt", ".safetensors")):
                paths.append(os.path.join(root, n))
    if not paths:
        print("没有找到 .ckpt / .safetensors")
        return 1

    results = {}
    print(f"{'模型':<34}{'类型':<22}关键参数")
    print("-" * 108)
    for p in sorted(paths):
        sd = model_io.load_any(p)
        info = classify(sd)
        name = os.path.splitext(os.path.basename(p))[0]
        results[name] = {"type": info.get("family"), **{k: v for k, v in info.items() if k != "family"}}
        print(f"{name:<34}{info.get('family', '?'):<22}{describe(info)}")
        del sd

    if args.json:
        with open(args.json, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        print(f"\n已写入 {args.json}")
    return 0


if __name__ == "__main__":
    sys.exit(main())