#!/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 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())