Snooow1029/itts-attribute-classifiers

ITTS Attribute Classifiers (accent + age)

Speaker-attribute classifiers used to score speaker diversity / controllability of instruction-TTS (ITTS) systems. They reproduce the Vox-Profile recipe (foundation-model encoder + LoRA adapter + light downstream head), each attribute served as a WavLM + Whisper ensemble (softmax-averaged), which is stronger than either single model.

This repo stores only the multi-GB *.pt weights. The model code, training and inference scripts live in the companion GitHub repo (Vox-Profile WavLMWrapper / WhisperWrapper).

Checkpoints

file	backbone	task	classes	test acc	macro-F1	UAR
accent_wavlm.pt	WavLM-Large	English accent	6	0.939	0.935	0.936
accent_whisper.pt	Whisper-Large-V3	English accent	6	0.940	0.942	0.950
age_wavlm.pt	WavLM-Large	speaker age	4	0.791	0.797	0.801
age_whisper.pt	Whisper-Large-V3	speaker age	4	0.856	0.862	0.865

Ensembles (softmax-average of the two members, on the held-out test split):

attribute	acc	macro-F1	UAR
accent (wavlm+whisper)	0.963	0.962	0.965
age (wavlm+whisper)	0.860	0.866	0.869

• Accent labels (6): North America, English, South Asia, Germanic, Oceania, South African
• Age labels (4): child, young_adult, middle_aged, older

Architecture

Each checkpoint is a Vox-Profile wrapper (src/model/accent/wavlm_accent.py / whisper_accent.py in the GitHub repo):

• Frozen foundation-model encoder (WavLM-Large or Whisper-Large-V3 encoder).
• LoRA rank 16 on the upper-half FFN layers (intermediate_dense + output_dense).
• Learnable softmax layer-weighted-sum over hidden states (WavLM); Whisper wrapper uses the last encoder layer as shipped.
• 3× point-wise Conv1d(→256) head + length-aware mean pooling + Linear(256,256)→ReLU→Linear(256,C).

Age and accent share the same wrapper and LoRA-r16 recipe; only the output class count and the training data differ. Age is a 4-class single-task head (child / young_adult / middle_aged / older).

Training data

Same source mix as Vox-Profile's released models: CommonVoice + TIMIT + VoxCeleb (age-enriched) for age; accent from the Vox-Profile English-accent sources. Audio filtered to 3–15 s, speaker-independent train/val/test splits. LoRA fine-tune: Adam lr 5e-4, effective batch 16, 10 epochs (accent 15), best checkpoint selected on validation UAR (age) / macro-F1 (accent).

Inference

import sys, torch, torchaudio, loralib as lora
# point these at the GitHub checkout (Vox-Profile wrappers)
sys.path.append("speaker_embedding_eval/vox-profile-release/src/model/accent")
from wavlm_accent import WavLMWrapper
from whisper_accent import WhisperWrapper
from huggingface_hub import hf_hub_download

REPO = "Snooow1029/itts-attribute-classifiers"
AGE_LABELS = ["child", "young_adult", "middle_aged", "older"]
device = "cuda"

def load(cls, fname, pretrain, n):
    ckpt = hf_hub_download(REPO, fname)
    m = cls(pretrain_model=pretrain, finetune_method="lora", lora_rank=16,
            output_class_num=n, apply_gradient_reversal=False).to(device)
    m.load_state_dict(torch.load(ckpt, map_location=device)["state_dict"])
    # IMPORTANT loralib double-merge fix: best.pt stores already-merged LoRA weights but the
    # `merged` flag is not in the state_dict. Mark layers merged so the following .eval() does
    # NOT merge a second time (W + 2*BA). Critical for Whisper (0.46 -> 0.94), +1pt for WavLM.
    for mod in m.modules():
        if isinstance(mod, lora.LoRALayer):
            mod.merged = True
    return m.eval()

wavlm   = load(WavLMWrapper,   "age_wavlm.pt",   "wavlm_large",   len(AGE_LABELS))
whisper = load(WhisperWrapper, "age_whisper.pt", "whisper_large", len(AGE_LABELS))

wav, sr = torchaudio.load("clip.wav")           # mono; resample to 16 kHz if needed
if sr != 16000:
    wav = torchaudio.functional.resample(wav, sr, 16000)
x = wav.mean(0, keepdim=True).to(device)        # [1, T]
length = torch.tensor([x.shape[-1]], device=device)

with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
    pw = torch.softmax(wavlm(x, length=length).float(), -1)
    ph = torch.softmax(whisper(x, length=length).float(), -1)
pred = ((pw + ph) / 2).argmax(-1).item()        # ensemble
print(AGE_LABELS[pred])

Swap age_* → accent_* and AGE_LABELS → the 6 accent labels for accent inference.

Citation

Built on Vox-Profile. If you use these classifiers, please cite:

@article{voxprofile2025,
  title={Vox-Profile: A Speech Foundation Model Benchmark for Characterizing Diverse Speaker and Speech Traits},
  journal={arXiv preprint arXiv:2505.14648},
  year={2025}
}