Feat/diarization multi speaker uncertainty

[Evan8456]

Description

Related Issue(s)

Motivation and Context

How Has This Been Tested?

Screenshots (if appropriate):

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Checklist:

• I have added tests to cover my changes.
• All new and existing tests passed.
• My code follows the code style of this project.

[wilke0818]

Code review — f17c6b06

Scope: git diff 22da0e40 f17c6b06 (the merge-base with alpha) — 14 files, +1097/−9. Four new diarization backends (VibeVoice-ASR-HF, USC-SAIL child-adult, MOSS-Transcribe-Diarize, DiariZen), the model_for_task dispatch extension, and transformers>=5.0 → >=5.3. Verified against upstream sources (usc-sail/child-adult-diarization, BUTSpeechFIT/DiariZen, OpenMOSS/MOSS-Transcribe-Diarize), the installed transformers 5.5.4 VibeVoice implementation, and the live PyTorch wheel indexes. 16 findings.

High

1. child_adult.py:45 — the torch==2.3.0 pin means the venv can never build on any host this backend supports.
ensure_venv is called without max_cuda_version, so Stage 1 installs through the host-selected PyTorch index (cu128/cu126/cu124/cu121, cuda_probe.py:32-41). torch 2.3.0 publishes wheels only on cu121/cu118 — I checked the indexes directly. On a CUDA ≥ 12.4 host Stage 1 raises SenselabCudaCompatibilityError, shutil.rmtrees the venv, and memoizes nothing, so every subsequent call repeats the whole failure. Since the backend is CUDA-only, those are the only hosts that could ever run it. scene_quality/brouhaha.py:204 is the existing precedent: pass max_cuda_version=(12, 1).

2. vibevoice.py:142 — inputs are moved to the device but never cast to the model dtype, so all CUDA inference fails.
_select_device_and_dtype (:59) returns torch.float16 for CUDA and the model loads with dtype=dtype (:80), but :142 only does v.to(vv_model.device). The processor emits input_values as float32 and VibeVoiceAsrForConditionalGeneration.forward passes it into get_audio_features with no internal cast → Input type (float) and weight type (half) should be the same. CPU works only because its dtype happens to be float32, which is why nothing caught it. Both the model card and the transformers forward docstring do inputs.to(model.device, dtype=model.dtype); BatchFeature.to casts only float tensors, leaving input_ids/padding_mask intact.

3. child_adult.py:89 — the clone guard wedges permanently and races across jobs.
Guard is if not (repo_dir / "whisper-modeling").is_dir(): git clone … repo_dir. (a) If a clone is interrupted — SIGKILL, or the 1200 s subprocess timeout firing mid-clone — repo_dir is left non-empty without that subdir: the guard stays false while git clone into a non-empty directory can never succeed, so the backend is permanently wedged with no cleanup path. (b) ensure_venv's FileLock is released before subprocess.run, so two concurrent jobs sharing $HOME both see the subdir missing and clone into the same path at once. Clone to a temp path then os.replace, under a lock.

4. api.py:119 — max_new_tokens is unreachable, and truncation is reported as success.
The 4096 default (vibevoice.py:92, moss.py:98) can't be raised: diarize_audios forwards only audios/model/device, and stage_diarization has no channel for it either. On a 20+ minute recording generation truncates mid-output; VibeVoice's JSON then fails to parse and is swallowed by the guard at vibevoice.py:151, returning an empty segment list as a successful result — indistinguishable downstream from "no speech detected". MOSS returns only the parsable prefix. Nothing records that truncation occurred.

Medium

5. vibevoice.py:59 — fp16 for a bfloat16-native 7B model. config.json declares "dtype": "bfloat16" and every upstream example uses dtype="auto". Even after (2) is fixed, running bf16-trained weights in fp16 loses ~4 exponent bits — the overflow→NaN path, which surfaces silently here as garbage swallowed into segments = [] by the :151 guard.

