|
|
| """
|
| 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__))
|
|
|
|
|
| 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)
|
|
|
|
|
| keys = all_keys
|
| info: dict = {}
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| 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)
|
|
|
| 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]
|
|
|
| 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())
|
|
|
