Fix server (STT, TTS) (#460)

* Add python-multipart dependency to pyproject.toml for multipart form handling

* Refactor tokenizer usage in Whisper model files

- Removed direct imports and references to the Tokenizer class in favor of using model methods for tokenizer retrieval.
- Updated function signatures to accept a generic tokenizer parameter instead of a specific Tokenizer type.
- Enhanced the `detect_language` and `get_suppress_tokens` functions to utilize the model's tokenizer method, improving flexibility and reducing dependencies.
- Deleted unused tokenizer asset files to streamline the codebase.

* Add instruct and verbose fields to SpeechRequest model

- Updated the SpeechRequest model to include an optional 'instruct' field for additional instructions.
- Added a 'verbose' field to control output verbosity during audio generation.
- Refactored the generate_audio function to utilize the new fields, enhancing flexibility in audio processing.

* Add TranscriptionRequest model and update stt_transcriptions function

- Introduced a new TranscriptionRequest model to encapsulate parameters for audio transcription, including fields for language, verbosity, and streaming options.
- Updated the stt_transcriptions function to utilize the new model, enhancing parameter management and flexibility.
- Implemented logic to handle streaming results and filter generation parameters based on the model's signature, improving the transcription process.

* Add streaming transcription support and refactor STT model handling

- Introduced a new `generate_transcription_stream` function to handle streaming transcription, yielding results in real-time and managing temporary file cleanup.
- Updated the `stt_transcriptions` endpoint to utilize the new streaming function, enhancing the responsiveness of audio transcription.
- Refactored the `Qwen3ASRModel` to support streaming transcriptions, including a new `StreamingResult` data class for structured output.
- Improved model remapping logic in `get_model_category` to prioritize explicit remapping matches, enhancing model loading flexibility.
- Suppressed warnings during tokenizer loading in both `Qwen3ASRModel` and `ForcedAlignerModel` to improve user experience during model initialization.

* format

* fix tests

* Enhance chunk processing in Qwen3ASRModel with progress indication

- Added a tqdm progress bar to the chunk processing loop in the Qwen3ASRModel, improving user feedback during audio processing.
- The progress bar is configurable based on verbosity and the number of chunks, enhancing the overall user experience.

* Refactor audio chunk parameters in Qwen3ASRModel and related functions

- Renamed `max_chunk_sec` and `min_chunk_sec` to `chunk_duration` and `min_chunk_duration` respectively for clarity and consistency across the codebase.
- Updated function signatures and internal logic to reflect the new parameter names, ensuring proper handling of audio chunk durations.
- Adjusted test cases to align with the updated parameter names, maintaining test integrity.

* format

Author: Blaizzy

mlx_audio/server.py

@@ -7,6 +7,7 @@
 
 import argparse
 import asyncio
+import inspect
 import io
 import json
 import os
@@ -153,6 +154,7 @@ def setup_cors(app: FastAPI, allowed_origins: List[str]):
 class SpeechRequest(BaseModel):
     model: str
     input: str
+    instruct: str | None = None
     voice: str | None = None
     speed: float | None = 1.0
     gender: str | None = "male"
@@ -165,6 +167,20 @@ class SpeechRequest(BaseModel):
     top_k: int | None = 40
     repetition_penalty: float | None = 1.0
     response_format: str | None = "mp3"
+    verbose: bool = False
+
+
+class TranscriptionRequest(BaseModel):
+    model: str
+    language: str | None = None
+    verbose: bool = False
+    max_tokens: int = 128
+    chunk_duration: float = 30.0
+    frame_threshold: int = 25
+    stream: bool = False
+    context: str | None = None
+    prefill_step_size: int = 2048
+    text: str | None = None
 
 
 # Initialize the ModelProvider
@@ -234,7 +250,7 @@ async def remove_model(model_name: str):
         raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found")
 
 
-async def generate_audio(model, payload: SpeechRequest, verbose: bool = False):
+async def generate_audio(model, payload: SpeechRequest):
     # Load reference audio if provided
     ref_audio = payload.ref_audio
     if ref_audio and isinstance(ref_audio, str):
@@ -258,13 +274,15 @@ async def generate_audio(model, payload: SpeechRequest, verbose: bool = False):
         speed=payload.speed,
         gender=payload.gender,
         pitch=payload.pitch,
+        instruct=payload.instruct,
         lang_code=payload.lang_code,
         ref_audio=ref_audio,
         ref_text=payload.ref_text,
         temperature=payload.temperature,
         top_p=payload.top_p,
         top_k=payload.top_k,
         repetition_penalty=payload.repetition_penalty,
+        verbose=payload.verbose,
     ):
 
         sample_rate = result.sample_rate
