feat(dub): optional speaker-count hint for diarization (#274) (#275)

When a clip has multiple speakers, pyannote's auto-detect sometimes collapses
them into a single "Speaker 1" — so the transcript merges turns and the dub
mixes voices. The diarization-consumption side is correct (overlap-weighted,
distinct Speaker N ids — pinned by a new test), so the collapse comes from
auto-detect itself.

Add an optional speaker-count hint (the reporter's own suggestion):
- backend: `/dub/transcribe-stream/{job_id}?num_speakers=N` (clamped 1–20;
  None → auto-detect) threaded to `diar_pipe(audio, num_speakers=N)`. Omitted
  entirely when unset so we don't depend on the kwarg in every pyannote build.
- frontend: `dubNumSpeakers` store field + a compact "Speakers" number input
  in the dub panel (placeholder "Auto") + i18n; `transcribeStreamUrl` appends
  the param; the SSE hook reads the hint at stream-open time.

Tests: tests/test_assign_speakers_from_diarization.py (multi-speaker split,
overlap weighting, label robustness, empty-result safety) +
dub.transcribeUrl.test.ts (param appended only for a positive int). Full
backend diarization + frontend suites pass; CJK i18n guard passes.

Does NOT close #274 — pending the reporter confirming that setting the count
resolves the collapse on their video (can't verify pyannote behaviour without
a CUDA box + the clip).

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

backend/api/routers/dub_core.py

@@ -366,14 +366,28 @@ _prep_event_helper = dub_pipeline.prep_event  # alias; we keep the module-local
 
 
 @router.get("/dub/transcribe-stream/{job_id}")
-async def dub_transcribe_stream(job_id: str):
+async def dub_transcribe_stream(job_id: str, num_speakers: Optional[int] = None):
     """Stream per-chunk segments via SSE, then emit diarized final pass.
 
     Pre-flight checks (missing job, missing audio, ASR not loaded) are emitted
     as in-stream `error` events rather than HTTP errors, because EventSource
     on the client can't read non-2xx response bodies — a 503 there surfaces
     as an opaque "network error" instead of the actionable message we want.
+
+    `num_speakers` is an optional hint passed straight to pyannote. Left unset,
+    pyannote auto-detects the count — but its auto-detect can collapse a
+    multi-speaker clip to a single speaker (issue #274). When the user knows
+    the exact count, supplying it forces pyannote to return that many speakers.
     """
+    # Clamp to a sane range; ignore anything non-positive / absurd so a bad
+    # query string can never break the diarization call. None → auto-detect.
+    if num_speakers is not None:
+        try:
+            num_speakers = int(num_speakers)
+            num_speakers = num_speakers if 1 <= num_speakers <= 20 else None
+        except (TypeError, ValueError):
+            num_speakers = None
+
     job = _get_job(job_id)
 
     preflight_error: Optional[str] = None
@@ -671,7 +685,15 @@ async def dub_transcribe_stream(job_id: str):
                     },
                 )
             try:
-                diar = diar_pipe(asr_audio_target)
+                # Pass the user's speaker-count hint through to pyannote when
+                # provided (#274). pyannote's apply() accepts num_speakers;
+                # omit it entirely when None so we don't depend on the kwarg
+                # existing in every pyannote build.
+                if num_speakers:
+                    logger.info("Diarizing with num_speakers=%d (user hint)", num_speakers)
+                    diar = diar_pipe(asr_audio_target, num_speakers=num_speakers)
+                else:
+                    diar = diar_pipe(asr_audio_target)
                 return assign_speakers_from_diarization(all_segments, diar), None
             except Exception as e:
                 logger.error(f"Diarization failed: {e}")

frontend/src/api/dub.transcribeUrl.test.ts

@@ -0,0 +1,24 @@
+import { describe, it, expect } from 'vitest';
+import { transcribeStreamUrl } from './dub';
+
+// #274: the optional speaker-count hint is appended only when it's a positive
+// integer; otherwise the backend auto-detects.
+describe('transcribeStreamUrl', () => {
+  it('omits num_speakers when not provided', () => {
+    expect(transcribeStreamUrl('job1')).toMatch(/\/dub\/transcribe-stream\/job1$/);
+  });
+
+  it('omits num_speakers for null / 0 / negative / NaN', () => {
+    for (const v of [null, undefined, 0, -3, NaN] as (number | null | undefined)[]) {
+      expect(transcribeStreamUrl('j', v)).not.toContain('num_speakers');
+    }
+  });
+
+  it('appends a positive integer hint', () => {
+    expect(transcribeStreamUrl('j', 3)).toContain('num_speakers=3');
+  });
+
+  it('floors a fractional hint', () => {
+    expect(transcribeStreamUrl('j', 2.9)).toContain('num_speakers=2');
+  });
+});

