diff --git a/mlx_audio/server.py b/mlx_audio/server.py index c5bc9110c..86ece3533 100644 --- a/mlx_audio/server.py +++ b/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( diff --git a/mlx_audio/stt/models/vibevoice_asr/__init__.py b/mlx_audio/stt/models/vibevoice_asr/__init__.py index e6ba49519..1d92c91d4 100644 --- a/mlx_audio/stt/models/vibevoice_asr/__init__.py +++ b/mlx_audio/stt/models/vibevoice_asr/__init__.py @@ -6,7 +6,7 @@ Qwen2Config, SemanticTokenizerConfig, ) -from .vibevoice_asr import Model +from .vibevoice_asr import Model, StreamingResult __all__ = [ "Model", @@ -14,4 +14,5 @@ "AcousticTokenizerConfig", "SemanticTokenizerConfig", "Qwen2Config", -] + "StreamingResult", +] \ No newline at end of file diff --git a/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py b/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py index 76bd2229a..05c243ded 100644 --- a/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py +++ b/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()