@@ -287,25 +305,87 @@ async def tts_speech(payload: SpeechRequest):
     )
 
 
+def generate_transcription_stream(stt_model, tmp_path: str, gen_kwargs: dict):
+    """Generator that yields transcription chunks and cleans up temp file."""
+    try:
+        # Call generate with stream=True (models handle streaming internally)
+        result = stt_model.generate(tmp_path, **gen_kwargs)
+
+        # Check if result is a generator (streaming mode)
+        if hasattr(result, "__iter__") and hasattr(result, "__next__"):
+            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}
+                else:
+                    # Structured chunk (e.g., Whisper streaming)
+                    chunk_data = {
+                        "text": chunk.text,
+                        "start": getattr(chunk, "start_time", None),
+                        "end": getattr(chunk, "end_time", None),
+                        "is_final": getattr(chunk, "is_final", None),
+                        "language": getattr(chunk, "language", None),
+                    }
+                yield json.dumps(sanitize_for_json(chunk_data)) + "\n"
+        else:
+            # Not a generator, yield the full result
+            yield json.dumps(sanitize_for_json(result)) + "\n"
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+
+
 @app.post("/v1/audio/transcriptions")
 async def stt_transcriptions(
     file: UploadFile = File(...),
     model: str = Form(...),
     language: Optional[str] = Form(None),
+    verbose: bool = Form(False),
+    max_tokens: int = Form(128),
+    chunk_duration: float = Form(30.0),
+    frame_threshold: int = Form(25),
+    stream: bool = Form(False),
+    context: Optional[str] = Form(None),
+    prefill_step_size: int = Form(2048),
+    text: Optional[str] = Form(None),
 ):
     """Transcribe audio using an STT model in OpenAI format."""
+    # Create TranscriptionRequest from form fields
+    payload = TranscriptionRequest(
+        model=model,
+        language=language,
+        verbose=verbose,
+        max_tokens=max_tokens,
+        chunk_duration=chunk_duration,
+        frame_threshold=frame_threshold,
+        stream=stream,
+        context=context,
+        prefill_step_size=prefill_step_size,
+        text=text,
+    )
+
     data = await file.read()
     tmp = io.BytesIO(data)
     audio, sr = audio_read(tmp, always_2d=False)
     tmp.close()
     tmp_path = f"/tmp/{time.time()}.mp3"
     audio_write(tmp_path, audio, sr)
 
-    stt_model = model_provider.load_model(model)
-    result = stt_model.generate(tmp_path)
-    os.remove(tmp_path)
-    # Sanitize NaN values for JSON serialization
-    return sanitize_for_json(result)
+    stt_model = model_provider.load_model(payload.model)
+
+    # Build kwargs for generate, filtering None values
+    gen_kwargs = payload.model_dump(exclude={"model"}, exclude_none=True)
+
+    # Filter kwargs to only include parameters the model's generate method accepts
+    signature = inspect.signature(stt_model.generate)
+    gen_kwargs = {k: v for k, v in gen_kwargs.items() if k in signature.parameters}
+
+    return StreamingResponse(
+        generate_transcription_stream(stt_model, tmp_path, gen_kwargs),
+        media_type="application/x-ndjson",
+    )
 
 
 @app.websocket("/v1/audio/transcriptions/realtime")

mlx_audio/stt/generate.py

@@ -83,7 +83,7 @@ def parse_args():
         "--gen-kwargs",
         type=json.loads,
         default=None,
-        help='Additional generate kwargs as JSON (e.g. \'{"max_chunk_sec": 600, "min_chunk_sec": 1.0}\')',
+        help="Additional generate kwargs as JSON (e.g. '{\"min_chunk_duration\": 1.0}')",
     )
     parser.add_argument(
         "--text",