frontend/src/api/dub.ts

@@ -39,8 +39,14 @@ export async function dubIngestUrl(
   );
 }
 
-export function transcribeStreamUrl(jobId: string): string {
+export function transcribeStreamUrl(jobId: string, numSpeakers?: number | null): string {
-  return `${API}/dub/transcribe-stream/${jobId}`;
+  const base = `${API}/dub/transcribe-stream/${jobId}`;
+  // Optional pyannote speaker-count hint (#274). Only appended when a positive
+  // integer; otherwise the backend auto-detects.
+  if (numSpeakers && Number.isFinite(numSpeakers) && numSpeakers > 0) {
+    return `${base}?num_speakers=${Math.floor(numSpeakers)}`;
+  }
+  return base;
 }
 
 export async function dubAbort(jobId: string): Promise<void> {

frontend/src/hooks/useDubWorkflow.js

@@ -68,7 +68,11 @@ export default function useDubWorkflow({ loadProjects, loadProfiles, loadDubHist
 
   // ── SSE: wait for transcription stream ──
   const _waitForTranscribe = useCallback((jobId, ctrl) => new Promise((resolve, reject) => {
-    const evt = new EventSource(transcribeStreamUrl(jobId));
+    // Read the optional speaker-count hint at stream-open time (#274) so the
+    // user's choice for this job is honoured without threading it through the
+    // three call sites. null → pyannote auto-detect.
+    const numSpeakers = useAppStore.getState().dubNumSpeakers;
+    const evt = new EventSource(transcribeStreamUrl(jobId, numSpeakers));
     let gotFinal = false;
     const close = () => { try { evt.close(); } catch {} };
     const onAbortSignal = () => { close(); reject(Object.assign(new Error('aborted'), { name: 'AbortError' })); };

frontend/src/i18n/locales/en.json

@@ -450,6 +450,9 @@
     "export_wav": "WAV",
     "export_srt": "SRT",
     "upload_transcribe": "Upload & Transcribe",
+    "num_speakers_label": "Speakers",
+    "num_speakers_auto": "Auto",
+    "num_speakers_help": "How many speakers are in this video? Leave blank to auto-detect. Set a number if auto-detect merges multiple speakers into one.",
     "multi_lang": "Multi-lang",
     "prep_download": "Downloading video…",
     "prep_extract": "Extracting audio…",

frontend/src/pages/DubTab.css

@@ -352,6 +352,8 @@
 
 .dub-change-row       { display: flex; gap: 8px; margin-top: 8px; align-items: center; }
 .dub-change-row__cta  { flex: 1; margin-top: 0; }
+.dub-speakers-hint    { display: inline-flex; align-items: center; gap: 5px; font-size: 12px; color: var(--muted, #a89984); white-space: nowrap; }
+.dub-speakers-input   { width: 52px; margin-left: 4px; padding: 4px 6px; border-radius: 6px; border: 1px solid var(--border, #3c3836); background: var(--input-bg, #282828); color: inherit; font-size: 12px; }
 
 /* Compact disabled inputs in the idle skeleton */
 .input-base--xs              { font-size: 0.65rem; }

frontend/src/pages/DubTab.jsx

@@ -2,7 +2,7 @@ import React, { Suspense, lazy, useState, useEffect, useCallback, useRef } from
 import { copyText } from "../utils/copyText";
 import { useTranslation } from 'react-i18next';
 import {
-  PanelLeftOpen, PanelLeftClose, Film, Save, UploadCloud, Sparkles, Loader, Square,
+  PanelLeftOpen, PanelLeftClose, Film, Save, UploadCloud, Sparkles, Loader, Square, Users,
   FileText, Play, DownloadIcon, Volume2, Link2,
   Languages, ChevronDown, ChevronUp, Wand2, Trash2, Check, Globe, UserSquare2, User, AlertCircle,
   ExternalLink, Copy,
@@ -109,6 +109,8 @@ export default function DubTab(props) {
   const setDubLang        = useAppStore(s => s.setDubLang);
   const dubLangCode       = useAppStore(s => s.dubLangCode);
   const setDubLangCode    = useAppStore(s => s.setDubLangCode);
+  const dubNumSpeakers    = useAppStore(s => s.dubNumSpeakers);
+  const setDubNumSpeakers = useAppStore(s => s.setDubNumSpeakers);
   const dubInstruct       = useAppStore(s => s.dubInstruct);
   const setDubInstruct    = useAppStore(s => s.setDubInstruct);
   const dubTracks         = useAppStore(s => s.dubTracks);
@@ -392,6 +394,23 @@ export default function DubTab(props) {
                         />
                       </label>
                     )}
+                    <label className="dub-speakers-hint" title={t('dub.num_speakers_help')}>
+                      <Users size={13} /> {t('dub.num_speakers_label')}
+                      <input
+                        type="number"
+                        min={1}
+                        max={20}
+                        step={1}
+                        className="dub-speakers-input"
+                        placeholder={t('dub.num_speakers_auto')}
+                        value={dubNumSpeakers ?? ''}
+                        disabled={dubStep === 'uploading' || dubStep === 'transcribing'}
+                        onChange={(e) => {
+                          const v = parseInt(e.target.value, 10);
+                          setDubNumSpeakers(Number.isFinite(v) && v > 0 ? Math.min(v, 20) : null);
+                        }}
+                      />
+                    </label>
                     <button className="btn-primary dub-change-row__cta"
                       onClick={handleDubUpload}
                       disabled={dubStep === 'uploading' || dubStep === 'transcribing'}>

frontend/src/store/dubSlice.ts

@@ -87,6 +87,11 @@ export interface DubSlice {
   dubLang: string;
   dubLangCode: string;
 
+  // Optional speaker-count hint passed to pyannote diarization (#274). null =
+  // let pyannote auto-detect; a positive int forces that many speakers when
+  // auto-detect collapses a multi-speaker clip to one.
+  dubNumSpeakers: number | null;
+
   // ── Generation options ────────────────────────────────────────────────
   dubInstruct: string;
   preserveBg: boolean;
@@ -136,6 +141,7 @@ export interface DubSlice {
   setDubTracks: (v: Updater<string[]>) => void;
   setDubLang: (v: Updater<string>) => void;
   setDubLangCode: (v: Updater<string>) => void;
+  setDubNumSpeakers: (v: Updater<number | null>) => void;
   setDubInstruct: (v: Updater<string>) => void;
   setPreserveBg: (v: Updater<boolean>) => void;
   setDefaultTrack: (v: Updater<string>) => void;
@@ -152,7 +158,7 @@ const INITIAL: Omit<DubSlice,
   | 'setDubPrepProgress' | 'setDubCurrentSegId'
   | 'setDubProgress' | 'setDubError' | 'setDubFailure' | 'setIsTranslating' | 'setDubSegments'
   | 'setDubTranscript' | 'setDubFilename' | 'setDubDuration' | 'setDubTracks'
-  | 'setDubLang' | 'setDubLangCode' | 'setDubInstruct' | 'setPreserveBg'
+  | 'setDubLang' | 'setDubLangCode' | 'setDubNumSpeakers' | 'setDubInstruct' | 'setPreserveBg'
   | 'setDefaultTrack' | 'setExportTracks' | 'setPreviewSegIds' | 'setSpeakerClones'
   | 'setSegmentEffectPreset' | 'setAvailableEffectPresets' | 'resetDubState'
 > = {
@@ -174,6 +180,7 @@ const INITIAL: Omit<DubSlice,
   dubTracks: [],
   dubLang: 'Auto',
   dubLangCode: 'en',
+  dubNumSpeakers: null,
   dubInstruct: '',
   preserveBg: true,
   defaultTrack: 'original',
@@ -205,6 +212,7 @@ export const createDubSlice: StateCreator<DubSlice, [], [], DubSlice> = (set, ge
   setDubTracks:    (v) => set((s) => ({ dubTracks:    resolve(v, s.dubTracks) })),
   setDubLang:      (v) => set((s) => ({ dubLang:      resolve(v, s.dubLang) })),
   setDubLangCode:  (v) => set((s) => ({ dubLangCode:  resolve(v, s.dubLangCode) })),
+  setDubNumSpeakers: (v) => set((s) => ({ dubNumSpeakers: resolve(v, s.dubNumSpeakers) })),
   setDubInstruct:  (v) => set((s) => ({ dubInstruct:  resolve(v, s.dubInstruct) })),
   setPreserveBg:   (v) => set((s) => ({ preserveBg:   resolve(v, s.preserveBg) })),
   setDefaultTrack: (v) => set((s) => ({ defaultTrack: resolve(v, s.defaultTrack) })),

tests/test_assign_speakers_from_diarization.py

@@ -0,0 +1,80 @@
+"""`assign_speakers_from_diarization` overlap-weighting + label handling (#274).
+
+When pyannote returns multiple speakers, each transcript segment must get the
+speaker whose turns overlap it most — so a 2-speaker diarization yields two
+distinct `Speaker N` ids, NOT a collapse to one. (The single-speaker collapse
+users see comes from pyannote's *auto-detect*, which the new `num_speakers`
+hint addresses; this test pins that the consumption side is correct.)
+"""
+from dataclasses import dataclass
+
+from services.segmentation import assign_speakers_from_diarization
+
+
+@dataclass
+class _Turn:
+    start: float
+    end: float
+
+
+class _FakeDiarization:
+    """Mimics pyannote's `Annotation.itertracks(yield_label=True)`."""
+    def __init__(self, turns):
+        self._turns = turns  # list of (Turn, track_name, speaker_label)
+
+    def itertracks(self, yield_label=False):
+        for turn, track, spk in self._turns:
+            yield (turn, track, spk) if yield_label else (turn, track)
+
+
+def _segs(*spans):
+    return [{"start": a, "end": b, "speaker_id": "Speaker 1"} for a, b in spans]
+
+
+def test_two_speakers_yield_two_distinct_ids():
+    # Speaker_0 owns 0–5s, Speaker_1 owns 5–10s.
+    diar = _FakeDiarization([
+        (_Turn(0.0, 5.0), "A", "SPEAKER_00"),
+        (_Turn(5.0, 10.0), "B", "SPEAKER_01"),
+    ])
+    segs = _segs((0.5, 2.0), (6.0, 9.0))
+    out = assign_speakers_from_diarization(segs, diar)
+    assert out[0]["speaker_id"] == "Speaker 1"   # SPEAKER_00 -> idx 1
+    assert out[1]["speaker_id"] == "Speaker 2"   # SPEAKER_01 -> idx 2
+    assert out[0]["speaker_id"] != out[1]["speaker_id"]
+
+
+def test_overlap_weighted_winner():
+    # Segment 2–8 overlaps SPEAKER_00 for 3s (2–5) and SPEAKER_01 for 3s...
+    # but SPEAKER_01 owns 5–9 so overlap is 3s each — tie broken by max(); make
+    # SPEAKER_01 clearly dominant by extending its turn.
+    diar = _FakeDiarization([
+        (_Turn(0.0, 5.0), "A", "SPEAKER_00"),
+        (_Turn(5.0, 12.0), "B", "SPEAKER_01"),
+    ])
+    out = assign_speakers_from_diarization(_segs((4.0, 11.0)), diar)
+    # overlap: SPEAKER_00 = 1s (4–5), SPEAKER_01 = 6s (5–11) → winner SPEAKER_01
+    assert out[0]["speaker_id"] == "Speaker 2"
+
+
+def test_midpoint_fallback_when_no_overlap():
+    # Segment sits fully inside a single turn; overlap path still catches it,
+    # but a zero-length segment exercises the midpoint fallback.
+    diar = _FakeDiarization([(_Turn(0.0, 10.0), "A", "SPEAKER_02")])
+    out = assign_speakers_from_diarization([{"start": 3.0, "end": 3.0, "speaker_id": "Speaker 1"}], diar)
+    assert out[0]["speaker_id"] == "Speaker 3"
+
+
+def test_non_underscore_label_kept_verbatim():
+    # A label that isn't `<prefix>_<int>` must not crash — kept as-is.
+    diar = _FakeDiarization([(_Turn(0.0, 5.0), "A", "narrator")])
+    out = assign_speakers_from_diarization(_segs((1.0, 2.0)), diar)
+    assert out[0]["speaker_id"] == "narrator"
+
+
+def test_empty_diarization_leaves_segment_untouched():
+    # No turns → no winner → segment keeps whatever it had (heuristic fallback
+    # upstream owns that case).
+    diar = _FakeDiarization([])
+    out = assign_speakers_from_diarization(_segs((1.0, 2.0)), diar)
+    assert out[0]["speaker_id"] == "Speaker 1"