Upload folder using huggingface_hub

diarize_audio.py

@@ -1,837 +1,913 @@
-"""Speaker diarization + gender estimation for subtitle cues.
-
-This module answers "who spoke when" using pyannote.audio and attaches a stable
-speaker label (and a coarse male/female/unknown gender guess) to every SRT cue.
-Gemma then uses those labels to keep Vietnamese forms of address / pronouns
-(ta, hắn, y, muội muội, tỷ tỷ, phu quân, ...) consistent for each character.
-
-Pipeline:
-    extract_audio (ffmpeg -> 16k mono wav) ->
-    run_diarization (pyannote community-1) -> list[SpeakerTurn] ->
-    estimate_speaker_genders (median F0 per speaker, optional) ->
-    assign_speakers_to_cues (max temporal overlap)
-
-Heavy dependencies (pyannote.audio, torch, torchaudio) are imported lazily so
-that the rest of the project keeps working even when they are not installed.
-The Hugging Face access token is required by pyannote and is read, in order,
-from the explicit argument, then the ``HF_TOKEN`` / ``HUGGINGFACE_TOKEN`` env
-vars.
-"""
-
-from __future__ import annotations
-
-import os
-import subprocess
-import sys
-import tempfile
-import time
-from dataclasses import dataclass
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
-
-if TYPE_CHECKING:  # pragma: no cover - typing only
-    from translate_srt import Cue
-
-# Preferred open-source pipeline (pyannote.audio 4.x). Falls back to the legacy
-# 3.1 pipeline if the newer one is unavailable for the installed version.
-DEFAULT_PIPELINE = "pyannote/speaker-diarization-community-1"
-FALLBACK_PIPELINE = "pyannote/speaker-diarization-3.1"
-
-def _load_default_token() -> str:
-    """Local convenience token, read from a sibling file that is NOT uploaded.
-
-    Put your token in ``hf_token.local`` (or ``hf_token.txt``) next to this file
-    to have local runs pick it up automatically. The file is git-ignored and is
-    never copied by ``prepare_upload.ps1`` so the public code stays secret-free.
-    """
-    for name in ("hf_token.local", "hf_token.txt"):
-        p = Path(__file__).with_name(name)
-        try:
-            if p.is_file():
-                tok = p.read_text(encoding="utf-8").strip()
-                if tok:
-                    return tok
-        except OSError:
-            pass
-    return ""
-
-
-# Built-in fallback Hugging Face token (used when no arg/env token is supplied).
-# Loaded from a local-only file so the public source never embeds a secret.
-DEFAULT_HF_TOKEN = _load_default_token()
-
-_SUBPROCESS_FLAGS = subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
-
-# Coarse pitch thresholds (Hz) for median fundamental frequency per speaker.
-# A deliberate "unknown" band avoids confidently mislabelling ambiguous voices.
-# Male speech F0 is typically ~85-180 Hz and female ~165-255 Hz, so the bands
-# overlap. We label only confident cases and leave the overlap as "unknown" so
-# the translator infers gender from dialogue/context instead of guessing wrong.
-_MALE_F0_MAX = 150.0
-_FEMALE_F0_MIN = 195.0
-_F0_MIN_HZ = 65.0
-_F0_MAX_HZ = 400.0
-
-# Dedicated speech model that predicts age (0..100) and gender (female/male/child).
-DEFAULT_AGE_GENDER_MODEL = "audeering/wav2vec2-large-robust-24-ft-age-gender"
-# The model's gender head order is [female, male, child].
-_AGE_GENDER_LABELS = ("female", "male", "child")
-
-# Cached (model, processor, device) so we only load the ~1 GB model once.
-_AGE_GENDER_CACHE: dict[str, Any] = {}
-
-
-def age_to_group(age_years: float | None, gender: str = "") -> str:
-    """Bucket an age (years) into a coarse life-stage label.
-
-    Returns one of "child", "teen", "young_adult", "middle_aged", "elderly" or
-    "" when age is unknown.
-    """
-    if gender == "child":
-        return "child"
-    if age_years is None:
-        return ""
-    if age_years < 13:
-        return "child"
-    if age_years < 20:
-        return "teen"
-    if age_years < 40:
-        return "young_adult"
-    if age_years < 60:
-        return "middle_aged"
-    return "elderly"
-
-
-LogFn = Callable[[str], None]
-
-
-def _noop(_msg: str) -> None:
-    pass
-
-
-def _fmt_dur(seconds: float) -> str:
-    """Human-readable duration, e.g. '1h 23m' or '4m 12s'."""
-    if seconds < 0:
-        seconds = 0.0
-    s = int(round(seconds))
-    if s >= 3600:
-        h, rem = divmod(s, 3600)
-        m, sec = divmod(rem, 60)
-        return f"{h}h {m}m" if sec < 30 else f"{h}h {m}m {sec}s"
-    if s >= 60:
-        m, sec = divmod(s, 60)
-        return f"{m}m {sec}s"
-    return f"{s}s"
-
-
-def _log_step(log: LogFn, step: int, total: int, msg: str) -> None:
-    log(f"[diarize] [{step}/{total}] {msg}")
-
-
-def _audio_duration_sec(waveform, sample_rate: int) -> float:
-    try:
-        n = int(waveform.shape[-1])
-        return n / float(sample_rate) if sample_rate else 0.0
-    except Exception:  # noqa: BLE001
-        return 0.0
-
-
-@dataclass
-class SpeakerTurn:
-    start: float
-    end: float
-    speaker: str
-
-    @property
-    def duration(self) -> float:
-        return max(0.0, self.end - self.start)
-
-
-# --------------------------------------------------------------------------- #
-# Token handling
-# --------------------------------------------------------------------------- #
-def resolve_hf_token(token: str | None) -> str | None:
-    if token and token.strip():
-        return token.strip()
-    for env in ("HF_TOKEN", "HUGGINGFACE_TOKEN", "HUGGING_FACE_HUB_TOKEN"):
-        val = os.environ.get(env)
-        if val and val.strip():
-            return val.strip()
-    if DEFAULT_HF_TOKEN and DEFAULT_HF_TOKEN.strip():
-        return DEFAULT_HF_TOKEN.strip()
-    return None
-
-
-# --------------------------------------------------------------------------- #
-# Audio extraction
-# --------------------------------------------------------------------------- #
-def extract_audio(
-    video: Path,
-    out_wav: Path,
-    sample_rate: int = 16000,
-    max_seconds: float | None = None,
-) -> Path:
-    """Extract a 16 kHz mono PCM wav from *video* (what pyannote expects).
-
-    When *max_seconds* is given, only that many seconds from the start are
-    extracted (so diarization scales with the number of cues being processed).
-    """
-    cmd = [
-        "ffmpeg",
-        "-hide_banner",
-        "-loglevel",
-        "error",
-        "-i",
-        str(video),
-        "-vn",
-        "-ac",
-        "1",
-        "-ar",
-        str(sample_rate),
-    ]
-    if max_seconds and max_seconds > 0:
-        cmd += ["-t", f"{max_seconds:.3f}"]
-    cmd += [
-        "-c:a",
-        "pcm_s16le",
-        "-y",
-        str(out_wav),
-    ]
-    proc = subprocess.run(
-        cmd, capture_output=True, timeout=600, creationflags=_SUBPROCESS_FLAGS
-    )
-    if proc.returncode != 0:
-        raise RuntimeError(
-            "ffmpeg audio extraction failed: "
-            + proc.stderr.decode("utf-8", "replace").strip()
-        )
-    if not out_wav.exists():
-        raise RuntimeError(f"Audio extraction produced no file at {out_wav}")
-    return out_wav
-
-
-# --------------------------------------------------------------------------- #
-# Diarization
-# --------------------------------------------------------------------------- #
-def _pick_device(preferred: str = "auto") -> str:
-    if preferred and preferred != "auto":
-        return preferred
-    try:
-        import torch  # noqa: PLC0415
-
-        return "cuda" if torch.cuda.is_available() else "cpu"
-    except Exception:  # noqa: BLE001
-        return "cpu"
-
-
-def _patch_speechbrain_lazy(log: LogFn = _noop) -> None:
-    """Work around a speechbrain LazyModule bug that breaks pyannote on Windows.
-
-    ``speechbrain.utils.importutils.LazyModule.ensure_module`` tries to skip the
-    lazy import when the caller is ``inspect.py`` (e.g. ``inspect.getmodule``
-    probing ``__file__``), but it only matches the POSIX ``/inspect.py`` path. On
-    Windows the path uses ``\\`` so the guard misses, the lazy module (e.g.
-    ``speechbrain.integrations.k2_fsa``) is force-imported, and it raises
-    ``ImportError`` — which ``hasattr`` propagates, crashing pipeline loading.
-
-    We re-implement ``ensure_module`` with an OS-agnostic ``inspect.py`` check.
-    """
-    try:
-        import importlib as _importlib  # noqa: PLC0415
-        import inspect as _inspect  # noqa: PLC0415
-        import os as _os  # noqa: PLC0415
-        import sys as _sys  # noqa: PLC0415
-
-        from speechbrain.utils import importutils as iu  # noqa: PLC0415
-    except Exception:  # noqa: BLE001 - speechbrain not used by this pipeline
-        return
-
-    if getattr(iu.LazyModule, "_gemma_patched", False):
-        return
-
-    def ensure_module(self, stacklevel: int):
-        importer_frame = None
-        try:
-            importer_frame = _inspect.getframeinfo(_sys._getframe(stacklevel + 1))
-        except (AttributeError, ValueError):
-            importer_frame = None
-        if (
-            importer_frame is not None
-            and _os.path.basename(importer_frame.filename) == "inspect.py"
-        ):
-            raise AttributeError()
-        if self.lazy_module is None:
-            try:
-                if self.package is None:
-                    self.lazy_module = _importlib.import_module(self.target)
-                else:
-                    self.lazy_module = _importlib.import_module(
-                        f".{self.target}", self.package
-                    )
-            except Exception as e:  # noqa: BLE001
-                raise ImportError(f"Lazy import of {self!r} failed") from e
-        return self.lazy_module
-
-    iu.LazyModule.ensure_module = ensure_module
-    iu.LazyModule._gemma_patched = True
-    log("[diarize] applied speechbrain LazyModule Windows compatibility patch.")
-
-
-def _load_pipeline(hf_token: str, device: str, log: LogFn):
-    try:
-        from pyannote.audio import Pipeline  # noqa: PLC0415
-    except ImportError as exc:  # pragma: no cover - env dependent
-        raise RuntimeError(
-            "pyannote.audio is not installed. Install diarization extras:\n"
-            "  pip install -r requirements-diarize.txt"
-        ) from exc
-
-    import torch  # noqa: PLC0415
-
-    _patch_speechbrain_lazy(log)
-
-    last_err: Exception | None = None
-    for model_id in (DEFAULT_PIPELINE, FALLBACK_PIPELINE):
-        try:
-            log(f"[diarize] loading pipeline {model_id} ...")
-            pipeline = Pipeline.from_pretrained(model_id, token=hf_token)
-            if pipeline is None:
-                raise RuntimeError(
-                    "Pipeline.from_pretrained returned None — the HF token is "
-                    "likely invalid or the model conditions were not accepted at "
-                    f"https://hf.co/{model_id}"
-                )
-            try:
-                pipeline.to(torch.device(device))
-            except Exception:  # noqa: BLE001 - keep CPU pipeline if .to fails
-                pass
-            log(f"[diarize] pipeline ready: {model_id} on {device}.")
-            return pipeline
-        except Exception as exc:  # noqa: BLE001
-            last_err = exc
-            log(f"[diarize] could not load {model_id}: {exc}")
-    raise RuntimeError(
-        "Failed to load any diarization pipeline. Most common causes:\n"
-        "  1) HF token's account has not accepted ALL gated model conditions (Agree):\n"
-        f"       https://hf.co/{DEFAULT_PIPELINE}\n"
-        f"       https://hf.co/{FALLBACK_PIPELINE}\n"
-        "       https://hf.co/pyannote/segmentation-3.0\n"
-        "  2) On Colab: pytorch-lightning version mismatch (community-1).\n"
-        f"Last error: {last_err}"
-    )
-
-
-def load_wav(wav: Path):
-    """Load a wav into an in-memory (channel, time) float32 tensor + sample rate.
-
-    Uses ``soundfile`` instead of ``torchaudio.load`` / torchcodec so audio
-    decoding does not depend on a working torchcodec/ffmpeg-DLL setup.
-    """
-    import soundfile as sf  # noqa: PLC0415
-    import torch  # noqa: PLC0415
-
-    data, sr = sf.read(str(wav), dtype="float32", always_2d=True)  # (time, channels)
-    waveform = torch.from_numpy(data.T).contiguous()  # (channels, time)
-    if waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0, keepdim=True)
-    return waveform, int(sr)
-
-
-def _as_annotation(diarization):
-    """Return a pyannote Annotation (with ``itertracks``) from a pipeline result.
-
-    pyannote.audio 4.x (community-1) returns a ``DiarizeOutput`` exposing
-    ``speaker_diarization`` and ``exclusive_speaker_diarization``; the legacy 3.1
-    pipeline returns an ``Annotation`` directly. Exclusive diarization assigns at
-    most one speaker per instant, which is ideal for tagging subtitle cues.
-    """
-    if hasattr(diarization, "itertracks"):
-        return diarization
-    for attr in ("exclusive_speaker_diarization", "speaker_diarization"):
-        candidate = getattr(diarization, attr, None)
-        if candidate is not None and hasattr(candidate, "itertracks"):
-            return candidate
-    raise RuntimeError(
-        f"Unexpected diarization output type without itertracks: {type(diarization)!r}"
-    )
-
-
-def run_diarization(
-    waveform,
-    sample_rate: int,
-    hf_token: str,
-    num_speakers: int | None = None,
-    min_speakers: int | None = None,
-    max_speakers: int | None = None,
-    device: str = "auto",
-    log: LogFn = _noop,
-) -> list[SpeakerTurn]:
-    """Run pyannote diarization and return merged speaker turns sorted by time."""
-    dev = _pick_device(device)
-    pipeline = _load_pipeline(hf_token, dev, log)
-
-    kwargs: dict[str, int] = {}
-    if num_speakers and num_speakers > 0:
-        kwargs["num_speakers"] = num_speakers
-    else:
-        if min_speakers and min_speakers > 0:
-            kwargs["min_speakers"] = min_speakers
-        if max_speakers and max_speakers > 0:
-            kwargs["max_speakers"] = max_speakers
-
-    dur = _audio_duration_sec(waveform, sample_rate)
-    kw_hint = ", ".join(f"{k}={v}" for k, v in kwargs.items()) or "auto-detect speakers"
-    log(
-        f"[diarize] running pyannote on {_fmt_dur(dur)} audio "
-        f"({dev}, {kw_hint}) — có thể mất vài phút đến ~1h với phim dài..."
-    )
-    t0 = time.time()
-    audio = {"waveform": waveform, "sample_rate": sample_rate}
-    diarization = pipeline(audio, **kwargs)
-    log(f"[diarize] pyannote inference done in {_fmt_dur(time.time() - t0)}.")
-    annotation = _as_annotation(diarization)
-
-    turns: list[SpeakerTurn] = []
-    for segment, _track, speaker in annotation.itertracks(yield_label=True):
-        turns.append(SpeakerTurn(float(segment.start), float(segment.end), str(speaker)))
-    turns.sort(key=lambda t: (t.start, t.end))
-    n_spk = len({t.speaker for t in turns})
-    log(f"[diarize] found {n_spk} speaker(s) across {len(turns)} turn(s).")
-    return turns
-
-
-# --------------------------------------------------------------------------- #
-# Gender estimation (coarse, pitch based)
-# --------------------------------------------------------------------------- #
-def estimate_speaker_genders(
-    waveform,
-    sr: int,
-    turns: list[SpeakerTurn],
-    *,
-    max_seconds_per_speaker: float = 30.0,
-    log: LogFn = _noop,
-) -> dict[str, str]:
-    """Estimate male/female/unknown per speaker from median fundamental frequency.
-
-    This is a lightweight heuristic (no extra model download) meant only as a
-    hint for the translator. Returns a mapping ``{speaker_label: gender}`` where
-    gender is one of ``"male"``, ``"female"`` or ``"unknown"``. *waveform* is a
-    ``(channel, time)`` float32 tensor as returned by :func:`load_wav`.
-    """
-    try:
-        import torch  # noqa: PLC0415
-        import torchaudio.functional as AF  # noqa: PLC0415
-    except Exception as exc:  # noqa: BLE001
-        log(f"[diarize] gender estimation skipped (torch/torchaudio missing): {exc}")
-        return {}
-
-    if waveform.ndim == 1:
-        waveform = waveform.unsqueeze(0)
-    elif waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0, keepdim=True)
-
-    by_speaker: dict[str, list[SpeakerTurn]] = {}
-    for t in turns:
-        by_speaker.setdefault(t.speaker, []).append(t)
-
-    genders: dict[str, str] = {}
-    for speaker, spk_turns in by_speaker.items():
-        spk_turns = sorted(spk_turns, key=lambda t: t.duration, reverse=True)
-        chunks: list["torch.Tensor"] = []
-        used = 0.0
-        for t in spk_turns:
-            if used >= max_seconds_per_speaker:
-                break
-            s = max(0, int(t.start * sr))
-            e = min(waveform.shape[-1], int(t.end * sr))
-            if e - s < int(0.20 * sr):
-                continue
-            chunks.append(waveform[..., s:e])
-            used += (e - s) / sr
-        if not chunks:
-            genders[speaker] = "unknown"
-            continue
-        audio = torch.cat(chunks, dim=-1)
-        try:
-            pitch = AF.detect_pitch_frequency(audio, sr)
-        except Exception as exc:  # noqa: BLE001
-            log(f"[diarize] pitch detection failed for {speaker}: {exc}")
-            genders[speaker] = "unknown"
-            continue
-        flat = pitch.flatten()
-        voiced = flat[(flat >= _F0_MIN_HZ) & (flat <= _F0_MAX_HZ)]
-        if voiced.numel() < 5:
-            genders[speaker] = "unknown"
-            continue
-        median_f0 = float(voiced.median().item())
-        if median_f0 <= _MALE_F0_MAX:
-            gender = "male"
-        elif median_f0 >= _FEMALE_F0_MIN:
-            gender = "female"
-        else:
-            gender = "unknown"
-        genders[speaker] = gender
-        log(f"[diarize] {speaker}: median F0 ~{median_f0:.0f} Hz -> {gender}")
-    return genders
-
-
-# --------------------------------------------------------------------------- #
-# Age + gender estimation (dedicated wav2vec2 model)
-# --------------------------------------------------------------------------- #
-def _collect_speaker_audio(waveform, sr, turns, max_seconds_per_speaker=30.0):
-    """Group turns per speaker and concatenate up to N seconds of their audio.
-
-    Returns ``{speaker: 1-D float32 torch.Tensor}``.
-    """
-    import torch  # noqa: PLC0415
-
-    if waveform.ndim == 1:
-        waveform = waveform.unsqueeze(0)
-    elif waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0, keepdim=True)
-
-    by_speaker: dict[str, list[SpeakerTurn]] = {}
-    for t in turns:
-        by_speaker.setdefault(t.speaker, []).append(t)
-
-    out: dict[str, Any] = {}
-    for speaker, spk_turns in by_speaker.items():
-        spk_turns = sorted(spk_turns, key=lambda t: t.duration, reverse=True)
-        chunks = []
-        used = 0.0
-        for t in spk_turns:
-            if used >= max_seconds_per_speaker:
-                break
-            s = max(0, int(t.start * sr))
-            e = min(waveform.shape[-1], int(t.end * sr))
-            if e - s < int(0.20 * sr):
-                continue
-            chunks.append(waveform[..., s:e])
-            used += (e - s) / sr
-        if chunks:
-            out[speaker] = torch.cat(chunks, dim=-1).flatten()
-    return out
-
-
-def _load_age_gender_model(model_name: str, device: str, log: LogFn):
-    """Load (and cache) the audeering age/gender wav2vec2 model + processor."""
-    cache_key = f"{model_name}@{device}"
-    if cache_key in _AGE_GENDER_CACHE:
-        return _AGE_GENDER_CACHE[cache_key]
-
-    import torch  # noqa: PLC0415
-    import torch.nn as nn  # noqa: PLC0415
-    from transformers import Wav2Vec2Processor  # noqa: PLC0415
-    from transformers.models.wav2vec2.modeling_wav2vec2 import (  # noqa: PLC0415
-        Wav2Vec2Model,
-        Wav2Vec2PreTrainedModel,
-    )
-
-    class ModelHead(nn.Module):
-        def __init__(self, config, num_labels):
-            super().__init__()
-            self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-            self.dropout = nn.Dropout(config.final_dropout)
-            self.out_proj = nn.Linear(config.hidden_size, num_labels)
-
-        def forward(self, features):
-            x = self.dropout(features)
-            x = self.dense(x)
-            x = torch.tanh(x)
-            x = self.dropout(x)
-            return self.out_proj(x)
-
-    class AgeGenderModel(Wav2Vec2PreTrainedModel):
-        # This model has no tied weights, so the mapping is always empty.
-        _tied_weights_keys = {}
-
-        def __init__(self, config):
-            super().__init__(config)
-            self.wav2vec2 = Wav2Vec2Model(config)
-            self.age = ModelHead(config, 1)
-            self.gender = ModelHead(config, 3)
-            # transformers >=5 builds `all_tied_weights_keys` in post_init();
-            # older releases only have init_weights(). Support both.
-            if hasattr(self, "post_init"):
-                self.post_init()
-            else:
-                self.init_weights()
-            # `from_pretrained` calls `all_tied_weights_keys.keys()` while
-            # loading; guarantee it is a dict so the load never crashes.
-            if not isinstance(getattr(self, "all_tied_weights_keys", None), dict):
-                self.all_tied_weights_keys = {}
-
-        def forward(self, input_values):
-            outputs = self.wav2vec2(input_values)
-            hidden_states = outputs[0]
-            hidden_states = torch.mean(hidden_states, dim=1)
-            logits_age = self.age(hidden_states)
-            logits_gender = torch.softmax(self.gender(hidden_states), dim=1)
-            return hidden_states, logits_age, logits_gender
-
-    log(f"[diarize] loading age/gender model {model_name} (first run downloads ~1 GB)...")
-    processor = Wav2Vec2Processor.from_pretrained(model_name)
-    model = AgeGenderModel.from_pretrained(model_name)
-    tied = getattr(model, "all_tied_weights_keys", None)
-    if not isinstance(tied, dict):
-        model.all_tied_weights_keys = {}
-    model = model.to(torch.device(device)).eval()
-    _AGE_GENDER_CACHE[cache_key] = (model, processor, device)
-    log(f"[diarize] age/gender model ready on {device}.")
-    return _AGE_GENDER_CACHE[cache_key]
-
-
-def estimate_speaker_demographics_model(
-    waveform,
-    sr: int,
-    turns: list[SpeakerTurn],
-    *,
-    model_name: str = DEFAULT_AGE_GENDER_MODEL,
-    device: str = "auto",
-    max_seconds_per_speaker: float = 30.0,
-    log: LogFn = _noop,
-) -> dict[str, dict[str, Any]]:
-    """Predict gender (female/male/child) + age per speaker with a dedicated model.
-
-    Returns ``{speaker: {"gender": str, "age": float|None, "age_group": str}}``.
-    """
-    import torch  # noqa: PLC0415
-
-    dev = _pick_device(device)
-    model, processor, dev = _load_age_gender_model(model_name, dev, log)
-    audio_by_speaker = _collect_speaker_audio(
-        waveform, sr, turns, max_seconds_per_speaker
-    )
-    n_spk = len(audio_by_speaker)
-    log(f"[diarize] age/gender: predicting {n_spk} speaker(s)...")
-
-    out: dict[str, dict[str, Any]] = {}
-    for i, (speaker, audio) in enumerate(audio_by_speaker.items(), start=1):
-        log(f"[diarize] age/gender: [{i}/{n_spk}] {speaker}...")
-        signal = audio.detach().cpu().numpy().astype("float32")
-        inputs = processor(signal, sampling_rate=sr)
-        values = torch.from_numpy(
-            inputs["input_values"][0].reshape(1, -1)
-        ).to(torch.device(dev))
-        with torch.no_grad():
-            _hidden, logits_age, logits_gender = model(values)
-        age_years = float(logits_age[0].item()) * 100.0
-        gender_idx = int(torch.argmax(logits_gender[0]).item())
-        gender = _AGE_GENDER_LABELS[gender_idx]
-        age_group = age_to_group(age_years, gender)
-        out[speaker] = {
-            "gender": gender,
-            "age": round(age_years, 1),
-            "age_group": age_group,
-        }
-        log(
-            f"[diarize] {speaker}: ~{age_years:.0f}y, gender={gender} "
-            f"({age_group or 'n/a'})"
-        )
-    return out
-
-
-def estimate_speaker_demographics(
-    waveform,
-    sr: int,
-    turns: list[SpeakerTurn],
-    *,
-    method: str = "auto",
-    model_name: str = DEFAULT_AGE_GENDER_MODEL,
-    device: str = "auto",
-    log: LogFn = _noop,
-) -> dict[str, dict[str, Any]]:
-    """Estimate per-speaker demographics, dispatching by *method*.
-
-    method:
-        "model"  -> dedicated wav2vec2 age/gender model (accurate, age + gender).
-        "pitch"  -> lightweight F0 heuristic (gender only, no age).
-        "auto"   -> try the model, fall back to pitch on any failure.
-    Returns ``{speaker: {"gender","age","age_group"}}``.
-    """
-    if method in ("model", "auto"):
-        try:
-            log(f"[diarize] age/gender: loading model ({model_name})...")
-            return estimate_speaker_demographics_model(
-                waveform, sr, turns, model_name=model_name, device=device, log=log
-            )
-        except Exception as exc:  # noqa: BLE001
-            log(f"[diarize] age/gender model failed, using pitch heuristic: {exc}")
-
-    log("[diarize] age/gender: pitch heuristic (F0) per speaker...")
-    genders = estimate_speaker_genders(waveform, sr, turns, log=log)
-    return {
-        spk: {"gender": g, "age": None, "age_group": ""} for spk, g in genders.items()
-    }
-
-
-# --------------------------------------------------------------------------- #
-# Cue assignment
-# --------------------------------------------------------------------------- #
-def _overlap(a_start: float, a_end: float, b_start: float, b_end: float) -> float:
-    return max(0.0, min(a_end, b_end) - max(a_start, b_start))
-
-
-def _speaker_gender(info: Any) -> str:
-    return info if isinstance(info, str) else (info or {}).get("gender", "")
-
-
-def _speaker_age_group(info: Any) -> str:
-    return "" if isinstance(info, str) else (info or {}).get("age_group", "")
-
-
-def assign_speakers_to_cues(
-    cues: list["Cue"],
-    turns: list[SpeakerTurn],
-    demographics: dict[str, Any] | None = None,
-    *,
-    log: LogFn = _noop,
-) -> int:
-    """Tag each cue with the speaker that overlaps it most in time.
-
-    *demographics* maps a speaker label to either a gender string or a dict
-    ``{"gender","age_group",...}``. Returns the number of cues tagged.
-    """
-    demographics = demographics or {}
-    tagged = 0
-    total = len(cues)
-    log(f"[diarize] gán speaker vào {total} cue(s)...")
-    report_every = max(200, total // 20) if total else 200
-    for idx, cue in enumerate(cues, start=1):
-        best_speaker = ""
-        best_overlap = 0.0
-        for t in turns:
-            if t.end <= cue.start:
-                continue
-            if t.start >= cue.end:
-                break
-            ov = _overlap(cue.start, cue.end, t.start, t.end)
-            if ov > best_overlap:
-                best_overlap = ov
-                best_speaker = t.speaker
-        if best_speaker:
-            info = demographics.get(best_speaker, {})
-            cue.speaker = best_speaker
-            cue.speaker_gender = _speaker_gender(info)
-            cue.speaker_age_group = _speaker_age_group(info)
-            tagged += 1
-        if idx % report_every == 0 or idx == total:
-            log(f"[diarize] gán cue: {idx}/{total} ({tagged} tagged)...")
-    return tagged
-
-
-# --------------------------------------------------------------------------- #
-# High level entry point
-# --------------------------------------------------------------------------- #
-def diarize_and_tag_cues(
-    video: Path,
-    cues: list["Cue"],
-    *,
-    hf_token: str | None = None,
-    num_speakers: int | None = None,
-    min_speakers: int | None = None,
-    max_speakers: int | None = None,
-    device: str = "auto",
-    detect_gender: bool = True,
-    gender_method: str = "auto",
-    age_gender_model: str = DEFAULT_AGE_GENDER_MODEL,
-    log: LogFn = _noop,
-) -> dict[str, Any]:
-    """Run the full diarization flow and tag *cues* in place.
-
-    Returns the ``{speaker: {"gender","age","age_group"}}`` registry (possibly
-    empty). Raises RuntimeError with an actionable message when the HF token is
-    missing.
-    """
-    token = resolve_hf_token(hf_token)
-    if not token:
-        raise RuntimeError(
-            "Hugging Face token required for diarization. Provide --hf-token, set "
-            "the HF_TOKEN environment variable, and accept the model conditions at "
-            f"https://hf.co/{DEFAULT_PIPELINE}"
-        )
-
-    total_steps = 5 if detect_gender else 4
-    t_all = time.time()
-    dev = _pick_device(device)
-    n_cues = len(cues)
-    _log_step(
-        log,
-        1,
-        total_steps,
-        f"bắt đầu Pass 0 — {n_cues} cue(s), device={dev}, gender={gender_method if detect_gender else 'off'}",
-    )
-
-    tmp_dir = Path(tempfile.mkdtemp(prefix="gemma_diarize_"))
-    wav = tmp_dir / "audio.wav"
-    max_seconds = None
-    if cues:
-        max_seconds = max((c.end for c in cues), default=0.0) + 5.0
-    _log_step(
-        log,
-        2,
-        total_steps,
-        f"tách audio ffmpeg (16 kHz mono, tối đa {_fmt_dur(max_seconds or 0)})...",
-    )
-    t0 = time.time()
-    extract_audio(video, wav, max_seconds=max_seconds)
-    wav_mb = wav.stat().st_size / (1024 * 1024) if wav.is_file() else 0.0
-    log(f"[diarize] audio wav ready: {wav_mb:.0f} MB in {_fmt_dur(time.time() - t0)}.")
-
-    _log_step(log, 3, total_steps, "nạp waveform + chạy pyannote (who spoke when)...")
-    t0 = time.time()
-    waveform, sr = load_wav(wav)
-    aud_dur = _audio_duration_sec(waveform, sr)
-    log(f"[diarize] waveform: {_fmt_dur(aud_dur)} @ {sr} Hz.")
-
-    turns = run_diarization(
-        waveform,
-        sr,
-        token,
-        num_speakers=num_speakers,
-        min_speakers=min_speakers,
-        max_speakers=max_speakers,
-        device=device,
-        log=log,
-    )
-    log(f"[diarize] bước pyannote xong trong {_fmt_dur(time.time() - t0)}.")
-    if not turns:
-        log("[diarize] no speaker turns detected; skipping speaker tags.")
-        return {}
-
-    demographics: dict[str, Any] = {}
-    if detect_gender:
-        _log_step(
-            log,
-            4,
-            total_steps,
-            f"ước lượng giới/tuổi ({gender_method}) cho {len({t.speaker for t in turns})} speaker(s)...",
-        )
-        t0 = time.time()
-        try:
-            demographics = estimate_speaker_demographics(
-                waveform,
-                sr,
-                turns,
-                method=gender_method,
-                model_name=age_gender_model,
-                device=device,
-                log=log,
-            )
-            log(f"[diarize] demographics xong trong {_fmt_dur(time.time() - t0)}.")
-        except Exception as exc:  # noqa: BLE001
-            log(
-                f"[diarize] demographics failed ({exc}) — gán speaker không có giới/tuổi."
-            )
-
-    assign_step = 5 if detect_gender else 4
-    _log_step(log, assign_step, total_steps, "gán speaker + giới/tuổi vào từng cue...")
-    t0 = time.time()
-    tagged = assign_speakers_to_cues(cues, turns, demographics, log=log)
-    n_spk = len({c.speaker for c in cues if c.speaker})
-    log(
-        f"[diarize] Pass 0 xong — {tagged}/{n_cues} cue tagged, "
-        f"{n_spk} speaker(s), tổng {_fmt_dur(time.time() - t_all)} "
-        f"(gán cue: {_fmt_dur(time.time() - t0)})."
-    )
-    return demographics
+"""Speaker diarization + gender estimation for subtitle cues.
+
+This module answers "who spoke when" using pyannote.audio and attaches a stable
+speaker label (and a coarse male/female/unknown gender guess) to every SRT cue.
+Gemma then uses those labels to keep Vietnamese forms of address / pronouns
+(ta, hắn, y, muội muội, tỷ tỷ, phu quân, ...) consistent for each character.
+
+Pipeline:
+    extract_audio (ffmpeg -> 16k mono wav) ->
+    run_diarization (pyannote community-1) -> list[SpeakerTurn] ->
+    estimate_speaker_genders (median F0 per speaker, optional) ->
+    assign_speakers_to_cues (max temporal overlap)
+
+Heavy dependencies (pyannote.audio, torch, torchaudio) are imported lazily so
+that the rest of the project keeps working even when they are not installed.
+The Hugging Face access token is required by pyannote and is read, in order,
+from the explicit argument, then the ``HF_TOKEN`` / ``HUGGINGFACE_TOKEN`` env
+vars.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from translate_srt import Cue
+
+# Preferred open-source pipeline (pyannote.audio 4.x). Falls back to the legacy
+# 3.1 pipeline if the newer one is unavailable for the installed version.
+DEFAULT_PIPELINE = "pyannote/speaker-diarization-community-1"
+FALLBACK_PIPELINE = "pyannote/speaker-diarization-3.1"
+
+def _load_default_token() -> str:
+    """Local convenience token, read from a sibling file that is NOT uploaded.
+
+    Put your token in ``hf_token.local`` (or ``hf_token.txt``) next to this file
+    to have local runs pick it up automatically. The file is git-ignored and is
+    never copied by ``prepare_upload.ps1`` so the public code stays secret-free.
+    """
+    for name in ("hf_token.local", "hf_token.txt"):
+        p = Path(__file__).with_name(name)
+        try:
+            if p.is_file():
+                tok = p.read_text(encoding="utf-8").strip()
+                if tok:
+                    return tok
+        except OSError:
+            pass
+    return ""
+
+
+# Built-in fallback Hugging Face token (used when no arg/env token is supplied).
+# Loaded from a local-only file so the public source never embeds a secret.
+DEFAULT_HF_TOKEN = _load_default_token()
+
+_SUBPROCESS_FLAGS = subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
+
+# Coarse pitch thresholds (Hz) for median fundamental frequency per speaker.
+# A deliberate "unknown" band avoids confidently mislabelling ambiguous voices.
+# Male speech F0 is typically ~85-180 Hz and female ~165-255 Hz, so the bands
+# overlap. We label only confident cases and leave the overlap as "unknown" so
+# the translator infers gender from dialogue/context instead of guessing wrong.
+_MALE_F0_MAX = 150.0
+_FEMALE_F0_MIN = 195.0
+_F0_MIN_HZ = 65.0
+_F0_MAX_HZ = 400.0
+
+# Dedicated speech model that predicts age (0..100) and gender (female/male/child).
+DEFAULT_AGE_GENDER_MODEL = "audeering/wav2vec2-large-robust-24-ft-age-gender"
+# The model's gender head order is [female, male, child].
+_AGE_GENDER_LABELS = ("female", "male", "child")
+
+# Cached (model, processor, device) so we only load the ~1 GB model once.
+_AGE_GENDER_CACHE: dict[str, Any] = {}
+
+
+def age_to_group(age_years: float | None, gender: str = "") -> str:
+    """Bucket an age (years) into a coarse life-stage label.
+
+    Returns one of "child", "teen", "young_adult", "middle_aged", "elderly" or
+    "" when age is unknown.
+    """
+    if gender == "child":
+        return "child"
+    if age_years is None:
+        return ""
+    if age_years < 13:
+        return "child"
+    if age_years < 20:
+        return "teen"
+    if age_years < 40:
+        return "young_adult"
+    if age_years < 60:
+        return "middle_aged"
+    return "elderly"
+
+
+LogFn = Callable[[str], None]
+
+
+def _noop(_msg: str) -> None:
+    pass
+
+
+def _fmt_dur(seconds: float) -> str:
+    """Human-readable duration, e.g. '1h 23m' or '4m 12s'."""
+    if seconds < 0:
+        seconds = 0.0
+    s = int(round(seconds))
+    if s >= 3600:
+        h, rem = divmod(s, 3600)
+        m, sec = divmod(rem, 60)
+        return f"{h}h {m}m" if sec < 30 else f"{h}h {m}m {sec}s"
+    if s >= 60:
+        m, sec = divmod(s, 60)
+        return f"{m}m {sec}s"
+    return f"{s}s"
+
+
+def _log_step(log: LogFn, step: int, total: int, msg: str) -> None:
+    log(f"[diarize] [{step}/{total}] {msg}")
+
+
+def _audio_duration_sec(waveform, sample_rate: int) -> float:
+    try:
+        n = int(waveform.shape[-1])
+        return n / float(sample_rate) if sample_rate else 0.0
+    except Exception:  # noqa: BLE001
+        return 0.0
+
+
+@dataclass
+class SpeakerTurn:
+    start: float
+    end: float
+    speaker: str
+
+    @property
+    def duration(self) -> float:
+        return max(0.0, self.end - self.start)
+
+
+# --------------------------------------------------------------------------- #
+# Token handling
+# --------------------------------------------------------------------------- #
+def resolve_hf_token(token: str | None) -> str | None:
+    if token and token.strip():
+        return token.strip()
+    for env in ("HF_TOKEN", "HUGGINGFACE_TOKEN", "HUGGING_FACE_HUB_TOKEN"):
+        val = os.environ.get(env)
+        if val and val.strip():
+            return val.strip()
+    if DEFAULT_HF_TOKEN and DEFAULT_HF_TOKEN.strip():
+        return DEFAULT_HF_TOKEN.strip()
+    return None
+
+
+# --------------------------------------------------------------------------- #
+# Audio extraction
+# --------------------------------------------------------------------------- #
+def extract_audio(
+    video: Path,
+    out_wav: Path,
+    sample_rate: int = 16000,
+    max_seconds: float | None = None,
+) -> Path:
+    """Extract a 16 kHz mono PCM wav from *video* (what pyannote expects).
+
+    When *max_seconds* is given, only that many seconds from the start are
+    extracted (so diarization scales with the number of cues being processed).
+    """
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-i",
+        str(video),
+        "-vn",
+        "-ac",
+        "1",
+        "-ar",
+        str(sample_rate),
+    ]
+    if max_seconds and max_seconds > 0:
+        cmd += ["-t", f"{max_seconds:.3f}"]
+    cmd += [
+        "-c:a",
+        "pcm_s16le",
+        "-y",
+        str(out_wav),
+    ]
+    proc = subprocess.run(
+        cmd, capture_output=True, timeout=600, creationflags=_SUBPROCESS_FLAGS
+    )
+    if proc.returncode != 0:
+        raise RuntimeError(
+            "ffmpeg audio extraction failed: "
+            + proc.stderr.decode("utf-8", "replace").strip()
+        )
+    if not out_wav.exists():
+        raise RuntimeError(f"Audio extraction produced no file at {out_wav}")
+    return out_wav
+
+
+# --------------------------------------------------------------------------- #
+# Diarization
+# --------------------------------------------------------------------------- #
+def _pick_device(preferred: str = "auto") -> str:
+    if preferred and preferred != "auto":
+        return preferred
+    try:
+        import torch  # noqa: PLC0415
+
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    except Exception:  # noqa: BLE001
+        return "cpu"
+
+
+def _patch_speechbrain_lazy(log: LogFn = _noop) -> None:
+    """Work around a speechbrain LazyModule bug that breaks pyannote on Windows.
+
+    ``speechbrain.utils.importutils.LazyModule.ensure_module`` tries to skip the
+    lazy import when the caller is ``inspect.py`` (e.g. ``inspect.getmodule``
+    probing ``__file__``), but it only matches the POSIX ``/inspect.py`` path. On
+    Windows the path uses ``\\`` so the guard misses, the lazy module (e.g.
+    ``speechbrain.integrations.k2_fsa``) is force-imported, and it raises
+    ``ImportError`` — which ``hasattr`` propagates, crashing pipeline loading.
+
+    We re-implement ``ensure_module`` with an OS-agnostic ``inspect.py`` check.
+    """
+    try:
+        import importlib as _importlib  # noqa: PLC0415
+        import inspect as _inspect  # noqa: PLC0415
+        import os as _os  # noqa: PLC0415
+        import sys as _sys  # noqa: PLC0415
+
+        from speechbrain.utils import importutils as iu  # noqa: PLC0415
+    except Exception:  # noqa: BLE001 - speechbrain not used by this pipeline
+        return
+
+    if getattr(iu.LazyModule, "_gemma_patched", False):
+        return
+
+    def ensure_module(self, stacklevel: int):
+        importer_frame = None
+        try:
+            importer_frame = _inspect.getframeinfo(_sys._getframe(stacklevel + 1))
+        except (AttributeError, ValueError):
+            importer_frame = None
+        if (
+            importer_frame is not None
+            and _os.path.basename(importer_frame.filename) == "inspect.py"
+        ):
+            raise AttributeError()
+        if self.lazy_module is None:
+            try:
+                if self.package is None:
+                    self.lazy_module = _importlib.import_module(self.target)
+                else:
+                    self.lazy_module = _importlib.import_module(
+                        f".{self.target}", self.package
+                    )
+            except Exception as e:  # noqa: BLE001
+                raise ImportError(f"Lazy import of {self!r} failed") from e
+        return self.lazy_module
+
+    iu.LazyModule.ensure_module = ensure_module
+    iu.LazyModule._gemma_patched = True
+    log("[diarize] applied speechbrain LazyModule Windows compatibility patch.")
+
+
+def _load_pipeline(hf_token: str, device: str, log: LogFn):
+    try:
+        from pyannote.audio import Pipeline  # noqa: PLC0415
+    except ImportError as exc:  # pragma: no cover - env dependent
+        raise RuntimeError(
+            "pyannote.audio is not installed. Install diarization extras:\n"
+            "  pip install -r requirements-diarize.txt"
+        ) from exc
+
+    import torch  # noqa: PLC0415
+
+    _patch_speechbrain_lazy(log)
+
+    last_err: Exception | None = None
+    for model_id in (DEFAULT_PIPELINE, FALLBACK_PIPELINE):
+        try:
+            log(f"[diarize] loading pipeline {model_id} ...")
+            pipeline = Pipeline.from_pretrained(model_id, token=hf_token)
+            if pipeline is None:
+                raise RuntimeError(
+                    "Pipeline.from_pretrained returned None — the HF token is "
+                    "likely invalid or the model conditions were not accepted at "
+                    f"https://hf.co/{model_id}"
+                )
+            try:
+                pipeline.to(torch.device(device))
+            except Exception:  # noqa: BLE001 - keep CPU pipeline if .to fails
+                pass
+            log(f"[diarize] pipeline ready: {model_id} on {device}.")
+            return pipeline
+        except Exception as exc:  # noqa: BLE001
+            last_err = exc
+            log(f"[diarize] could not load {model_id}: {exc}")
+    raise RuntimeError(
+        "Failed to load any diarization pipeline. Most common causes:\n"
+        "  1) HF token's account has not accepted ALL gated model conditions (Agree):\n"
+        f"       https://hf.co/{DEFAULT_PIPELINE}\n"
+        f"       https://hf.co/{FALLBACK_PIPELINE}\n"
+        "       https://hf.co/pyannote/segmentation-3.0\n"
+        "  2) On Colab: pytorch-lightning version mismatch (community-1).\n"
+        f"Last error: {last_err}"
+    )
+
+
+def load_wav(wav: Path):
+    """Load a wav into an in-memory (channel, time) float32 tensor + sample rate.
+
+    Uses ``soundfile`` instead of ``torchaudio.load`` / torchcodec so audio
+    decoding does not depend on a working torchcodec/ffmpeg-DLL setup.
+    """
+    import soundfile as sf  # noqa: PLC0415
+    import torch  # noqa: PLC0415
+
+    data, sr = sf.read(str(wav), dtype="float32", always_2d=True)  # (time, channels)
+    waveform = torch.from_numpy(data.T).contiguous()  # (channels, time)
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    return waveform, int(sr)
+
+
+def _as_annotation(diarization):
+    """Return a pyannote Annotation (with ``itertracks``) from a pipeline result.
+
+    pyannote.audio 4.x (community-1) returns a ``DiarizeOutput`` exposing
+    ``speaker_diarization`` and ``exclusive_speaker_diarization``; the legacy 3.1
+    pipeline returns an ``Annotation`` directly. Exclusive diarization assigns at
+    most one speaker per instant, which is ideal for tagging subtitle cues.
+    """
+    if hasattr(diarization, "itertracks"):
+        return diarization
+    for attr in ("exclusive_speaker_diarization", "speaker_diarization"):
+        candidate = getattr(diarization, attr, None)
+        if candidate is not None and hasattr(candidate, "itertracks"):
+            return candidate
+    raise RuntimeError(
+        f"Unexpected diarization output type without itertracks: {type(diarization)!r}"
+    )
+
+
+def run_diarization(
+    waveform,
+    sample_rate: int,
+    hf_token: str,
+    num_speakers: int | None = None,
+    min_speakers: int | None = None,
+    max_speakers: int | None = None,
+    device: str = "auto",
+    log: LogFn = _noop,
+) -> list[SpeakerTurn]:
+    """Run pyannote diarization and return merged speaker turns sorted by time."""
+    dev = _pick_device(device)
+    pipeline = _load_pipeline(hf_token, dev, log)
+
+    kwargs: dict[str, int] = {}
+    if num_speakers and num_speakers > 0:
+        kwargs["num_speakers"] = num_speakers
+    else:
+        if min_speakers and min_speakers > 0:
+            kwargs["min_speakers"] = min_speakers
+        if max_speakers and max_speakers > 0:
+            kwargs["max_speakers"] = max_speakers
+
+    dur = _audio_duration_sec(waveform, sample_rate)
+    kw_hint = ", ".join(f"{k}={v}" for k, v in kwargs.items()) or "auto-detect speakers"
+    log(
+        f"[diarize] running pyannote on {_fmt_dur(dur)} audio "
+        f"({dev}, {kw_hint}) — có thể mất vài phút đến ~1h với phim dài..."
+    )
+    t0 = time.time()
+    audio = {"waveform": waveform, "sample_rate": sample_rate}
+    diarization = pipeline(audio, **kwargs)
+    log(f"[diarize] pyannote inference done in {_fmt_dur(time.time() - t0)}.")
+    annotation = _as_annotation(diarization)
+
+    turns: list[SpeakerTurn] = []
+    for segment, _track, speaker in annotation.itertracks(yield_label=True):
+        turns.append(SpeakerTurn(float(segment.start), float(segment.end), str(speaker)))
+    turns.sort(key=lambda t: (t.start, t.end))
+    n_spk = len({t.speaker for t in turns})
+    log(f"[diarize] found {n_spk} speaker(s) across {len(turns)} turn(s).")
+    return turns
+
+
+# --------------------------------------------------------------------------- #
+# Gender estimation (coarse, pitch based)
+# --------------------------------------------------------------------------- #
+def estimate_speaker_genders(
+    waveform,
+    sr: int,
+    turns: list[SpeakerTurn],
+    *,
+    max_seconds_per_speaker: float = 30.0,
+    log: LogFn = _noop,
+) -> dict[str, str]:
+    """Estimate male/female/unknown per speaker from median fundamental frequency.
+
+    This is a lightweight heuristic (no extra model download) meant only as a
+    hint for the translator. Returns a mapping ``{speaker_label: gender}`` where
+    gender is one of ``"male"``, ``"female"`` or ``"unknown"``. *waveform* is a
+    ``(channel, time)`` float32 tensor as returned by :func:`load_wav`.
+    """
+    try:
+        import torch  # noqa: PLC0415
+        import torchaudio.functional as AF  # noqa: PLC0415
+    except Exception as exc:  # noqa: BLE001
+        log(f"[diarize] gender estimation skipped (torch/torchaudio missing): {exc}")
+        return {}
+
+    if waveform.ndim == 1:
+        waveform = waveform.unsqueeze(0)
+    elif waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+
+    by_speaker: dict[str, list[SpeakerTurn]] = {}
+    for t in turns:
+        by_speaker.setdefault(t.speaker, []).append(t)
+
+    genders: dict[str, str] = {}
+    for speaker, spk_turns in by_speaker.items():
+        spk_turns = sorted(spk_turns, key=lambda t: t.duration, reverse=True)
+        chunks: list["torch.Tensor"] = []
+        used = 0.0
+        for t in spk_turns:
+            if used >= max_seconds_per_speaker:
+                break
+            s = max(0, int(t.start * sr))
+            e = min(waveform.shape[-1], int(t.end * sr))
+            if e - s < int(0.20 * sr):
+                continue
+            chunks.append(waveform[..., s:e])
+            used += (e - s) / sr
+        if not chunks:
+            genders[speaker] = "unknown"
+            continue
+        audio = torch.cat(chunks, dim=-1)
+        try:
+            pitch = AF.detect_pitch_frequency(audio, sr)
+        except Exception as exc:  # noqa: BLE001
+            log(f"[diarize] pitch detection failed for {speaker}: {exc}")
+            genders[speaker] = "unknown"
+            continue
+        flat = pitch.flatten()
+        voiced = flat[(flat >= _F0_MIN_HZ) & (flat <= _F0_MAX_HZ)]
+        if voiced.numel() < 5:
+            genders[speaker] = "unknown"
+            continue
+        median_f0 = float(voiced.median().item())
+        if median_f0 <= _MALE_F0_MAX:
+            gender = "male"
+        elif median_f0 >= _FEMALE_F0_MIN:
+            gender = "female"
+        else:
+            gender = "unknown"
+        genders[speaker] = gender
+        log(f"[diarize] {speaker}: median F0 ~{median_f0:.0f} Hz -> {gender}")
+    return genders
+
+
+# --------------------------------------------------------------------------- #
+# Age + gender estimation (dedicated wav2vec2 model)
+# --------------------------------------------------------------------------- #
+def _weighted_median(values: list[float], weights: list[float]) -> float:
+    """Weighted median — robust to a few outlier windows (e.g. noisy segments).
+
+    Falls back to the plain mean when weights are missing/zero.
+    """
+    pairs = [(v, w) for v, w in zip(values, weights) if w and w > 0]
+    if not pairs:
+        return sum(values) / len(values) if values else 0.0
+    pairs.sort(key=lambda p: p[0])
+    total = sum(w for _, w in pairs)
+    half = total / 2.0
+    acc = 0.0
+    for v, w in pairs:
+        acc += w
+        if acc >= half:
+            return v
+    return pairs[-1][0]
+
+
+def _collect_speaker_windows(
+    waveform,
+    sr,
+    turns,
+    *,
+    max_seconds_per_speaker: float = 45.0,
+    window_sec: float = 8.0,
+    min_sec: float = 0.6,
+    max_windows: int = 12,
+):
+    """Group turns per speaker and split them into short coherent windows.
+
+    Returns ``{speaker: [(1-D float32 torch.Tensor, duration_sec), ...]}``.
+
+    Unlike concatenating one long 30 s clip and mean-pooling it once (which
+    biases the predicted age toward the dataset mean and lets a couple of noisy
+    seconds flip the gender), we keep each window separate so the caller can
+    predict per window and aggregate robustly (weighted median age, weighted
+    mean gender probability). Longest turns are used first.
+    """
+    import torch  # noqa: PLC0415
+
+    if waveform.ndim == 1:
+        waveform = waveform.unsqueeze(0)
+    elif waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+
+    by_speaker: dict[str, list[SpeakerTurn]] = {}
+    for t in turns:
+        by_speaker.setdefault(t.speaker, []).append(t)
+
+    win = max(1, int(window_sec * sr))
+    min_len = max(1, int(min_sec * sr))
+    total_samples = waveform.shape[-1]
+
+    out: dict[str, list] = {}
+    for speaker, spk_turns in by_speaker.items():
+        spk_turns = sorted(spk_turns, key=lambda t: t.duration, reverse=True)
+        windows: list = []
+        used = 0.0
+        for t in spk_turns:
+            if used >= max_seconds_per_speaker or len(windows) >= max_windows:
+                break
+            s = max(0, int(t.start * sr))
+            e = min(total_samples, int(t.end * sr))
+            pos = s
+            while pos + min_len <= e:
+                if used >= max_seconds_per_speaker or len(windows) >= max_windows:
+                    break
+                w_end = min(e, pos + win)
+                if w_end - pos < min_len:
+                    break
+                seg = waveform[..., pos:w_end].flatten()
+                dur = (w_end - pos) / sr
+                windows.append((seg, dur))
+                used += dur
+                pos = w_end
+        if windows:
+            out[speaker] = windows
+    return out
+
+
+def _load_age_gender_model(model_name: str, device: str, log: LogFn):
+    """Load (and cache) the audeering age/gender wav2vec2 model + processor."""
+    cache_key = f"{model_name}@{device}"
+    if cache_key in _AGE_GENDER_CACHE:
+        return _AGE_GENDER_CACHE[cache_key]
+
+    import torch  # noqa: PLC0415
+    import torch.nn as nn  # noqa: PLC0415
+    from transformers import Wav2Vec2Processor  # noqa: PLC0415
+    from transformers.models.wav2vec2.modeling_wav2vec2 import (  # noqa: PLC0415
+        Wav2Vec2Model,
+        Wav2Vec2PreTrainedModel,
+    )
+
+    class ModelHead(nn.Module):
+        def __init__(self, config, num_labels):
+            super().__init__()
+            self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+            self.dropout = nn.Dropout(config.final_dropout)
+            self.out_proj = nn.Linear(config.hidden_size, num_labels)
+
+        def forward(self, features):
+            x = self.dropout(features)
+            x = self.dense(x)
+            x = torch.tanh(x)
+            x = self.dropout(x)
+            return self.out_proj(x)
+
+    class AgeGenderModel(Wav2Vec2PreTrainedModel):
+        # This model has no tied weights, so the mapping is always empty.
+        _tied_weights_keys = {}
+
+        def __init__(self, config):
+            super().__init__(config)
+            self.wav2vec2 = Wav2Vec2Model(config)
+            self.age = ModelHead(config, 1)
+            self.gender = ModelHead(config, 3)
+            # transformers >=5 builds `all_tied_weights_keys` in post_init();
+            # older releases only have init_weights(). Support both.
+            if hasattr(self, "post_init"):
+                self.post_init()
+            else:
+                self.init_weights()
+            # `from_pretrained` calls `all_tied_weights_keys.keys()` while
+            # loading; guarantee it is a dict so the load never crashes.
+            if not isinstance(getattr(self, "all_tied_weights_keys", None), dict):
+                self.all_tied_weights_keys = {}
+
+        def forward(self, input_values):
+            outputs = self.wav2vec2(input_values)
+            hidden_states = outputs[0]
+            hidden_states = torch.mean(hidden_states, dim=1)
+            logits_age = self.age(hidden_states)
+            logits_gender = torch.softmax(self.gender(hidden_states), dim=1)
+            return hidden_states, logits_age, logits_gender
+
+    log(f"[diarize] loading age/gender model {model_name} (first run downloads ~1 GB)...")
+    processor = Wav2Vec2Processor.from_pretrained(model_name)
+    model = AgeGenderModel.from_pretrained(model_name)
+    tied = getattr(model, "all_tied_weights_keys", None)
+    if not isinstance(tied, dict):
+        model.all_tied_weights_keys = {}
+    model = model.to(torch.device(device)).eval()
+    _AGE_GENDER_CACHE[cache_key] = (model, processor, device)
+    log(f"[diarize] age/gender model ready on {device}.")
+    return _AGE_GENDER_CACHE[cache_key]
+
+
+def estimate_speaker_demographics_model(
+    waveform,
+    sr: int,
+    turns: list[SpeakerTurn],
+    *,
+    model_name: str = DEFAULT_AGE_GENDER_MODEL,
+    device: str = "auto",
+    max_seconds_per_speaker: float = 30.0,
+    log: LogFn = _noop,
+) -> dict[str, dict[str, Any]]:
+    """Predict gender (female/male/child) + age per speaker with a dedicated model.
+
+    Returns ``{speaker: {"gender": str, "age": float|None, "age_group": str}}``.
+    """
+    import torch  # noqa: PLC0415
+
+    dev = _pick_device(device)
+    model, processor, dev = _load_age_gender_model(model_name, dev, log)
+    windows_by_speaker = _collect_speaker_windows(
+        waveform, sr, turns, max_seconds_per_speaker=max_seconds_per_speaker
+    )
+    n_spk = len(windows_by_speaker)
+    log(f"[diarize] age/gender: predicting {n_spk} speaker(s) (per-window)...")
+
+    out: dict[str, dict[str, Any]] = {}
+    for i, (speaker, windows) in enumerate(windows_by_speaker.items(), start=1):
+        ages: list[float] = []
+        weights: list[float] = []
+        gender_prob_sum = torch.zeros(len(_AGE_GENDER_LABELS))
+        for seg, dur in windows:
+            signal = seg.detach().cpu().numpy().astype("float32")
+            inputs = processor(signal, sampling_rate=sr)
+            values = torch.from_numpy(
+                inputs["input_values"][0].reshape(1, -1)
+            ).to(torch.device(dev))
+            with torch.no_grad():
+                _hidden, logits_age, logits_gender = model(values)
+            ages.append(float(logits_age[0].item()) * 100.0)
+            weights.append(dur)
+            gender_prob_sum += logits_gender[0].detach().cpu().float() * dur
+
+        total_w = sum(weights)
+        if total_w <= 0:
+            continue
+        mean_gender = gender_prob_sum / total_w
+        gender_idx = int(torch.argmax(mean_gender).item())
+        gender = _AGE_GENDER_LABELS[gender_idx]
+        gender_conf = float(mean_gender[gender_idx].item())
+        # Weighted median age is robust to a few outlier windows; mean-pooling
+        # one long clip instead pulled every speaker toward the dataset mean.
+        age_years = _weighted_median(ages, weights)
+        age_group = age_to_group(age_years, gender)
+        out[speaker] = {
+            "gender": gender,
+            "age": round(age_years, 1),
+            "age_group": age_group,
+        }
+        log(
+            f"[diarize] age/gender: [{i}/{n_spk}] {speaker}: ~{age_years:.0f}y, "
+            f"gender={gender} (conf {gender_conf:.2f}) ({age_group or 'n/a'}) "
+            f"from {len(windows)} window(s)/{total_w:.0f}s"
+        )
+    return out
+
+
+def estimate_speaker_demographics(
+    waveform,
+    sr: int,
+    turns: list[SpeakerTurn],
+    *,
+    method: str = "auto",
+    model_name: str = DEFAULT_AGE_GENDER_MODEL,
+    device: str = "auto",
+    log: LogFn = _noop,
+) -> dict[str, dict[str, Any]]:
+    """Estimate per-speaker demographics, dispatching by *method*.
+
+    method:
+        "model"  -> dedicated wav2vec2 age/gender model (accurate, age + gender).
+        "pitch"  -> lightweight F0 heuristic (gender only, no age).
+        "auto"   -> try the model, fall back to pitch on any failure.
+    Returns ``{speaker: {"gender","age","age_group"}}``.
+    """
+    if method in ("model", "auto"):
+        try:
+            log(f"[diarize] age/gender: loading model ({model_name})...")
+            return estimate_speaker_demographics_model(
+                waveform, sr, turns, model_name=model_name, device=device, log=log
+            )
+        except Exception as exc:  # noqa: BLE001
+            log(f"[diarize] age/gender model failed, using pitch heuristic: {exc}")
+
+    log("[diarize] age/gender: pitch heuristic (F0) per speaker...")
+    genders = estimate_speaker_genders(waveform, sr, turns, log=log)
+    return {
+        spk: {"gender": g, "age": None, "age_group": ""} for spk, g in genders.items()
+    }
+
+
+# --------------------------------------------------------------------------- #
+# Cue assignment
+# --------------------------------------------------------------------------- #
+def _overlap(a_start: float, a_end: float, b_start: float, b_end: float) -> float:
+    return max(0.0, min(a_end, b_end) - max(a_start, b_start))
+
+
+def _speaker_gender(info: Any) -> str:
+    return info if isinstance(info, str) else (info or {}).get("gender", "")
+
+
+def _speaker_age_group(info: Any) -> str:
+    return "" if isinstance(info, str) else (info or {}).get("age_group", "")
+
+
+def assign_speakers_to_cues(
+    cues: list["Cue"],
+    turns: list[SpeakerTurn],
+    demographics: dict[str, Any] | None = None,
+    *,
+    nearest_gap_tolerance: float = 0.4,
+    log: LogFn = _noop,
+) -> int:
+    """Tag each cue with the speaker that overlaps it most in time.
+
+    When a cue overlaps no turn (a small diarization gap), it is assigned to the
+    nearest turn within *nearest_gap_tolerance* seconds instead of being left
+    blank. *demographics* maps a speaker label to either a gender string or a
+    dict ``{"gender","age_group",...}``. Returns the number of cues tagged.
+    """
+    demographics = demographics or {}
+    tagged = 0
+    total = len(cues)
+    log(f"[diarize] gán speaker vào {total} cue(s)...")
+    report_every = max(200, total // 20) if total else 200
+    for idx, cue in enumerate(cues, start=1):
+        best_speaker = ""
+        best_overlap = 0.0
+        nearest_speaker = ""
+        nearest_gap = float("inf")
+        for t in turns:
+            if t.end <= cue.start:
+                gap = cue.start - t.end
+                if gap < nearest_gap:
+                    nearest_gap = gap
+                    nearest_speaker = t.speaker
+                continue
+            if t.start >= cue.end:
+                gap = t.start - cue.end
+                if gap < nearest_gap:
+                    nearest_gap = gap
+                    nearest_speaker = t.speaker
+                break
+            ov = _overlap(cue.start, cue.end, t.start, t.end)
+            if ov > best_overlap:
+                best_overlap = ov
+                best_speaker = t.speaker
+        if not best_speaker and nearest_speaker and nearest_gap <= nearest_gap_tolerance:
+            best_speaker = nearest_speaker
+        if best_speaker:
+            info = demographics.get(best_speaker, {})
+            cue.speaker = best_speaker
+            cue.speaker_gender = _speaker_gender(info)
+            cue.speaker_age_group = _speaker_age_group(info)
+            tagged += 1
+        if idx % report_every == 0 or idx == total:
+            log(f"[diarize] gán cue: {idx}/{total} ({tagged} tagged)...")
+    return tagged
+
+
+# --------------------------------------------------------------------------- #
+# High level entry point
+# --------------------------------------------------------------------------- #
+def diarize_and_tag_cues(
+    video: Path,
+    cues: list["Cue"],
+    *,
+    hf_token: str | None = None,
+    num_speakers: int | None = None,
+    min_speakers: int | None = None,
+    max_speakers: int | None = None,
+    device: str = "auto",
+    detect_gender: bool = True,
+    gender_method: str = "auto",
+    age_gender_model: str = DEFAULT_AGE_GENDER_MODEL,
+    log: LogFn = _noop,
+) -> dict[str, Any]:
+    """Run the full diarization flow and tag *cues* in place.
+
+    Returns the ``{speaker: {"gender","age","age_group"}}`` registry (possibly
+    empty). Raises RuntimeError with an actionable message when the HF token is
+    missing.
+    """
+    token = resolve_hf_token(hf_token)
+    if not token:
+        raise RuntimeError(
+            "Hugging Face token required for diarization. Provide --hf-token, set "
+            "the HF_TOKEN environment variable, and accept the model conditions at "
+            f"https://hf.co/{DEFAULT_PIPELINE}"
+        )
+
+    total_steps = 5 if detect_gender else 4
+    t_all = time.time()
+    dev = _pick_device(device)
+    n_cues = len(cues)
+    _log_step(
+        log,
+        1,
+        total_steps,
+        f"bắt đầu Pass 0 — {n_cues} cue(s), device={dev}, gender={gender_method if detect_gender else 'off'}",
+    )
+
+    tmp_dir = Path(tempfile.mkdtemp(prefix="gemma_diarize_"))
+    wav = tmp_dir / "audio.wav"
+    max_seconds = None
+    if cues:
+        max_seconds = max((c.end for c in cues), default=0.0) + 5.0
+    _log_step(
+        log,
+        2,
+        total_steps,
+        f"tách audio ffmpeg (16 kHz mono, tối đa {_fmt_dur(max_seconds or 0)})...",
+    )
+    t0 = time.time()
+    extract_audio(video, wav, max_seconds=max_seconds)
+    wav_mb = wav.stat().st_size / (1024 * 1024) if wav.is_file() else 0.0
+    log(f"[diarize] audio wav ready: {wav_mb:.0f} MB in {_fmt_dur(time.time() - t0)}.")
+
+    _log_step(log, 3, total_steps, "nạp waveform + chạy pyannote (who spoke when)...")
+    t0 = time.time()
+    waveform, sr = load_wav(wav)
+    aud_dur = _audio_duration_sec(waveform, sr)
+    log(f"[diarize] waveform: {_fmt_dur(aud_dur)} @ {sr} Hz.")
+
+    turns = run_diarization(
+        waveform,
+        sr,
+        token,
+        num_speakers=num_speakers,
+        min_speakers=min_speakers,
+        max_speakers=max_speakers,
+        device=device,
+        log=log,
+    )
+    log(f"[diarize] bước pyannote xong trong {_fmt_dur(time.time() - t0)}.")
+    if not turns:
+        log("[diarize] no speaker turns detected; skipping speaker tags.")
+        return {}
+
+    demographics: dict[str, Any] = {}
+    if detect_gender:
+        _log_step(
+            log,
+            4,
+            total_steps,
+            f"ước lượng giới/tuổi ({gender_method}) cho {len({t.speaker for t in turns})} speaker(s)...",
+        )
+        t0 = time.time()
+        try:
+            demographics = estimate_speaker_demographics(
+                waveform,
+                sr,
+                turns,
+                method=gender_method,
+                model_name=age_gender_model,
+                device=device,
+                log=log,
+            )
+            log(f"[diarize] demographics xong trong {_fmt_dur(time.time() - t0)}.")
+        except Exception as exc:  # noqa: BLE001
+            log(
+                f"[diarize] demographics failed ({exc}) — gán speaker không có giới/tuổi."
+            )
+
+    assign_step = 5 if detect_gender else 4
+    _log_step(log, assign_step, total_steps, "gán speaker + giới/tuổi vào từng cue...")
+    t0 = time.time()
+    tagged = assign_speakers_to_cues(cues, turns, demographics, log=log)
+    n_spk = len({c.speaker for c in cues if c.speaker})
+    log(
+        f"[diarize] Pass 0 xong — {tagged}/{n_cues} cue tagged, "
+        f"{n_spk} speaker(s), tổng {_fmt_dur(time.time() - t_all)} "
+        f"(gán cue: {_fmt_dur(time.time() - t0)})."
+    )
+    return demographics

requirements-diarize.txt

@@ -1,24 +1,24 @@
-# Optional dependencies for speaker diarization (Pass 0).
-# Only needed when you enable "Phân tích giọng nói nhân vật" / --diarize.
-#
-# Setup (Windows, NVIDIA GPU recommended):
-#   1. Install PyTorch with CUDA matching your driver, e.g.:
-#        pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu124
-#      (CPU-only also works but is much slower:)
-#        pip install torch torchaudio
-#   2. pip install -r requirements-diarize.txt
-#   3. Create a Hugging Face token at https://hf.co/settings/tokens
-#   4. Accept the model conditions at:
-#        https://hf.co/pyannote/speaker-diarization-community-1
-#   5. Set HF_TOKEN env var or paste the token into the app's "HF token" field.
-#
-# torchcodec is used by pyannote.audio 4.x for audio decoding and needs ffmpeg
-# available on PATH (this project already ships/uses ffmpeg).
-
-# Audio is decoded with soundfile (not torchcodec) for a robust Windows setup.
-pyannote.audio>=4.0
-torch
-torchaudio
-soundfile>=0.11
-# transformers is needed for the audeering age/gender model (--gender-method model).
-transformers>=4.40
+# Optional dependencies for speaker diarization (Pass 0).
+# Only needed when you enable "Phân tích giọng nói nhân vật" / --diarize.
+#
+# Setup (Windows, NVIDIA GPU recommended):
+#   1. Install PyTorch with CUDA matching your driver, e.g.:
+#        pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu124
+#      (CPU-only also works but is much slower:)
+#        pip install torch torchaudio
+#   2. pip install -r requirements-diarize.txt
+#   3. Create a Hugging Face token at https://hf.co/settings/tokens
+#   4. Accept the model conditions at:
+#        https://hf.co/pyannote/speaker-diarization-community-1
+#   5. Set HF_TOKEN env var or paste the token into the app's "HF token" field.
+#
+# torchcodec is used by pyannote.audio 4.x for audio decoding and needs ffmpeg
+# available on PATH (this project already ships/uses ffmpeg).
+
+# Audio is decoded with soundfile (not torchcodec) for a robust Windows setup.
+pyannote.audio>=4.0
+torch
+torchaudio
+soundfile>=0.11
+# transformers is needed for the audeering age/gender model (--gender-method model).
+transformers>=4.40

translate_srt.py

@@ -814,19 +814,22 @@ def build_translation_prompt(
         speaker_block = (
             "\nSPEAKER GUIDE (from voice diarization — who is speaking):\n"
             + speaker_registry
-            + "\n- Each cue's 'speaker' field tells you WHICH character is talking plus "
-            "their detected gender and age group: gender nam=male, nữ=female, "
-            "trẻ em=child, chưa rõ=unknown; age thiếu niên=teen, thanh niên=young adult, "
-            "trung niên=middle-aged, lớn tuổi=elderly.\n"
-            "- Pick ONE consistent persona (gender + age + pronouns/forms of address) for "
-            "each speaker label and reuse it every time that speaker talks.\n"
-            "- Use BOTH gender and age for natural Vietnamese address: e.g. elderly male "
-            "→ ông/lão/cụ/bác; elderly female → bà/cụ; middle-aged → chú/bác/cô/dì; young "
-            "adult male → anh/cậu; young adult female → cô/chị; child → bé/cháu/con. Match "
-            "older-vs-younger relationships (cha/mẹ vs con, bác vs cháu) to the age gaps.\n"
-            "- Do not flip a character's gender or age between cues.\n"
-            "- 'chưa rõ' or a missing age means it is uncertain: infer from dialogue, "
-            "names and the scene image instead of guessing blindly.\n"
+            + "\n- Each cue's 'speaker' field tells you WHICH character is talking. "
+            "Labels: gender nam=male, nữ=female, trẻ em=child, chưa rõ=unknown; age "
+            "thiếu niên=teen, thanh niên=young adult, trung niên=middle-aged, lớn "
+            "tuổi=elderly.\n"
+            "- GENDER is the reliable anchor: keep each speaker's gender consistent and "
+            "use it to pick the right gendered Vietnamese address — nam → anh/cậu/chú/"
+            "ông/lão...; nữ → cô/chị/dì/bà...; never flip a character's gender between "
+            "cues.\n"
+            "- AGE group is only a ROUGH HINT from voice and may be wrong. Treat it as a "
+            "clue, not a rule: YOU decide the most natural address by combining the age "
+            "hint with the dialogue, names, relationships and the scene image. If the "
+            "age hint clearly conflicts with the context, trust the context.\n"
+            "- 'chưa rõ' gender or a missing age means it is uncertain: infer everything "
+            "from dialogue, names and the scene image instead of guessing blindly.\n"
+            "- Pick ONE consistent persona (forms of address / pronouns) per speaker "
+            "label and reuse it every time that speaker talks.\n"
         )
     return (
         f"Translate the following subtitle cues from {src} into {target_lang}.\n"