Update custom model files, README, and requirements

.gitattributes

@@ -1,36 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-tokenizer.json filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+tokenizer_config.json -filter -diff -merge text

README.md

@@ -75,52 +75,51 @@ for r in results:
 
 ### Word-Level Timestamps
 
-Timestamps come from forced alignment of the transcript against the audio and are returned under a `words` key:
+`return_timestamps=True` times every word with [Qwen3-ForcedAligner-0.6B](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B-hf) and returns them under a `words` key. Audio of any length works: the pipeline transcribes it in 8-18 s chunks cut at pauses (the model trained on clips up to 19 s), aligns each chunk against its own transcript in batches, and places every word on the recording's timeline.
 
 ```python
 result = pipe("audio.wav", return_timestamps=True)
 print(result["text"])
 for word in result["words"]:
     print(word)
+#   {'word': 'Hi,', 'start': 0.0, 'end': 0.24}
 ```
 
 ### Speaker Diarization
 
-`return_speakers=True` adds speaker labels. It first transcribes and aligns words (it implies `return_timestamps=True`). It then splits the audio into speech regions with [TEN VAD](https://github.com/TEN-framework/ten-vad), embeds each region with [ECAPA-TDNN](https://huggingface.co/speechbrain/spkrec-ecapa-voxceleb), and clusters the embeddings into speakers. Each word is assigned to the speaker whose segment it overlaps.
+`return_speakers=True` labels every word with a speaker (it implies `return_timestamps=True`). Speakers come from [Nemotron-3-Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization), run once over the whole recording, so labels stay consistent across hour-long meetings; it handles overlapping speech and up to 8 speakers. Each word goes to the speaker most active while it was spoken.
 
-Diarization needs a few extra packages:
+Nemotron-3-Diarization is in transformers `main`, not yet a release:
 
 ```bash
-pip install speechbrain ten-vad scikit-learn scipy
+pip install git+https://github.com/huggingface/transformers
 ```
 
 ```python
-result = pipe("meeting.wav", return_speakers=True, num_speakers=2)
+result = pipe("meeting.wav", return_speakers=True)
 
 print(result["text"])
 
-# Speaker turns
+# Speaker turns (may overlap)
 for seg in result["speaker_segments"]:
     print(f"{seg['start']:6.2f}-{seg['end']:6.2f}  {seg['speaker']}")
 #   0.00-  2.90  SPEAKER_0
 #   3.36-  6.47  SPEAKER_1
-#   6.93-  9.81  SPEAKER_0
 
 # Words with timestamps and speakers
 for w in result["words"]:
     print(f"{w['start']:6.2f}-{w['end']:6.2f}  {w['speaker']}  {w['word']}")
 #   0.00-  0.19  SPEAKER_0  Hi,
 #   0.22-  0.69  SPEAKER_0  Daniel.
-#   ...
 ```
 
-**Setting the speaker count.** Pass `num_speakers` when you know it. Otherwise the number of speakers is estimated from the audio, and on short clips that estimate tends to split one voice into several speakers. You can also bound the estimate with `min_speakers` / `max_speakers`:
+Speakers are numbered by when they first speak. The number of speakers is detected automatically. If you know it, pass `num_speakers` (exact) or `max_speakers` (an upper bound); any extra voices' words go to the closest kept speaker:
 
 ```python
-result = pipe("panel.wav", return_speakers=True, min_speakers=2, max_speakers=4)
+result = pipe("call.wav", return_speakers=True, num_speakers=2)
 ```
 
-If diarization fails (for example, a missing package), transcription still succeeds. The error is returned under `result["diarization_error"]` instead of raising.
+If alignment or diarization fails (for example, transformers without Nemotron), the transcript is still returned and the error is reported under `result["timestamp_error"]` or `result["diarization_error"]`.
 
 ### GPU Inference
 

alignment.py

@@ -1,4 +1,4 @@
-"""Forced alignment for word-level timestamps using Wav2Vec2."""
+"""Forced alignment for word-level timestamps: Qwen3-ForcedAligner, wav2vec2 as fallback."""
 
 import numpy as np
 import torch
@@ -240,3 +240,112 @@ class ForcedAligner:
             )
 
         return word_timestamps
