Fix VibeVoice-ASR streaming via generate() and server mapping

mlx_audio/server.py

@@ -15,6 +15,7 @@
 import subprocess
 import time
 import webbrowser
+from collections.abc import Iterator
 from dataclasses import asdict, is_dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional
@@ -319,22 +320,34 @@ def generate_transcription_stream(stt_model, tmp_path: str, gen_kwargs: dict):
         result = stt_model.generate(tmp_path, **gen_kwargs)
 
         # Check if result is a generator (streaming mode)
-        if hasattr(result, "__iter__") and hasattr(result, "__next__"):
+        if isinstance(result, Iterator):
             accumulated_text = ""
             for chunk in result:
                 # Handle different chunk types (string tokens vs structured chunks)
                 if isinstance(chunk, str):
                     accumulated_text += chunk
                     chunk_data = {"text": chunk, "accumulated": accumulated_text}
+                elif isinstance(chunk, dict):
+                    text = chunk.get("text")
+                    if isinstance(text, str):
+                        accumulated_text += text
+                    chunk_data = dict(chunk)
+                    if isinstance(text, str):
+                        chunk_data.setdefault("accumulated", accumulated_text)
                 else:
                     # Structured chunk (e.g., Whisper streaming)
+                    text = getattr(chunk, "text", None)
+                    if isinstance(text, str):
+                        accumulated_text += text
                     chunk_data = {
-                        "text": chunk.text,
-                        "start": getattr(chunk, "start_time", None),
-                        "end": getattr(chunk, "end_time", None),
+                        "text": text,
+                        "start": getattr(chunk, "start_time", getattr(chunk, "start", None)),
+                        "end": getattr(chunk, "end_time", getattr(chunk, "end", None)),
                         "is_final": getattr(chunk, "is_final", None),
                         "language": getattr(chunk, "language", None),
                     }
+                    if isinstance(text, str):
+                        chunk_data["accumulated"] = accumulated_text
                 yield json.dumps(sanitize_for_json(chunk_data)) + "\n"
         else:
             # Not a generator, yield the full result
@@ -387,6 +400,14 @@ async def stt_transcriptions(
 
     # Filter kwargs to only include parameters the model's generate method accepts
     signature = inspect.signature(stt_model.generate)
+    # Map OpenAI-style stream flag to generation_stream when needed
+    if (
+        "generation_stream" in signature.parameters
+        and "stream" in gen_kwargs
+        and "stream" not in signature.parameters
+    ):
+        gen_kwargs["generation_stream"] = bool(gen_kwargs["stream"])
+        gen_kwargs.pop("stream", None)
     gen_kwargs = {k: v for k, v in gen_kwargs.items() if k in signature.parameters}
 
     return StreamingResponse(

mlx_audio/stt/models/vibevoice_asr/__init__.py

@@ -6,12 +6,13 @@
     Qwen2Config,
     SemanticTokenizerConfig,
 )
-from .vibevoice_asr import Model
+from .vibevoice_asr import Model, StreamingResult
 
 __all__ = [
     "Model",
     "ModelConfig",
     "AcousticTokenizerConfig",
     "SemanticTokenizerConfig",
     "Qwen2Config",
-]
+    "StreamingResult",
+]

mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py

@@ -3,8 +3,9 @@
 import re
 import time
 import warnings
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -16,6 +17,29 @@
 from .config import ModelConfig
 
 
+@dataclass
+class StreamingResult:
+    """Result object for streaming transcription.
+
+    Attributes:
+        text: Decoded text for this emission.
+        is_final: True if this is a final (committed) result, False if partial.
+        start_time: Start timestamp in seconds.
+        end_time: End timestamp in seconds.
+        language: Language of the transcription.
+        prompt_tokens: Total prompt tokens (only set on final result).
+        generation_tokens: Total generation tokens (only set on final result).
+    """
+
+    text: str
+    is_final: bool
+    start_time: float
+    end_time: float
+    language: str = "en"
+    prompt_tokens: int = 0
+    generation_tokens: int = 0
+
+
 class SpeechConnector(nn.Module):
     """
     MLP connector to project speech features to LM hidden dimension.
@@ -595,8 +619,9 @@ def generate(
         prefill_step_size: int = 2048,
         generation_stream: bool = False,
         verbose: bool = False,
+        stream: bool = False,
         **kwargs,
-    ) -> STTOutput:
+    ) -> Union[STTOutput, Generator[StreamingResult, None, None]]:
         """
         Generate transcription from audio.
 
@@ -614,12 +639,29 @@ def generate(
             prefill_step_size: Chunk size for prompt prefill (reduces peak memory)
             generation_stream: Enable streaming
             verbose: Print progress
+            stream: If True, return a generator that yields StreamingResult objects
 
         Returns:
-            STTOutput with transcription text and segments
+            STTOutput with transcription text and segments, or Generator[StreamingResult]
         """
         from mlx_lm.sample_utils import make_logits_processors, make_sampler
 
+        if stream:
+            return self.stream_transcribe(
+                audio,
+                context=context,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                min_p=min_p,
+                min_tokens_to_keep=min_tokens_to_keep,
+                repetition_penalty=repetition_penalty,
+                repetition_context_size=repetition_context_size,
+                prefill_step_size=prefill_step_size,
+                verbose=verbose,
+            )
+
         start_time = time.time()
 
         # Preprocess audio
@@ -705,7 +747,7 @@ def stream_transcribe(
         repetition_context_size: int = 100,
         prefill_step_size: int = 2048,
         verbose: bool = False,
-    ) -> Generator[str, None, None]:
+    ) -> Generator[StreamingResult, None, None]:
         """
         Stream transcription token-by-token from audio.
 
@@ -724,7 +766,7 @@ def stream_transcribe(
             verbose: Print progress
 
         Yields:
-            Decoded text chunks as they are generated.
+            StreamingResult objects with text, timing, and status information.
         """
         from mlx_lm.sample_utils import make_logits_processors, make_sampler
 
@@ -756,6 +798,8 @@ def stream_transcribe(
         )
 
         # Stream tokens
+        token_count = 0
+        total_prompt_tokens = input_ids.shape[1]
         for token, _ in self.stream_generate(
             input_ids=input_ids,
             speech_features=speech_features,
@@ -767,7 +811,30 @@ def stream_transcribe(
             verbose=verbose,
         ):
             text = self.tokenizer.decode([token])
-            yield text
+            prev_progress = token_count / max(max_tokens, 1)
+            token_count += 1
+            curr_progress = min(token_count / max(max_tokens, 1), 1.0)
+
+            estimated_start = audio_duration * prev_progress
+            estimated_end = audio_duration * curr_progress
+
+            yield StreamingResult(
+                text=text,
+                is_final=False,
+                start_time=estimated_start,
+                end_time=estimated_end,
+                language="en",
+            )
+
+        yield StreamingResult(
+            text="",
+            is_final=True,
+            start_time=0.0,
+            end_time=audio_duration,
+            language="en",
+            prompt_tokens=total_prompt_tokens,
+            generation_tokens=token_count,
+        )
 
         mx.clear_cache()