| """Standalone v5 demo pass: one STRONG utterance per emotion + neutral + zh clip.""" |
| import glob, json, os, sys |
| import torch |
| import torch.nn.functional as F |
| sys.path.insert(0, "/tmp/work/s2st_emotion") |
| from compare_v5 import (demo_plot, new_target, load_wav, EMO_FIELD) |
|
|
| def main(): |
| from funasr import AutoModel |
| from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder |
| DEV = "cuda:0" |
| fun = AutoModel(model="/tmp/work/emotion2vec_plus_base", trust_remote_code=False, |
| device=DEV, disable_update=True) |
| teacher = fun.model.eval() |
| payload = torch.load("/tmp/work/surgical_out_v5/surgical_emotion_module.pt", |
| map_location=DEV, weights_only=False) |
| cfg = SurgicalConfig() |
| cfg.group_size = int((payload.get("config") or {}).get("group_size", 8)) |
| enc = SurgicalEmotionEncoder(teacher, cfg).to(DEV) |
| enc.adapter.load_state_dict(payload["adapter_state"]); enc.eval() |
| torch.set_grad_enabled(False) |
|
|
| files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True)) |
| speech = [f for f in files if os.path.basename(f).split("-")[1] == "01" |
| and 19 <= int(os.path.basename(f).split("-")[-1].split(".")[0]) <= 24] |
| by_emo = {} |
| for f in speech: |
| by_emo.setdefault(os.path.basename(f).split("-")[2], []).append(f) |
| demo_files = [] |
| for e in sorted(by_emo): |
| if e == "01": continue |
| strong = [f for f in by_emo[e] if os.path.basename(f).split("-")[3] == "02"] |
| demo_files.append(strong[0] if strong else by_emo[e][0]) |
| neutral = by_emo.get("01", []) |
| if neutral: demo_files = [neutral[0]] + demo_files |
| for f in demo_files: |
| wav = load_wav(f).cpu() |
| b = os.path.basename(f) |
| out = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) |
| src = F.layer_norm(wav[None], wav[None].shape).to(DEV) |
| states = teacher.extract_features(src, padding_mask=None, mask=False)["x"] |
| probs_t9 = torch.softmax(teacher.proj(states), -1)[0] |
| tgt50 = new_target(probs_t9, wav) |
| emo = EMO_FIELD.get(b.split("-")[2], "?") |
| ilbl = b.split("-")[3] |
| demo_plot(f"{emo}_{b[:20]}", wav, out, probs_t9, tgt50, cfg.frame_rate_hz, |
| {"ravdess_label": emo, |
| "intensity_label": "strong" if ilbl == "02" else "normal"}) |
| print("demo done:", emo, flush=True) |
| print("DEMOS_DONE", flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