+
+
+class QwenForcedAligner:
+    """Word timestamps from Qwen3-ForcedAligner-0.6B (transformers >= 5.17).
+
+    Trained for timestamping rather than repurposed from CTC ASR, it aligns
+    cased, punctuated text with numbers directly. One forward pass covers up to
+    5 minutes of speech, so long recordings are aligned as chunks, each against
+    its own transcript (`align_chunks`), batched.
+
+    Returns `[{"word", "start", "end"}]` for every word it can align, in order,
+    with the ORIGINAL word (punctuation, casing) as "word". The aligner itself
+    sees words with punctuation stripped and drops any that strip to nothing
+    ("--"); the split is mirrored word by word so the pairing stays exact.
+    """
+
+    MODEL_ID = "Qwen/Qwen3-ForcedAligner-0.6B-hf"
+    MAX_SECONDS = 300.0
+    BATCH_SIZE = 8
+    _model = None
+    _processor = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._model is None:
+            from transformers import AutoProcessor, Qwen3ASRForTokenClassification
+
+            device = _get_device()
+            # Batches are padded, and Metal sdpa returns NaN for fully masked rows.
+            attn = "eager" if device == "mps" else "sdpa"
+            model = Qwen3ASRForTokenClassification.from_pretrained(
+                cls.MODEL_ID, dtype=torch.bfloat16, attn_implementation=attn
+            )
+            cls._model = model.to(device).eval()
+            cls._processor = AutoProcessor.from_pretrained(cls.MODEL_ID)
+        return cls._model, cls._processor
+
+    @staticmethod
+    def _to_16k(audio, sample_rate: int) -> np.ndarray:
+        if isinstance(audio, torch.Tensor):
+            audio = audio.cpu().numpy()
+        audio = np.asarray(audio, dtype=np.float32).reshape(-1)
+        if sample_rate != 16000:
+            import torchaudio
+
+            audio = torchaudio.functional.resample(
+                torch.from_numpy(audio), sample_rate, 16000
+            ).numpy()
+        return audio
+
+    @classmethod
+    @torch.inference_mode()
+    def align_chunks(
+        cls,
+        chunks: list[tuple[np.ndarray, str]],
+        sample_rate: int = 16000,
+        language: str = "English",
+    ) -> list[list[dict]]:
+        """Align each `(audio, text)` pair (each <= 5 min); one word list per pair."""
+        from transformers.models.qwen3_asr.processing_qwen3_asr import _clean_tokens
+
+        results: list[list[dict]] = [[] for _ in chunks]
+        todo = []
+        for i, (raw, text) in enumerate(chunks):
+            kept = [w for w in text.split() if _clean_tokens([w])]
+            if not kept:
+                continue
+            audio = cls._to_16k(raw, sample_rate)
+            if len(audio) / 16000 > cls.MAX_SECONDS:
+                raise ValueError(
+                    f"chunk {i} is {len(audio) / 16000:.0f} s; Qwen3-ForcedAligner takes at most "
+                    f"{cls.MAX_SECONDS:.0f} s -- cut the audio first"
+                )
+            todo.append((i, audio, kept))
+        if not todo:  # nothing to time: don't load the model
+            return results
+        model, processor = cls.get_instance()
+        for b in range(0, len(todo), cls.BATCH_SIZE):
+            batch = todo[b : b + cls.BATCH_SIZE]
+            inputs, word_lists = processor.prepare_forced_aligner_inputs(
+                [audio for _, audio, _ in batch],
+                [" ".join(kept) for _, _, kept in batch],
+                language=language,
+                return_tensors="pt",
+                processor_kwargs={"padding": True},
+            )
+            inputs = {
+                k: (
+                    v.to(model.device, dtype=model.dtype)
+                    if v.is_floating_point()
+                    else v.to(model.device)
+                )
+                for k, v in inputs.items()
+            }
+            logits = model(**inputs).logits
+            timed = processor.decode_forced_alignment(
+                logits, inputs["input_ids"], word_lists, model.config.timestamp_token_id
+            )
+            for (i, _, kept), words in zip(batch, timed, strict=True):
+                results[i] = [
+                    {"word": w, "start": float(t["start_time"]), "end": float(t["end_time"])}
+                    for w, t in zip(kept, words, strict=True)
+                ]
+        return results
+
+    @classmethod
+    def align(cls, audio, text: str, sample_rate: int = 16000, language: str = "English"):
+        """ForcedAligner.align's interface for one clip of at most 5 minutes."""
+        return cls.align_chunks([(audio, text)], sample_rate, language)[0]

asr_pipeline.py

@@ -10,18 +10,49 @@ import transformers
 from transformers.pipelines.audio_utils import ffmpeg_read
 
 try:
