stream_pipeline.py

from datasets import streaming
import gradio as gr
from httpx import stream
import numpy as np
from constants import STREAMER_CLASSES
from stt_streamers import DeepgramStreamer
from sdk import SDKWrapper, SDKParams
from typing import Any
from ui import VAD_OFF_HTML, VAD_ON_HTML
# ----------------------------
# Global transcript store (UI pulls from this)
# ----------------------------
_ENHANCED_TRANSCRIPT: str = ""
_RAW_TRANSCRIPT: str = ""

def _set_transcript_enhanced(text: str) -> None:
    """Deepgram callback: update latest transcript text (no printing)."""
    global _ENHANCED_TRANSCRIPT
    _ENHANCED_TRANSCRIPT = text

def _set_transcript_raw(text: str) -> None:
    """Deepgram callback: update latest transcript text (no printing)."""
    global _RAW_TRANSCRIPT
    _RAW_TRANSCRIPT = text


def get_live_transcripts() -> tuple[str, str]:
    """Return current enhanced and raw transcript for live UI updates."""
    return _ENHANCED_TRANSCRIPT, _RAW_TRANSCRIPT


SDK_STREAMING = SDKWrapper()


# Created on first start_recording (lazy) to avoid Soniox "No audio received" timeout at app load
Streamer_enhanced = None
Streamer_raw = None
_streamer_generation = 0
_last_stop_generation = 1  # so first stop doesn't skip (1 > 1 is False)


def init_streaming_sdk(sample_rate: int, enhancement_level: float):
    """Initialize SDK processor and STT streamers. Call on first start_recording to avoid Soniox timeout at app load."""
    sdk_params = SDKParams(
        sample_rate=sample_rate,
        enhancement_level=enhancement_level,
        allow_variable_frames=False,
        num_channels=1,
        sync=True, 
    )
    SDK_STREAMING.init_processor(sdk_params)

def _to_float32_mono(y: np.ndarray) -> np.ndarray:
    # Gradio liefert int16 (oder (samples, channels)). citeturn1view4
    y = np.asarray(y)
    if y.ndim > 1:
        y = y.mean(axis=1)
    if y.dtype == np.int16:
        y = (y.astype(np.float32) / 32768.0)
    else:
        y = y.astype(np.float32)
    return y


def transcribe_stream(current_sr: int | None, new_chunk, enhancement_level, input_gain_db: float = 0.0, stt_streamer: str = "deepgram") -> tuple[int | None, str, str, Any]:
    if new_chunk is None or new_chunk[1] is None:
        return None, _ENHANCED_TRANSCRIPT, _RAW_TRANSCRIPT, VAD_OFF_HTML 
    sr, y = new_chunk
    if current_sr != sr:
        init_streaming_sdk(sample_rate=sr, enhancement_level=enhancement_level/100.0)
        set_stt_streamer(stt_streamer, sr)
        current_sr = sr
    if SDK_STREAMING.enhancement_level != enhancement_level:
        SDK_STREAMING.change_enhancement_level(enhancement_level/100.0)
    y = _to_float32_mono(y)
    if input_gain_db is not None and input_gain_db > 0:
        gain_linear = np.float32(10.0 ** (float(input_gain_db) / 20.0))
        y = (y * gain_linear).astype(np.float32)
        y = np.clip(y, -1.0, 1.0)
    y = np.asarray(y, dtype=np.float32).flatten()
    enhanced_chunk_16k, vad_detected = SDK_STREAMING.process_with_vad(y)
    out_1d = np.asarray(enhanced_chunk_16k, dtype=np.float32).flatten()
    if (
        Streamer_enhanced is not None
        and Streamer_raw is not None
    ):
        Streamer_raw.process_chunk(y)
        Streamer_enhanced.process_chunk(out_1d)
    vad_led = VAD_ON_HTML if vad_detected else VAD_OFF_HTML
    return current_sr, _ENHANCED_TRANSCRIPT, _RAW_TRANSCRIPT, vad_led


def shutdown_streamers(from_stop_recording: bool = False):
    """Shut down STT streamers. If from_stop_recording, skip when streamers were
    created after the last stop (avoids delayed stop killing new streamers)."""
    global Streamer_enhanced, Streamer_raw, _streamer_generation, _last_stop_generation
    if from_stop_recording and _streamer_generation > _last_stop_generation:
        return
    gen = _streamer_generation
    try:
        if Streamer_enhanced is not None and Streamer_enhanced.ws is not None:
            Streamer_enhanced.shutdown()
        if Streamer_raw is not None and Streamer_raw.ws is not None:
            Streamer_raw.shutdown()
    except Exception:
        print("Failed to shutdown streamers.")
    finally:
        Streamer_enhanced = None
        Streamer_raw = None
        if from_stop_recording:
            _last_stop_generation = gen
    
def on_stop_recording():
    shutdown_streamers()
    return VAD_OFF_HTML, None


def clear_ui():
    global _ENHANCED_TRANSCRIPT, _RAW_TRANSCRIPT
    _ENHANCED_TRANSCRIPT = ""
    _RAW_TRANSCRIPT = ""
    return _ENHANCED_TRANSCRIPT, _RAW_TRANSCRIPT


def set_stt_streamer(model_name, fs_hz):
    StreamerCls = STREAMER_CLASSES.get(model_name, DeepgramStreamer)
    global Streamer_enhanced, Streamer_raw, _streamer_generation
    # Shut down current streamers first so we don't leak
    if Streamer_enhanced is not None or Streamer_raw is not None:
        shutdown_streamers()
    # Create both before assigning so transcribe_stream never sees one new and one old
    new_enhanced = StreamerCls(
        fs_hz=fs_hz,
        stream_name="enhanced",
        on_update=_set_transcript_enhanced,
    )
    new_raw = StreamerCls(
        fs_hz=fs_hz,
        stream_name="raw",
        on_update=_set_transcript_raw,
    )
    _streamer_generation += 1
    Streamer_enhanced = new_enhanced
    Streamer_raw = new_raw