Add VibeVoice-ASR backend
Third ASR backend option alongside faster-whisper and qwen3. Builds microsoft/VibeVoice from source (pinned to a specific commit, since it's custom modeling code not in transformers' Auto* registry) into its own venv, inheriting the base image's torch/CUDA. Comes with built-in speaker diarization (VibeVoiceASRProcessor.post_process_transcription returns per-segment speaker ids directly, no separate pyannote pass needed). Verified end-to-end against a real recording: 200 OK, correct Korean transcription, speaker labels populated. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -19,6 +19,7 @@ DEFAULT_LANGUAGE = env_str("DEFAULT_LANGUAGE", "ko")
|
||||
|
||||
FASTER_WHISPER_URL = env_str("FASTER_WHISPER_URL", "http://127.0.0.1:8001")
|
||||
QWEN3_URL = env_str("QWEN3_URL", "http://127.0.0.1:8004")
|
||||
VIBEVOICE_URL = env_str("VIBEVOICE_URL", "http://127.0.0.1:8006")
|
||||
|
||||
PYANNOTE_HF_TOKEN = env_str("PYANNOTE_HF_TOKEN", "")
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ from asr.config import (
|
||||
DEFAULT_MODEL,
|
||||
FASTER_WHISPER_URL,
|
||||
QWEN3_URL,
|
||||
VIBEVOICE_URL,
|
||||
RESULT_DIR,
|
||||
UPLOAD_DIR,
|
||||
ensure_runtime_dirs,
|
||||
@@ -29,6 +30,7 @@ ws_router = APIRouter()
|
||||
BACKENDS = {
|
||||
"faster-whisper": FASTER_WHISPER_URL,
|
||||
"qwen3": QWEN3_URL,
|
||||
"vibevoice": VIBEVOICE_URL,
|
||||
}
|
||||
|
||||
REALTIME_SAMPLE_RATE = 16000
|
||||
@@ -50,6 +52,9 @@ def config() -> Dict[str, Any]:
|
||||
"qwen3": {
|
||||
"models": ["Qwen/Qwen3-ASR-0.6B-hf", "Qwen/Qwen3-ASR-1.7B-hf"],
|
||||
},
|
||||
"vibevoice": {
|
||||
"models": ["microsoft/VibeVoice-ASR"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -168,6 +173,10 @@ async def transcribe(
|
||||
"language": effective_language,
|
||||
"task": task,
|
||||
}
|
||||
elif backend == "vibevoice":
|
||||
data = {
|
||||
"model": model,
|
||||
}
|
||||
else:
|
||||
# faster-whisper
|
||||
data = {
|
||||
|
||||
164
app/asr/workers/vibevoice_worker.py
Normal file
164
app/asr/workers/vibevoice_worker.py
Normal file
@@ -0,0 +1,164 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, "/app")
|
||||
from asr.config import DEVICE, MODEL_CACHE, ensure_runtime_dirs
|
||||
|
||||
import os
|
||||
# 컨테이너는 983 유저로 실행되는데 HOME(/app)이 root 소유라 numba/matplotlib/torch가
|
||||
# 각자 ~/.cache, ~/.config 아래에 쓰려다 실패한다. HOME을 쓰기 가능한 곳으로 돌린다.
|
||||
os.environ["HOME"] = "/tmp" # 컨테이너 기본 HOME=/app은 983 유저가 쓰기 불가 (setdefault로는 덮어쓰기 안 됨)
|
||||
os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache")
|
||||
|
||||
app = FastAPI(title="ASR VibeVoice Worker")
|
||||
|
||||
_MODEL_CACHE: Dict[str, Any] = {}
|
||||
|
||||
|
||||
def _log(msg: str) -> None:
|
||||
print(f"[vibevoice] {msg}", flush=True)
|
||||
|
||||
|
||||
def _load_model(model_id: str) -> Any:
|
||||
if model_id in _MODEL_CACHE:
|
||||
return _MODEL_CACHE[model_id]
|
||||
|
||||
import torch
|
||||
from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
|
||||
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
|
||||
|
||||
_log(f"loading model {model_id}")
|
||||
processor = VibeVoiceASRProcessor.from_pretrained(
|
||||
model_id,
|
||||
language_model_pretrained_name="Qwen/Qwen2.5-1.5B",
|
||||
cache_dir=str(MODEL_CACHE),
|
||||
)
|
||||
model = VibeVoiceASRForConditionalGeneration.from_pretrained(
|
||||
model_id,
|
||||
dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
|
||||
attn_implementation="sdpa",
|
||||
trust_remote_code=True,
|
||||
cache_dir=str(MODEL_CACHE),
|
||||
)
|
||||
model.to(DEVICE)
|
||||
model.eval()
|
||||
|
||||
_MODEL_CACHE[model_id] = (model, processor)
|
||||
_log(f"model {model_id} loaded")
|
||||
return _MODEL_CACHE[model_id]
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup() -> None:
|
||||
ensure_runtime_dirs()
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> Dict[str, Any]:
|
||||
return {"status": "ok", "device": DEVICE, "loaded_models": list(_MODEL_CACHE.keys())}
|
||||
|
||||
|
||||
@app.post("/transcribe")
|
||||
async def transcribe(
|
||||
file: UploadFile = File(...),
|
||||
model: str = Form("microsoft/VibeVoice-ASR"),
|
||||
max_new_tokens: int = Form(4096),
|
||||
) -> JSONResponse:
|
||||
ensure_runtime_dirs()
|
||||
suffix = Path(file.filename or "audio.bin").suffix or ".wav"
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(await file.read())
|
||||
tmp_path = Path(tmp.name)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
vv_model, processor = _load_model(model)
|
||||
|
||||
inputs = processor(
|
||||
audio=[str(tmp_path)],
|
||||
sampling_rate=None,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
inputs = {k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
|
||||
|
||||
gen_config = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"pad_token_id": processor.pad_id,
|
||||
"eos_token_id": processor.tokenizer.eos_token_id,
|
||||
"do_sample": False,
|
||||
}
|
||||
|
||||
_log(f"transcribing model={model}")
|
||||
with torch.no_grad():
|
||||
output_ids = vv_model.generate(**inputs, **gen_config)
|
||||
|
||||
input_length = inputs["input_ids"].shape[1]
|
||||
generated_ids = output_ids[0, input_length:]
|
||||
raw_text = processor.decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
try:
|
||||
raw_segments = processor.post_process_transcription(raw_text)
|
||||
except Exception as e:
|
||||
_log(f"post_process_transcription failed: {e}")
|
||||
raw_segments = []
|
||||
|
||||
segments: List[Dict[str, Any]] = []
|
||||
full_text_parts: List[str] = []
|
||||
for i, seg in enumerate(raw_segments):
|
||||
text = (seg.get("text") or "").strip()
|
||||
speaker_id = seg.get("speaker_id")
|
||||
start = seg.get("start_time")
|
||||
end = seg.get("end_time")
|
||||
segments.append({
|
||||
"id": i,
|
||||
"start": round(float(start), 3) if start is not None else None,
|
||||
"end": round(float(end), 3) if end is not None else None,
|
||||
"text": text,
|
||||
"speaker": f"SPEAKER_{speaker_id:02d}" if speaker_id is not None else "UNKNOWN",
|
||||
})
|
||||
if text:
|
||||
full_text_parts.append(text)
|
||||
|
||||
duration = segments[-1]["end"] if segments and segments[-1]["end"] is not None else None
|
||||
full_text = " ".join(full_text_parts)
|
||||
|
||||
return JSONResponse({
|
||||
"backend": "vibevoice",
|
||||
"model": model,
|
||||
"language": None,
|
||||
"duration": duration,
|
||||
"text": full_text,
|
||||
"segments": segments,
|
||||
"diarized": True,
|
||||
})
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
_log(f"error: {type(e).__name__}: {e}")
|
||||
raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
|
||||
finally:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8006)
|
||||
args = parser.parse_args()
|
||||
_log(f"starting host={args.host} port={args.port} device={DEVICE}")
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
@@ -26,6 +26,17 @@ stderr_logfile=/dev/fd/2
|
||||
stderr_logfile_maxbytes=0
|
||||
priority=20
|
||||
|
||||
[program:vibevoice]
|
||||
command=/opt/venvs/vibevoice/bin/python /app/asr/workers/vibevoice_worker.py --host 0.0.0.0 --port 8006
|
||||
directory=/app
|
||||
autostart=true
|
||||
autorestart=true
|
||||
stdout_logfile=/dev/fd/1
|
||||
stdout_logfile_maxbytes=0
|
||||
stderr_logfile=/dev/fd/2
|
||||
stderr_logfile_maxbytes=0
|
||||
priority=25
|
||||
|
||||
[program:xtts]
|
||||
command=/opt/venvs/xtts/bin/python /app/tts/workers/xtts_worker.py --host 0.0.0.0 --port 8005
|
||||
directory=/app
|
||||
|
||||
@@ -14,6 +14,7 @@ async function loadConfig() {
|
||||
backends: {
|
||||
'faster-whisper': { models: ['tiny','base','small','medium','large-v3','large-v2','turbo'] },
|
||||
'qwen3': { models: ['Qwen/Qwen3-ASR-0.6B-hf','Qwen/Qwen3-ASR-1.7B-hf'] },
|
||||
'vibevoice': { models: ['microsoft/VibeVoice-ASR'] },
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -52,8 +53,8 @@ function populateModels(backend) {
|
||||
}
|
||||
|
||||
function onBackendChange() {
|
||||
const isQwen = backendSel.value === 'qwen3';
|
||||
fwOptions.style.display = isQwen ? 'none' : '';
|
||||
const isFasterWhisper = backendSel.value === 'faster-whisper';
|
||||
fwOptions.style.display = isFasterWhisper ? '' : 'none';
|
||||
populateModels(backendSel.value);
|
||||
}
|
||||
|
||||
|
||||
@@ -47,6 +47,7 @@
|
||||
<select id="backend" name="backend">
|
||||
<option value="faster-whisper" selected>faster-whisper</option>
|
||||
<option value="qwen3">Qwen3-ASR</option>
|
||||
<option value="vibevoice">VibeVoice-ASR</option>
|
||||
</select>
|
||||
</label>
|
||||
<label>모델
|
||||
|
||||
Reference in New Issue
Block a user