-    from .alignment import ForcedAligner
+    from .alignment import ForcedAligner, QwenForcedAligner
     from .asr_modeling import ASRModel
     from .asr_processing import prepend_lead_in
-    from .diarization import SpeakerDiarizer
+    from .diarization import NemotronDiarizer, SpeakerDiarizer
 except ImportError:
-    from alignment import ForcedAligner  # type: ignore[no-redef]
+    from alignment import ForcedAligner, QwenForcedAligner  # type: ignore[no-redef]
     from asr_modeling import ASRModel  # type: ignore[no-redef]
     from asr_processing import prepend_lead_in  # type: ignore[no-redef]
-    from diarization import SpeakerDiarizer  # type: ignore[no-redef]
+    from diarization import NemotronDiarizer, SpeakerDiarizer  # type: ignore[no-redef]
 
 # Re-export for backwards compatibility
-__all__ = ["ASRPipeline", "ForcedAligner", "SpeakerDiarizer"]
+__all__ = [
+    "ASRPipeline",
+    "ForcedAligner",
+    "NemotronDiarizer",
+    "QwenForcedAligner",
+    "SpeakerDiarizer",
+]
+
+# Timestamps and speakers transcribe in chunks cut at the quietest point between
+# these lengths. The model trained on clips of at most 19 s; 18 leaves room for
+# the inference lead-in. Short clips are one chunk, so their text is unchanged.
+CHUNK_MAX_S = 18.0
+CHUNK_MIN_S = 8.0
+
+
+def chunk_bounds(
+    audio: np.ndarray, sample_rate: int, max_s: float = CHUNK_MAX_S, min_s: float = CHUNK_MIN_S
+) -> list[tuple[int, int]]:
+    """Sample ranges of at most `max_s`, each cut at the quietest 100 ms frame after `min_s`."""
+    frame = int(0.1 * sample_rate)
+    bounds, start, n = [], 0, len(audio)
+    while n - start > max_s * sample_rate:
+        lo = start + int(min_s * sample_rate)
+        hi = start + int(max_s * sample_rate)
+        k = (hi - lo) // frame
+        rms = np.sqrt(np.mean(np.square(audio[lo : lo + k * frame].reshape(k, frame)), axis=1))
+        cut = lo + int(np.argmin(rms)) * frame + frame // 2
+        bounds.append((start, cut))
+        start = cut
+    bounds.append((start, n))
+    return bounds
+
 
 _THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*", flags=re.DOTALL)
 _MIN_REPEATS = 3
@@ -73,25 +104,28 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
 
         Args:
             inputs: Audio input (file path, dict with array/sampling_rate, etc.)
-            return_timestamps: If True, return word-level timestamps using forced alignment
-            return_speakers: If True, return speaker labels for each word
+            return_timestamps: If True, return word-level timestamps (Qwen3-ForcedAligner)
+            return_speakers: If True, label every word with a speaker (Nemotron-3-Diarization)
             user_prompt: Custom transcription prompt (default: "Transcribe: ")
-            num_speakers: Exact number of speakers (if known, for diarization)
-            min_speakers: Minimum number of speakers (for diarization)
-            max_speakers: Maximum number of speakers (for diarization)
+            num_speakers: Exact number of speakers, if known
+            max_speakers: Upper bound on the number of speakers
             **kwargs: Additional arguments passed to the pipeline
 
         Returns:
             Dict with 'text' key, 'words' key if return_timestamps=True,
