Qwen-ASR-VLM / app.py
aspirant312's picture
Update app.py
cecf4c0 verified
Raw
History Blame Contribute Delete
27.4 kB
"""
Qwen3-ASR-1.7B vLLM Streaming Server - HuggingFace Space
DashScope-compatible WebSocket protocol with server-side VAD.
Your existing QwenCloudASRSTTService pipecat client works by just changing the URL.
Endpoints:
GET /health - Health check
POST /v1/audio/transcriptions - Batch file transcription
WS /v1/realtime - Streaming ASR (DashScope protocol)
"""
import os
try:
from dotenv import load_dotenv
load_dotenv()
except ImportError:
pass
import sys
import json
import time
import base64
import asyncio
import threading
import logging
import tempfile
import uuid
import copy
import concurrent.futures
import numpy as np
import soundfile as sf
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, HTTPException, Form
from fastapi.responses import JSONResponse
import uvicorn
# ============================================================
# Logging
# ============================================================
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("qwen3-asr-vllm")
# ============================================================
# Configuration
# ============================================================
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-1.7B")
GPU_MEMORY_UTILIZATION = float(os.getenv("GPU_MEMORY_UTILIZATION", "0.80"))
STREAMING_MAX_NEW_TOKENS = int(os.getenv("STREAMING_MAX_NEW_TOKENS", "1024"))
# Caps vLLM's context window. Lower = smaller KV cache = much less VRAM.
# ASR rarely needs >2-4k. Set "" to let vLLM auto-pick (will be huge).
MAX_MODEL_LEN = int(os.getenv("MAX_MODEL_LEN", "0")) or None
# Skip CUDA graph capture. Saves ~1-2 GB VRAM, slightly slower per-step.
ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "false").lower() in ("1", "true", "yes")
# Server port (Dockerfile EXPOSE keeps 7860; override locally if taken).
PORT = int(os.getenv("PORT", "7860"))
# Streaming chunk config — smaller = faster partials, larger = more context
CHUNK_SIZE_SEC = float(os.getenv("CHUNK_SIZE_SEC", "4.0"))
UNFIXED_CHUNK_NUM = int(os.getenv("UNFIXED_CHUNK_NUM", "5"))
UNFIXED_TOKEN_NUM = int(os.getenv("UNFIXED_TOKEN_NUM", "15"))
SAMPLE_RATE = 16000
LANGUAGE = os.getenv("LANGUAGE", "English")
# # Hardcoded system prompt biasing the model toward Chughtai Lab patient queries
# # (test names, sample collection, reports). Fed to init_streaming_state / transcribe.
# CHUGHTAI_CONTEXT = (
# "Phone calls to Chughtai Lab, a Pakistani diagnostic laboratory. Callers "
# "speak Urdu/Hindi mixed with English medical and booking terms.\n"
# "Tests: CBC, LFT, RFT, KFT, HbA1c, FBS, RBS, BSR, Lipid Profile, "
# "Vitamin D, Vitamin B12, TSH, T3, T4, Free T3, Free T4, Thyroid Profile, "
# "Iron, TIBC, Ferritin, Creatinine, Urea, Uric Acid, Electrolytes, "
# "ESR, CRP, D-Dimer, PT, APTT, INR, Urine Complete Examination, Urine R/E, "
# "Stool R/E, Blood Group, Cross Match, Beta HCG, PSA, Amylase, Lipase, "
# "SGPT, ALT, SGOT, AST, Bilirubin, Albumin, Calcium, Magnesium, Phosphorus.\n"
# "Diseases and panels: COVID, PCR, Dengue, Dengue NS1, Typhoid, Widal, "
# "Malaria, MP, Hepatitis B, Hepatitis C, HBsAg, Anti HCV, HIV, "
# "H Pylori, Brucella.\n"
# "Imaging: ECG, EKG, X-Ray, Ultrasound, Echo, MRI, CT Scan, Mammogram.\n"
# "Specialists: cardiologist, general physician, pediatrician, gynecologist, "
# "dermatologist, neurologist, endocrinologist, urologist, nephrologist.\n"
# "Service terms: appointment, booking, available, cancel, reschedule, "
# "confirm, timing, slot, address, sample, sampling, home sampling, "
# "home collection, phlebotomist, report, result, test, profile, panel, "
# "fasting, non-fasting, price, rate, discount, branch, collection point, "
# "WhatsApp, online, portal."
# )
# VAD defaults (can be overridden per-session via session.update)
VAD_THRESHOLD = float(os.getenv("VAD_THRESHOLD", "0.7"))
VAD_MIN_SILENCE_MS = int(os.getenv("VAD_MIN_SILENCE_MS", "800"))
VAD_SPEECH_PAD_MS = int(os.getenv("VAD_SPEECH_PAD_MS", "300"))
# Hallucination filtering
HALLUCINATION_PHRASES = {
"transcript", "transcription", "thank you", "thanks for watching",
"you", "bye", "goodbye", "the end", "subtitle", "subtitles",
}
log.info("=" * 60)
log.info("Qwen3-ASR vLLM Streaming Server Config:")
log.info(f" MODEL_ID = {MODEL_ID}")
log.info(f" GPU_MEMORY_UTILIZATION = {GPU_MEMORY_UTILIZATION}")
log.info(f" STREAMING_MAX_TOKENS = {STREAMING_MAX_NEW_TOKENS}")
log.info(f" CHUNK_SIZE_SEC = {CHUNK_SIZE_SEC}s")
log.info(f" UNFIXED_CHUNK_NUM = {UNFIXED_CHUNK_NUM} (revise last {CHUNK_SIZE_SEC * UNFIXED_CHUNK_NUM:.0f}s)")
log.info(f" UNFIXED_TOKEN_NUM = {UNFIXED_TOKEN_NUM}")
log.info(f" LANGUAGE = {LANGUAGE}")
log.info(f" VAD_THRESHOLD = {VAD_THRESHOLD}")
log.info(f" VAD_MIN_SILENCE_MS = {VAD_MIN_SILENCE_MS}")
log.info(f" MAX_MODEL_LEN = {MAX_MODEL_LEN or 'auto'}")
log.info(f" ENFORCE_EAGER = {ENFORCE_EAGER}")
log.info(f" PORT = {PORT}")
log.info("=" * 60)
# Thread pool for running synchronous model inference off the event loop
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
# ============================================================
# FastAPI App
# ============================================================
app = FastAPI(title="Qwen3-ASR vLLM Streaming", version="2.0.0")
# ============================================================
# ASR Model Loading (singleton, thread-safe)
# ============================================================
_asr_model = None
_asr_lock = threading.Lock()
_model_ready = threading.Event()
def _get_dtype():
try:
import torch
if torch.cuda.is_available():
cap = torch.cuda.get_device_capability()
if cap[0] * 10 + cap[1] >= 80:
return "bfloat16"
except Exception:
pass
return "half"
def get_asr_model():
global _asr_model
if _asr_model is None:
with _asr_lock:
if _asr_model is None:
log.info(f"Loading {MODEL_ID}...")
start = time.time()
dtype = _get_dtype()
from qwen_asr import Qwen3ASRModel
llm_kwargs = {
"model": MODEL_ID,
"gpu_memory_utilization": GPU_MEMORY_UTILIZATION,
"dtype": dtype,
"max_new_tokens": STREAMING_MAX_NEW_TOKENS,
"enforce_eager": ENFORCE_EAGER,
}
if MAX_MODEL_LEN:
llm_kwargs["max_model_len"] = MAX_MODEL_LEN
_asr_model = Qwen3ASRModel.LLM(**llm_kwargs)
log.info(f"Model loaded in {time.time() - start:.1f}s (dtype={dtype})")
_model_ready.set()
return _asr_model
# ============================================================
# Silero VAD (per-connection instances via deep copy)
# ============================================================
_vad_model_template = None
_vad_utils = None
_vad_lock = threading.Lock()
def _ensure_vad_loaded():
"""Load the VAD model template once (thread-safe)."""
global _vad_model_template, _vad_utils
if _vad_model_template is None:
with _vad_lock:
if _vad_model_template is None:
import torch
_vad_model_template, _vad_utils = torch.hub.load(
"snakers4/silero-vad", "silero_vad",
trust_repo=True, verbose=False,
)
log.info("Silero VAD model loaded")
def create_vad(threshold=VAD_THRESHOLD, min_silence_ms=VAD_MIN_SILENCE_MS,
speech_pad_ms=VAD_SPEECH_PAD_MS):
"""Create a new VAD iterator for a WebSocket connection."""
try:
_ensure_vad_loaded()
model_copy = copy.deepcopy(_vad_model_template)
VADIterator = _vad_utils[3]
return VADIterator(
model_copy,
threshold=threshold,
sampling_rate=SAMPLE_RATE,
min_silence_duration_ms=min_silence_ms,
speech_pad_ms=speech_pad_ms,
)
except Exception as e:
log.warning(f"Silero VAD failed ({e}), using RMS fallback")
return RMSVad(threshold=0.01, silence_frames=int(min_silence_ms / 20))
class RMSVad:
"""Simple energy-based VAD fallback."""
def __init__(self, threshold=0.01, silence_frames=30, speech_frames=3):
self.threshold = threshold
self.silence_frames = silence_frames
self.speech_frames = speech_frames
self.is_speaking = False
self._silent_count = 0
self._speech_count = 0
def __call__(self, audio_chunk):
import torch
if isinstance(audio_chunk, np.ndarray):
audio_chunk = torch.from_numpy(audio_chunk)
rms = float(torch.sqrt(torch.mean(audio_chunk.float() ** 2)))
if rms > self.threshold:
self._speech_count += 1
self._silent_count = 0
if not self.is_speaking and self._speech_count >= self.speech_frames:
self.is_speaking = True
return {"start": 0}
else:
self._silent_count += 1
self._speech_count = 0
if self.is_speaking and self._silent_count >= self.silence_frames:
self.is_speaking = False
return {"end": 0}
return None
def reset_states(self):
self.is_speaking = False
self._silent_count = 0
self._speech_count = 0
# ============================================================
# Streaming Session (per-utterance ASR state)
# ============================================================
def _is_hallucination(text: str) -> bool:
return text.lower().strip().rstrip(".!?,;:") in HALLUCINATION_PHRASES
class UtteranceSession:
"""Manages vLLM streaming state for a single utterance."""
def __init__(self, model, language=None):
self.model = model
self.state = model.init_streaming_state(
context="",
language=None,
unfixed_chunk_num=UNFIXED_CHUNK_NUM,
unfixed_token_num=UNFIXED_TOKEN_NUM,
chunk_size_sec=CHUNK_SIZE_SEC,
)
self.is_finalized = False
self.last_text = ""
self.total_audio_sec = 0.0
def feed(self, audio: np.ndarray) -> dict:
"""Feed audio chunk, return {text, delta, is_final}."""
if self.is_finalized:
return {"text": self.last_text, "delta": "", "is_final": True}
audio = np.asarray(audio, dtype=np.float32)
self.total_audio_sec += len(audio) / SAMPLE_RATE
self.model.streaming_transcribe(audio, self.state)
current = self.state.text or ""
if current and _is_hallucination(current):
current = self.last_text
# Compute delta (new text since last update)
delta = ""
if current and current != self.last_text:
if current.startswith(self.last_text):
delta = current[len(self.last_text):]
else:
delta = current # Full revision occurred
self.last_text = current
return {"text": current, "delta": delta, "is_final": False}
def finalize(self) -> dict:
"""Finalize utterance, return final transcript."""
if self.is_finalized:
return {"text": self.last_text, "is_final": True}
fallback = self.state.text or ""
self.model.finish_streaming_transcribe(self.state)
text = self.state.text or fallback
if text and _is_hallucination(text):
text = ""
self.is_finalized = True
self.last_text = text
log.info(f"Utterance finalized: {self.total_audio_sec:.1f}s audio -> '{text[:100]}'")
return {"text": text, "is_final": True}
# ============================================================
# Realtime Connection Manager
# ============================================================
class RealtimeConnection:
"""
Per-WebSocket-connection state.
Manages VAD + streaming ASR sessions across multiple utterances.
"""
# Silero VAD requires at least 512 samples (32ms @ 16kHz)
MIN_VAD_SAMPLES = 512
def __init__(self, asr_model, vad, language=None):
self.asr_model = asr_model
self.vad = vad
self.language = language # Set by client via session.update
self.utterance = None # Current UtteranceSession
self.audio_buffer = np.array([], dtype=np.float32)
self.vad_buffer = np.array([], dtype=np.float32)
self.is_speaking = False
self.chunk_samples = int(CHUNK_SIZE_SEC * SAMPLE_RATE)
def process_audio(self, audio_f32: np.ndarray) -> list[dict]:
"""
Process audio chunk through VAD + ASR.
Returns list of DashScope-compatible events to send to client.
"""
import torch
events = []
# Buffer audio for VAD — Silero requires EXACTLY 512 samples at 16kHz
self.vad_buffer = np.concatenate([self.vad_buffer, audio_f32])
# Process in exact 512-sample windows (Silero rejects any other size)
while len(self.vad_buffer) >= self.MIN_VAD_SAMPLES:
vad_frame = self.vad_buffer[:self.MIN_VAD_SAMPLES]
self.vad_buffer = self.vad_buffer[self.MIN_VAD_SAMPLES:]
vad_result = self.vad(torch.from_numpy(vad_frame))
# Speech started
if vad_result and "start" in vad_result and not self.is_speaking:
self.is_speaking = True
self.utterance = UtteranceSession(self.asr_model, language=self.language)
self.audio_buffer = np.array([], dtype=np.float32)
events.append({"type": "input_audio_buffer.speech_started"})
# Accumulate audio during speech
if self.is_speaking and self.utterance:
self.audio_buffer = np.concatenate([self.audio_buffer, vad_frame])
# Feed chunks to ASR model when we have enough
while len(self.audio_buffer) >= self.chunk_samples:
chunk = self.audio_buffer[:self.chunk_samples]
self.audio_buffer = self.audio_buffer[self.chunk_samples:]
result = self.utterance.feed(chunk)
if result["delta"]:
events.append({
"type": "conversation.item.input_audio_transcription.delta",
"delta": result["delta"],
})
if result["text"]:
events.append({
"type": "conversation.item.input_audio_transcription.text",
"text": result["text"],
})
# Speech ended
if vad_result and "end" in vad_result and self.is_speaking:
events.extend(self._finalize_utterance())
return events
def _finalize_utterance(self) -> list[dict]:
"""Finalize current utterance and return events."""
events = []
if not self.utterance:
return events
# Feed remaining buffered audio
if len(self.audio_buffer) > 0:
result = self.utterance.feed(self.audio_buffer)
self.audio_buffer = np.array([], dtype=np.float32)
if result["text"]:
events.append({
"type": "conversation.item.input_audio_transcription.text",
"text": result["text"],
})
# Finalize
result = self.utterance.finalize()
events.append({"type": "input_audio_buffer.speech_stopped"})
events.append({
"type": "conversation.item.input_audio_transcription.completed",
"transcript": result["text"],
})
self.utterance = None
self.is_speaking = False
return events
def commit(self) -> list[dict]:
"""Force finalize current utterance (manual commit)."""
if self.utterance and not self.utterance.is_finalized:
return self._finalize_utterance()
return []
def clear(self) -> list[dict]:
"""Clear audio buffer and discard current utterance."""
self.audio_buffer = np.array([], dtype=np.float32)
if self.utterance and not self.utterance.is_finalized:
self.utterance.is_finalized = True
self.utterance = None
self.is_speaking = False
return [{"type": "input_audio_buffer.cleared"}]
# ============================================================
# HTTP Endpoints
# ============================================================
@app.get("/health")
def health():
return {
"status": "ok",
"model": MODEL_ID,
"model_status": "loaded" if _asr_model else "loading",
"vad": "silero" if _vad_model_template else "rms_fallback",
"streaming_config": {
"chunk_size_sec": CHUNK_SIZE_SEC,
"unfixed_chunk_num": UNFIXED_CHUNK_NUM,
"unfixed_token_num": UNFIXED_TOKEN_NUM,
},
}
@app.get("/")
def root():
return {
"service": "Qwen3-ASR vLLM Streaming Server",
"model": MODEL_ID,
"protocol": "DashScope-compatible (OpenAI realtime=v1)",
"endpoints": {
"/health": "GET - Health check",
"/v1/audio/transcriptions": "POST - Batch file transcription",
"/v1/realtime": "WebSocket - Streaming with server-side VAD",
},
"websocket_protocol": {
"url": "wss://YOUR-SPACE.hf.space/v1/realtime",
"input_format": "PCM int16, 16kHz, mono, base64-encoded",
"events_in": [
"session.update",
"input_audio_buffer.append",
"input_audio_buffer.commit",
"input_audio_buffer.clear",
],
"events_out": [
"session.created",
"session.updated",
"input_audio_buffer.speech_started",
"input_audio_buffer.speech_stopped",
"conversation.item.input_audio_transcription.delta",
"conversation.item.input_audio_transcription.text",
"conversation.item.input_audio_transcription.completed",
"error",
],
},
}
@app.post("/v1/audio/transcriptions")
async def transcribe(file: UploadFile = File(...), language: str = Form(None)):
"""Batch transcription (OpenAI-compatible)."""
if not _model_ready.is_set():
raise HTTPException(503, "Model still loading")
lang = language or LANGUAGE
suffix = os.path.splitext(file.filename or "")[1] or ".wav"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
content = await file.read()
tmp.write(content)
tmp_path = tmp.name
try:
model = get_asr_model()
log.info("Transcribing with language=None")
result = model.transcribe([tmp_path], language=[None], context=[""])
text = result[0].text if result else ""
if text and _is_hallucination(text):
text = ""
return {"text": text, "model": MODEL_ID, "language": lang}
except Exception as e:
log.error(f"Transcription error: {e}")
raise HTTPException(500, str(e))
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
# ============================================================
# WebSocket Streaming (DashScope-compatible protocol)
# ============================================================
@app.websocket("/v1/realtime")
async def realtime_stream(ws: WebSocket):
await ws.accept()
client = ws.client
log.info(f"WebSocket connected: {client}")
# Wait for model
if not _model_ready.is_set():
await ws.send_json({"type": "error", "error": {"message": "Model loading..."}})
loop = asyncio.get_event_loop()
ready = await loop.run_in_executor(None, _model_ready.wait, 300)
if not ready:
await ws.close(1011, "Model load timeout")
return
# Send session.created
session_id = f"sess_{uuid.uuid4().hex[:12]}"
await ws.send_json({
"type": "session.created",
"session": {"id": session_id, "model": MODEL_ID},
})
model = get_asr_model()
# Create per-connection VAD (deep copy of shared template)
vad = create_vad()
# Create connection manager
conn = RealtimeConnection(model, vad)
try:
while True:
message = await ws.receive()
if message["type"] == "websocket.disconnect":
break
# --- JSON messages ---
if "text" in message and message["text"]:
try:
data = json.loads(message["text"])
except json.JSONDecodeError:
await ws.send_json({
"type": "error",
"error": {"message": "Invalid JSON"},
})
continue
event_type = data.get("type", "")
if event_type == "session.update":
# Accept VAD + language config from client
session_cfg = data.get("session", {})
# Language (from input_audio_transcription.language)
iat = session_cfg.get("input_audio_transcription", {})
client_lang = iat.get("language")
if client_lang:
# Map ISO codes to full names for the model
lang_map = {"en": "English", "ms": "Malay", "zh": "Chinese",
"ja": "Japanese", "ko": "Korean", "hi": "Hindi",
"ar": "Arabic", "id": "Indonesian", "th": "Thai"}
conn.language = lang_map.get(client_lang, client_lang)
log.info(f"Language updated: {client_lang} -> {conn.language}")
td = session_cfg.get("turn_detection", {})
if td.get("type") == "server_vad":
threshold = td.get("threshold", VAD_THRESHOLD)
silence_ms = td.get("silence_duration_ms", VAD_MIN_SILENCE_MS)
# Recreate VAD with client-specified params
vad = create_vad(
threshold=threshold,
min_silence_ms=silence_ms,
)
conn.vad = vad
log.info(f"VAD updated: threshold={threshold}, silence={silence_ms}ms")
await ws.send_json({
"type": "session.updated",
"session": {"id": session_id},
})
elif event_type == "input_audio_buffer.append":
audio_b64 = data.get("audio", "")
if not audio_b64:
continue
raw = base64.b64decode(audio_b64)
chunk_int16 = np.frombuffer(raw, dtype=np.int16)
chunk_f32 = chunk_int16.astype(np.float32) / 32768.0
# Run VAD + ASR in thread pool (avoid blocking event loop)
events = await asyncio.get_event_loop().run_in_executor(
_executor, conn.process_audio, chunk_f32,
)
for ev in events:
await ws.send_json(ev)
elif event_type == "input_audio_buffer.commit":
events = await asyncio.get_event_loop().run_in_executor(
_executor, conn.commit,
)
for ev in events:
await ws.send_json(ev)
elif event_type == "input_audio_buffer.clear":
events = conn.clear()
for ev in events:
await ws.send_json(ev)
# --- Binary audio (alternative to base64 JSON) ---
elif "bytes" in message and message["bytes"]:
raw = message["bytes"]
chunk_int16 = np.frombuffer(raw, dtype=np.int16)
chunk_f32 = chunk_int16.astype(np.float32) / 32768.0
events = await asyncio.get_event_loop().run_in_executor(
_executor, conn.process_audio, chunk_f32,
)
for ev in events:
await ws.send_json(ev)
except WebSocketDisconnect:
log.info(f"WebSocket disconnected: {client}")
except Exception as e:
log.error(f"WebSocket error: {e}", exc_info=True)
finally:
# Finalize any in-progress utterance
if conn.utterance and not conn.utterance.is_finalized:
try:
events = await asyncio.get_event_loop().run_in_executor(
_executor, conn.commit,
)
for ev in events:
try:
await ws.send_json(ev)
except Exception:
pass
except Exception:
pass
log.info(f"WebSocket session ended: {client}")
# ============================================================
# Main
# ============================================================
if __name__ == "__main__":
log.info("=" * 60)
log.info("Qwen3-ASR vLLM Streaming Server")
log.info("=" * 60)
if not os.environ.get("OMP_NUM_THREADS"):
os.environ["OMP_NUM_THREADS"] = "4"
# Log system info
try:
import torch
log.info(f"Python: {sys.version.split()[0]}")
log.info(f"PyTorch: {torch.__version__}")
log.info(f"CUDA: {torch.cuda.is_available()}")
if torch.cuda.is_available():
log.info(f"CUDA version: {torch.version.cuda}")
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
mem = getattr(props, "total_memory", 0) or getattr(props, "total_mem", 0)
log.info(f"GPU {i}: {props.name} ({mem / (1024**3):.1f} GB)")
except Exception as e:
log.warning(f"System info: {e}")
# Pre-load ASR model (blocking)
log.info("Loading ASR model...")
try:
get_asr_model()
log.info("ASR model ready")
except Exception as e:
log.error(f"ASR model load failed: {e}", exc_info=True)
# Pre-load VAD model template
log.info("Loading Silero VAD...")
try:
_ensure_vad_loaded()
log.info("Silero VAD ready")
except Exception as e:
log.warning(f"Silero VAD failed, RMS fallback will be used: {e}")
log.info(f"Starting server on 0.0.0.0:{PORT}")
log.info(f" Health: http://0.0.0.0:{PORT}/health")
log.info(f" Batch: http://0.0.0.0:{PORT}/v1/audio/transcriptions")
log.info(f" Streaming: ws://0.0.0.0:{PORT}/v1/realtime")
uvicorn.run(app, host="0.0.0.0", port=PORT)