6. vibevoice.py:151 — parse guard is narrower than its documented contract. It catches only json.JSONDecodeError, but the contract is "Empty for an audio whose output didn't parse into structured segments (logged as a warning, not raised)." A KeyError/IndexError/TypeError/plain ValueError from processor.decode(..., return_format="parsed") propagates out of the per-audio loop and aborts the whole diarize_audios call, discarding every other audio in the batch. json.JSONDecodeError subclasses ValueError, so widening is safe. float(seg["Start"]) at :166 is also outside the guard.

7. child_adult.py:179 — the model argument is accepted but ignored, and explicit revisions are dropped. The worker always gets the hardcoded _CHILD_ADULT_HF_REPO + weights filename (:208-209) and staging hardcodes "main" (:218-223). Since api.py dispatches on the prefix, HFModel(path_or_uri="AlexXu811/whisper-child-adult-v2", revision="<sha>") silently gets a different model than requested. moss.py:164 and diarizen.py:217 do forward path_or_uri but also hardcode "main", so a revision pin is silently dropped there too. canary_qwen.py:252 shows the intended pattern.

8. child_adult.py:153 — the docstring understates how much audio is dropped. It says only "a trailing partial window shorter than 10s is dropped by upstream's own chunking loop", but upstream's loop is while start + WINDOW_SIZE < length (strict <, step 10). A 4.9 s clip → zero windows → []; exactly 10.0 s → []; 20.0 s → one window, so only 0–10 s is analysed. For a backend answering "was an adult voice present at all", [] is indistinguishable from "no adult" — a silent false negative precisely on short pediatric prompt recordings.

9. speaker_diarization_test.py:23 — the skip gate is wrong twice over and the test it guards is vacuous. (a) _CHILD_ADULT_VENV_ROOT hardcodes Path.home()/".cache"/"senselab"/"venvs"/…, ignoring the override subprocess_venv._cache_dir() honors. (b) test_diarize_audios_with_child_adult (:212) gates only on venv presence, but the backend raises without CUDA (_select_device_and_dtype(user_preference=CUDA, compatible_devices=[CUDA]) raises ValueError on a CPU host, device.py:130-137), so with $HOME shared between GPU nodes and a CPU login node it runs and fails on the login node instead of skipping. (c) The fixture is 4.92 s, under the 10 s window, so per (8) the backend returns [] and both all(...) assertions pass over an empty list — the test can never detect a broken backend even when it does run.

10. vibevoice.py:36 — a 7B class-level cache with no eviction, in a stage that never unloads. VibeVoiceDiarization._models has no eviction and no release hook, and stage_diarization (stages.py:136-160) runs each diar model in turn and returns without dropping references or calling torch.cuda.empty_cache(). Configured alongside a large in-process ASR model, ~14 GB stays resident and the second load OOMs on a 24–40 GB card. The other three new backends are subprocess-hosted and free their memory on exit; VibeVoice is the one that doesn't.

11. diarizen.py:217 — the offline staging is all-or-nothing and fails silently. hf_subprocess_env wraps staging in try/except Exception: return env with no logging (dependencies.py:568-572). If pyannote/wespeaker-voxceleb-resnet34-LM can't be staged, the offline flag is dropped for both repos and the worker reverts to online Hub loading — the exact 429 path this PR's final commit exists to remove, with nothing in the logs. Also worth double-checking the docstring claim at :26-28 that this repo needs no HF_TOKEN: it's a pyannote/* repo, and pyannote.py:67-76 threads a token explicitly because the others are gated. (resolve_model falls back to the env token, so this only bites tokenless hosts.)

Low

12. api.py:15 — _VIBEVOICE_PREFIXES = ("microsoft/VibeVoice",) is too broad. microsoft/VibeVoice-1.5B and -Large are TTS checkpoints sharing that prefix; passing one for diarization now routes to VibeVoiceAsrForConditionalGeneration.from_pretrained and fails with an opaque config error instead of reaching the clear NotImplementedError at :148. Same breadth in model_for_task (utils/data_structures/model.py:334). Narrow to "microsoft/VibeVoice-ASR".