-            and speaker labels on words if return_speakers=True
+            and speaker labels on words plus 'speaker_segments' if return_speakers=True
         """
         # Extract our params before super().__call__ (which will also call _sanitize_parameters)
         return_timestamps = kwargs.pop("return_timestamps", False)
         return_speakers = kwargs.pop("return_speakers", False)
         user_prompt = kwargs.pop("user_prompt", None)
+        if kwargs.pop("min_speakers", None) is not None:
+            raise ValueError(
+                "min_speakers is not supported: Nemotron-3-Diarization decides how many "
+                "speakers it hears. Pass num_speakers (exact) or max_speakers instead."
+            )
         diarization_params = {
             "num_speakers": kwargs.pop("num_speakers", None),
-            "min_speakers": kwargs.pop("min_speakers", None),
             "max_speakers": kwargs.pop("max_speakers", None),
         }
 
@@ -109,9 +143,10 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             self.model.TRANSCRIBE_PROMPT = user_prompt
 
         try:
-            return self._transcribe(
+            if not return_timestamps:
+                return super().__call__(inputs, **kwargs)
+            return self._transcribe_timed(
                 inputs,
-                return_timestamps=return_timestamps,
                 return_speakers=return_speakers,
                 diarization_params=diarization_params,
                 **kwargs,
@@ -120,60 +155,59 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             if original_prompt is not None:
                 self.model.TRANSCRIBE_PROMPT = original_prompt
 
-    def _transcribe(
+    def _transcribe_timed(
         self,
         inputs,
         *,
-        return_timestamps: bool,
         return_speakers: bool,
         diarization_params: dict,
         **kwargs,
     ):
-        """Transcribe, then attach timestamps and speakers if requested."""
-        # Decode once for timestamp alignment and diarization, then hand the
-        # decoded array to the parent so it doesn't run ffmpeg over the same
-        # input a second time.
-        current_audio = (
-            self._extract_audio(inputs) if (return_timestamps or return_speakers) else None
-        )
+        """Transcribe in chunks, time every word, and label speakers if requested.
+
+        Chunks of 8-18 s cut at quiet points are transcribed one by one and each
+        is aligned against its own transcript (batched), then offset onto the
+        recording's timeline -- so audio of any length gets every word timed,
+        and the model never sees a clip longer than it trained on. Speakers come
+        from ONE Nemotron pass over the whole recording, so labels are
+        recording-level: chunking never touches speaker identity.
+        """
+        audio = self._extract_audio(inputs)
+        if audio is None:
+            raise ValueError(f"Cannot read audio from {type(inputs).__name__} for timestamps")
+        array = np.asarray(audio["array"], dtype=np.float32)
+        sr = audio.get("sampling_rate", 16000)
+
+        bounds = chunk_bounds(array, sr)
+        texts = [
+            super(ASRPipeline, self).__call__({"raw": array[s:e], "sampling_rate": sr}, **kwargs)[
+                "text"
+            ]
+            for s, e in bounds
+        ]
+        result: dict[str, Any] = {"text": " ".join(t for t in texts if t)}
 
-        # Run standard transcription
-        result = super().__call__(current_audio if current_audio is not None else inputs, **kwargs)
-
-        # Add timestamps if requested
-        if return_timestamps and current_audio is not None:
-            text = result.get("text", "")
-            if text:
-                try:
-                    words = ForcedAligner.align(
-                        current_audio["array"],
-                        text,
-                        sample_rate=current_audio.get("sampling_rate", 16000),
-                    )
-                    result["words"] = words
-                except Exception as e:
-                    result["words"] = []
-                    result["timestamp_error"] = str(e)
-            else:
-                result["words"] = []
-
-        # Add speaker diarization if requested
-        if return_speakers and current_audio is not None:
+        try:
+            aligned = QwenForcedAligner.align_chunks(
+                [(array[s:e], text) for (s, e), text in zip(bounds, texts)], sample_rate=sr
+            )
+            result["words"] = [
+                {**w, "start": w["start"] + s / sr, "end": w["end"] + s / sr}
+                for (s, _), chunk_words in zip(bounds, aligned)
+                for w in chunk_words
+            ]
+        except Exception as e:
+            result["words"] = []
+            result["timestamp_error"] = str(e)
+
+        if return_speakers:
             try:
-                # Run diarization
-                speaker_segments = SpeakerDiarizer.diarize(
-                    current_audio["array"],
-                    sample_rate=current_audio.get("sampling_rate", 16000),
-                    **{k: v for k, v in diarization_params.items() if v is not None},
+                activity = NemotronDiarizer.activity(array, sample_rate=sr)
+                keep = NemotronDiarizer.top_speakers(activity, **diarization_params)
+                result["speaker_segments"] = NemotronDiarizer.segments(activity, keep)
+                result["words"] = NemotronDiarizer.assign_speakers_to_words(
+                    result["words"], activity, keep
                 )
-                result["speaker_segments"] = speaker_segments
-
-                # Assign speakers to words
-                if result.get("words"):
-                    result["words"] = SpeakerDiarizer.assign_speakers_to_words(
-                        result["words"],
-                        speaker_segments,
-                    )
             except Exception as e:
                 result["speaker_segments"] = []
                 result["diarization_error"] = str(e)

diarization.py

@@ -1,4 +1,4 @@
-"""Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.
+"""Speaker diarization: Nemotron-3-Diarization, and TEN-VAD + ECAPA-TDNN + spectral clustering.
 
 Spectral clustering implementation adapted from FunASR/3D-Speaker:
 https://github.com/alibaba-damo-academy/FunASR
@@ -663,3 +663,125 @@ class SpeakerDiarizer:
             word["speaker"] = best_speaker
 
         return words
+
+
+class NemotronDiarizer:
+    """Speaker activity from NVIDIA Nemotron-3-Diarization, per 10 ms frame.
+
+    One offline pass over the whole recording: the model chunks internally
+    (27.2 s + 3.2 s look-ahead) and carries a speaker cache across chunks, so
+    speaker ids are recording-level for audio of any length, with no
+    embedding or clustering step. Overlap is native (several speakers can be
+    active in one frame). Tracks at most 8 speakers, numbered by arrival.
+
+    Needs a transformers build with `nemotron3_diarization` (main as of
+    2026-10, after 5.17).
+
+    Example:
+        >>> activity = NemotronDiarizer.activity(audio)  # (frames, 8) probabilities
+        >>> keep = NemotronDiarizer.top_speakers(activity, max_speakers=2)
+        >>> segments = NemotronDiarizer.segments(activity, keep)
+        >>> words = NemotronDiarizer.assign_speakers_to_words(words, activity, keep)
+    """
+
+    MODEL_ID = "nvidia/Nemotron-3-Diarization"
+    FRAME_S = 0.01
+    SEGMENT_THRESHOLD = 0.5  # a speaker's segment = frames above this
+    MIN_WORD_ACTIVITY = 0.1  # below this nobody is heard: the word inherits its neighbour
+    _model = None
+    _processor = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._model is None:
+            import importlib.util
+
+            if importlib.util.find_spec("transformers.models.nemotron3_diarization") is None:
+                raise ImportError(
+                    "Speaker diarization uses Nemotron-3-Diarization, which needs a transformers "
+                    "build with `nemotron3_diarization`: "
+                    "pip install git+https://github.com/huggingface/transformers"
+                )
+            from transformers import AutoModelForAudioFrameClassification, AutoProcessor
+
+            model = AutoModelForAudioFrameClassification.from_pretrained(cls.MODEL_ID)
+            cls._model = model.to(_get_device()).eval()
+            cls._processor = AutoProcessor.from_pretrained(cls.MODEL_ID)
+        return cls._model, cls._processor
+
+    @classmethod
+    @torch.inference_mode()
+    def activity(cls, audio: np.ndarray, sample_rate: int = 16000) -> np.ndarray:
+        """(frames, 8) speech probabilities at 10 ms; columns in arrival order."""
+        if sample_rate != 16000:
+            import torchaudio
+
+            audio = torchaudio.functional.resample(
+                torch.as_tensor(audio, dtype=torch.float32), sample_rate, 16000
+            ).numpy()
+        model, processor = cls.get_instance()
+        inputs = processor(np.asarray(audio, dtype=np.float32), sampling_rate=16000)
+        inputs = inputs.to(model.device, dtype=model.dtype)
+        return model(**inputs).logits[0].float().sigmoid().cpu().numpy()
+
+    @classmethod
+    def top_speakers(
+        cls,
+        activity: np.ndarray,
+        num_speakers: int | None = None,
+        max_speakers: int | None = None,
+    ) -> list[int]:
+        """Columns that count as speakers, in arrival order.
+
+        Every column that is ever above SEGMENT_THRESHOLD, capped by
+        `num_speakers` (exact, when known) or `max_speakers` to the columns
+        with the most speech. Capping is how a known count is applied:
+        a dropped column's words go to the most active remaining speaker.
+        """
+        active = [
+            c for c in range(activity.shape[1]) if (activity[:, c] > cls.SEGMENT_THRESHOLD).any()
+        ]
+        cap = num_speakers or max_speakers
+        if cap is not None and len(active) > cap:
+            mass = activity.sum(axis=0)
+            active = sorted(sorted(active, key=lambda c: -mass[c])[:cap])
+        return active or [0]
+
+    @classmethod
+    def segments(cls, activity: np.ndarray, keep: list[int]) -> list[dict]:
+        """Speaker turns `{"speaker", "start", "end"}` by start time; overlaps allowed."""
+        names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
+        out = []
+        for c in keep:
+            for on, s, e in _label_runs(activity[:, c] > cls.SEGMENT_THRESHOLD):
+                if on:
+                    out.append(
+                        {"speaker": names[c], "start": s * cls.FRAME_S, "end": e * cls.FRAME_S}
+                    )
+        return sorted(out, key=lambda seg: (seg["start"], seg["end"]))
+
+    @classmethod
+    def assign_speakers_to_words(
+        cls, words: list[dict], activity: np.ndarray, keep: list[int]
+    ) -> list[dict]:
+        """Give each word the kept speaker most active over its span.
+
+        A word where no kept speaker reaches MIN_WORD_ACTIVITY (aligner put it
+        in a pause) takes the previous word's speaker -- the next word's at
+        the start of the recording.
+        """
+        names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
+        cols = np.asarray(keep)
+        picked: list[str | None] = []
+        for word in words:
+            lo = min(int(word["start"] / cls.FRAME_S), len(activity) - 1)
+            hi = max(lo + 1, min(int(np.ceil(word["end"] / cls.FRAME_S)), len(activity)))
+            mean = activity[lo:hi, cols].mean(axis=0)
+            best = int(mean.argmax())
+            picked.append(names[keep[best]] if mean[best] >= cls.MIN_WORD_ACTIVITY else None)
+        known = [p for p in picked if p is not None]
+        last = known[0] if known else names[keep[0]]
+        for word, p in zip(words, picked):
+            last = p if p is not None else last
+            word["speaker"] = last
+        return words

handler.py

@@ -0,0 +1,85 @@
+"""Custom inference handler for HuggingFace Inference Endpoints."""
+
+from typing import Any
+
+import torch
+
+try:
+    # For remote execution, imports are relative
+    from .asr_modeling import ASRModel
+    from .asr_pipeline import ASRPipeline
+except ImportError:
+    # For local execution, imports are not relative
+    from asr_modeling import ASRModel  # type: ignore[no-redef]
+    from asr_pipeline import ASRPipeline  # type: ignore[no-redef]
+
+
+def _best_device() -> torch.device:
+    """Best available inference device (cuda > mps > cpu)."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+class EndpointHandler:
+    """HuggingFace Inference Endpoints handler for ASR model.
+
+    Handles model loading, warmup, and inference requests for deployment
+    on HuggingFace Inference Endpoints or similar services.
+    """
+
+    def __init__(self, path: str = ""):
+        """Initialize the endpoint handler.
+
+        Args:
+            path: Path to model directory or HuggingFace model ID
+        """
+        import os
+
+        import nltk
+
+        nltk.download("punkt_tab", quiet=True)
+
+        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+        # `ASRModel.from_pretrained` constructs its own submodules and forwards
+        # **kwargs to `ASRModel.__init__`, which discards them -- no loader
+        # reads `device_map`, `torch_dtype` or `low_cpu_mem_usage`, and it sets
+        # `_is_loading_from_pretrained` precisely to keep `device_map="auto"`
+        # out of the sub-model loaders. Passing them here did nothing, so the
+        # model stayed on CPU and a GPU Inference Endpoint silently decoded on
+        # CPU. Place it explicitly instead. dtype comes from
+        # `config.model_dtype` and the attention backend from
+        # `config.attn_implementation`, which already downgrades FA2 when
+        # flash_attn is missing.
+        self.model = ASRModel.from_pretrained(path)
+        self.device = _best_device()
+        self.model.to(self.device)
+        self.model.eval()
+
+        self.pipe = ASRPipeline(
+            model=self.model,
+            feature_extractor=self.model.feature_extractor,
+            tokenizer=self.model.tokenizer,
+            device=self.device,
+        )
+
+    def __call__(self, data: dict[str, Any]) -> dict[str, Any] | list[dict[str, Any]]:
+        """Process an inference request.
+
+        Args:
+            data: Request data containing 'inputs' (audio path/bytes) and optional 'parameters'
+
+        Returns:
+            Transcription result with 'text' key
+        """
+        inputs = data.get("inputs")
+        if inputs is None:
+            raise ValueError("Missing 'inputs' in request data")
+
+        # Pass through any parameters from request, let model config provide defaults
+        params = data.get("parameters", {})
+
+        return self.pipe(inputs, **params)

requirements.txt

@@ -0,0 +1,10 @@
+# Core dependencies for tiny-audio model inference
+# This file is pushed to HuggingFace for model repository
+
+# Transcription and word timestamps (Qwen3-ForcedAligner) need transformers >= 5.17.
+# Speaker diarization (return_speakers=True) uses Nemotron-3-Diarization, which is in
+# transformers main after 5.17 -- pinned to the commit the pipeline was verified on.
+transformers @ git+https://github.com/huggingface/transformers@91ea1192a00e0ebb684400434375fbd2345c8449
+torch
+torchaudio
+peft