13. api.py:119 — num_speakers/min_speakers/max_speakers are silently dropped by all four new branches. Documented as Pyannote-only and consistent with the existing Sortformer branch, but a caller passing num_speakers=2 to DiariZen gets neither an effect nor a warning.

Integration with the restructured analyze_audio

Worth stating because the model_for_task extension (model.py:331-338) makes all four backends reachable from analyze_audio immediately — stage_diarization → model_for_task → diarize_audios, no stages.py change needed, and they inherit the cached-inference layer and provenance for free. That's a nice property, but the "diarization" slot is the wrong shape for three of them:

14. clustering.py:256 — child-adult's role labels break the identity axis by construction. It emits CHILD/ADULT/OVERLAP, not speaker identities, while cluster_speaker_labels_by_embedding averages all segments sharing a label: with two children the CHILD centroid is a blend of two speakers, and OVERLAP is a mixture that snaps to whichever centroid is nearest. identity.py:297's __cross_diar_label_disagreement__ then reads that as genuine cross-model disagreement in every overlap bucket, and the embedding-free fallback (clustering.py:230-237) is worse still, using the raw role label directly as a cluster id. Its real value — "is an unconsented adult audible here" — is a file-level role signal that wants its own stage, not a diar voter slot. (The fused n_pred_speakers at adaptive/evaluate.py:168 counts fusion.py:268 clusters, so it's affected only indirectly.)

15. stage_context.py:129 — joint ASR+diar models pay twice and lose their transcript. vibevoice.py:168 and moss.py:183 populate ScriptLine.text, but diar consumers read only speaker/start/end (identity.py:190 via diar_speaker_label_in_window), so the transcript reaches the sidecar JSON and nothing else. cache_key includes task, so registering the same model under --asr-models too means a second full 7B inference rather than a shared one. The right shape is a joint stage emitting into both summary["diarization"]["by_model"] and summary["asr"]["by_model"] from one inference — every current stage is single-task, so that's a contract change, not a model-id addition. The payoff is real: utterance.py votes on transcripts and no current ASR model gives speaker attribution.

16. stage_context.py:42 — STAGE_VERSIONS["diarization"] is one counter for all diar models. Cache keys are per-model_id, so entries don't collide, but bumping the counter because one new backend's outcome shape changed invalidates every other diar model's cached results, including expensive pyannote and Sortformer entries. Adding four backends to this stage makes that shared blast radius much likelier to be triggered.

Suggested sequencing. DiariZen is the only clean drop-in — speaker labels only, subprocess-hosted, and architecturally independent of both current voters, so it actually decorrelates the pool rather than double-weighting one family (compute.py:118-121 is the existing synthetic-voter precedent). VibeVoice and MOSS are worth landing as backends but not as diar voters; the joint stage is the follow-up, and VibeVoice needs (2) and (5) before it runs on CUDA at all. child-adult shouldn't be wired as a diar voter in any form, and can't be relied on until (1) and (8) are fixed.

Checked and found correct

The 1-D generated_ids slice at vibevoice.py is right (tokenizer.decode returns a str for 1-D input and extract_speaker_dict then returns list[dict], so the model card's [0] would be wrong here); AlexXu811/whisper-child-adult having no config.json doesn't break resolve_model; lightning>=2.0.1 in _DIARIZEN_REQUIREMENTS is fine (2.6.5 is live, contradicting the stale note in nvidia.py); DiariZen's pyannote-audio is a real subdirectory, not a submodule; the un-+SKIP'd doctest is harmless with no --doctest-modules configured; the model_registry.md YAMNet edit correctly re-syncs md to the already-updated yaml; and MOSS's worker (resolve_device, build_transcription_messages, generate_transcription, parse_transcript, auto_map) and DiariZen's worker (DiariZenPipeline.from_pretrained, itertracks) both check out against upstream.

[wilke0818]

Follow-up review — 85947c9a

Re-reviewed against the merge-base 22da0e40 after 3da5961b / e901c54e / 85947c9a. All four highs are closed, and two of the fixes are better than what I suggested. 7 new findings, none high — all are fixes that are incomplete or overshoot.

Status of the previous 16

#	Status	Evidence
1 torch pin	fixed	_CHILD_ADULT_MAX_CUDA_VERSION = (12, 1) (child_adult.py:59) passed at :232; ensure_venv accepts it at subprocess_venv.py:198
2 dtype cast	fixed, improved	model_dtype = next(vv_model.parameters()).dtype, cast gated on v.is_floating_point() — ties to what actually loaded rather than assuming fp16, which is better than the model.dtype I proposed
3 clone wedge	fixed	fcntl.flock + mkdtemp + os.replace (see finding 4 below for the residue)
4 max_new_tokens	fixed, exceeded	Exposed at api.py:44, forwarded at :146/:162 — and you added the truncation warning at vibevoice.py:180-186, which closes the silent-truncation half I flagged but didn't ask for
5 fp16 vs bf16	fixed	dtype="auto" at :86 with the rationale in-comment; _select_device_and_dtype's dtype now discarded
6 parse guard	fixed	Widened at :193 (see finding 5 below)
7 revision dropped	fixed	child_adult threads hf_revision → hf_hub_download; moss forwards model_revision to both from_pretrained calls; diarizen can't forward (upstream DiariZenPipeline.from_pretrained has no revision param) so it now warns explicitly at :193-200 instead of dropping silently — good resolution
8 window loop	documented, not fixed	:183-200 now states the strict <, that a ≤10 s clip analyzes zero windows, and that the final 10 s block is always dropped. Behavior unchanged — which matters for finding 2 below
9 test gate	fixed	_cache_dir() at :27, CUDA gate added (see finding 7 below)
10 7B cache	fixed, but see finding 1	release_all() added, wrong call site
11 silent staging	fixed	Warning at dependencies.py:576 before returning the online env. The no-HF_TOKEN claim also checks out — I queried the Hub: BUT-FIT/diarizen-wavlm-large-s80-md and pyannote/wespeaker-voxceleb-resnet34-LM are both gated=False, versus gated="auto" for pyannote/speaker-diarization-community-1, so the contrast the docstring draws is accurate
12 prefix breadth	fixed	("microsoft/VibeVoice-ASR",) at api.py:16 and model.py:334
13 speaker hints	fixed, incomplete	See finding 3 below
14 role labels	fixed, incomplete	See finding 2 below
15 joint ASR+diar	not fixed	Unchanged, and reasonably so — it's a stages.py contract change, not a patch
16 shared counter	documented, not fixed	A 10-line note explaining the grown blast radius and deferring the per-backend-family split until a concrete bump forces it. Agreed as the right call

New findings

1. stages.py:175 — medium: release_all() is called per clip, so the model cache is dead code.
stage_diarization runs once per pass/clip via run_pass, so the class-level cache is emptied after every clip and the 7B (~14 GB) model is fully re-loaded on the next one. A 500-clip run with --diarization-models microsoft/VibeVoice-ASR-HF does 500 from_pretrained + .to(cuda) cycles instead of 1, and VibeVoiceDiarization's documented "repeated calls with the same configuration reuse the initialized model" never holds in the only in-repo caller. It also fires on a run_task_cached hit, where the model was never loaded — a sign the release is keyed to the wrong event. The OOM concern is real; the release belongs at end-of-run, or gated on a subsequent large in-process stage.

2. presence.py:201 — medium: the role-label guard missed the third diar consumer.
ROLE_LABEL_ONLY_PREFIXES was applied to clustering.py (both paths) and identity.py's diar_ok, but not presence.py's diar_ok, which also treats every status-ok diar model as a peer voter (:314, via diar_speaks_in_window). Combined with finding 8 above — child-adult emits zero segments for any clip ≤ 10 s and always drops the final 10 s block — and harvesters.py:32 mapping "no overlapping segment" to speaks=False, wiring child-adult into --diarization-models makes it actively vote "no speech" over regions that contain speech, rather than abstaining. Same misconfiguration the other two guards were added to prevent.

3. api.py:138 — low: _warn_if_speaker_hints_ignored skips the pre-existing Sortformer branch.
It was added to all four new backends, but diarize_audios_with_nvidia_sortformer (nvidia.py:93) also accepts no speaker-hint parameters, so diarize_audios(..., model=HFModel("nvidia/diar_sortformer_4spk-v1"), num_speakers=2) still drops the hint silently — the exact behavior the warning was introduced to eliminate. The max_speakers docstring ("NVIDIA Sortformer is limited to 4 speakers") compounds it by reading as though the hint were honored.

4. child_adult.py:112 — low: the atomic clone leaks a temp dir per failed attempt.
mkdtemp(prefix=".child-adult-clone-", dir=repo_dir.parent) is followed by sp.run([...git clone...], check=True); on clone failure the raise propagates while finally releases only the flock. The partial .child-adult-clone-XXXX directory stays in the venv dir and one more accumulates per retry. Wrap the clone and shutil.rmtree(tmp_clone_dir, ignore_errors=True) on failure.

5. vibevoice.py:193 — low: the widened guard now also swallows API-contract breakage.
Catching KeyError/IndexError/TypeError/ValueError around processor.decode covers truncated output, but if a future transformers revision changes decode's signature or the return_format="parsed" contract, every audio yields segments = [], the stage records a status-ok outcome with an empty result, and run_task_cached persists it — so a total backend break presents as "no speech detected" and survives subsequent runs. Consider re-raising when every audio in a batch fails to parse.

6. clustering.py:37 — low: the role-only prefix list is now duplicated in four places.
ROLE_LABEL_ONLY_PREFIXES restates api.py's _CHILD_ADULT_PREFIXES as an independent literal, and stages.py:175 / model.py:335 similarly hand-copy "microsoft/VibeVoice-ASR". The plural name invites a second role-only backend; when one arrives, a maintainer wiring dispatch in api.py gets no signal that clustering.py, identity.py and presence.py also need updating — silently reintroducing the identity-axis corruption this fix set removed. Deriving the workflow-side guard from the dispatch constants would make the coupling explicit.

7. speaker_diarization_test.py:27 — low: _cache_dir() added an import-time filesystem side effect.
_cache_dir() calls cache.mkdir(parents=True, exist_ok=True), which the old Path.home() / ".cache" / … expression did not. On a host where the venv-cache root isn't writable (read-only HOME, sandboxed CI), collection of the whole module raises instead of the tests skipping. The gate only needs the path, not its creation.

Also verified, no problem found

hf_subprocess_env's new warning (logger is defined at dependencies.py:14; rid/rev are always bound because repos always contains the primary repo); MOSS's new result["generated_tokens"] truncation check (upstream generate_transcription does return that key, counting new tokens only); VibeVoice's output_ids[0, inputs["input_ids"].shape[1]:] slice and the Start/End/Speaker/Content keys (both match the official model-card example); concatenate_audios reachable via the star-import in preprocessing/__init__.py; uv.lock pins transformers 5.5.4, consistent with the >=5.3 floor; MOSS's own metadata pins transformers>=5.6,<6 so the un-pinned venv requirement resolves correctly; ruff check passes on all touched paths.

The test suite was not run — both interpreters available in the review environment have incomplete dependency sets (.venv missing typing_extensions, the py3.10 env missing dotenv), so the CUDA-gated and venv-gated tests remain unexercised here.