This commit is contained in:
@@ -0,0 +1,358 @@
|
||||
"""
|
||||
ASR Worker for real-time transcription.
|
||||
|
||||
Wraps the existing VibeVoiceASRInference for async/streaming operation.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import queue
|
||||
from typing import AsyncGenerator, Optional, List, Callable
|
||||
from dataclasses import dataclass
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Add parent directory and demo directory to path for importing existing code
|
||||
_parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, _parent_dir)
|
||||
sys.path.insert(0, os.path.join(_parent_dir, "demo"))
|
||||
|
||||
from .models import (
|
||||
TranscriptionResult,
|
||||
TranscriptionSegment,
|
||||
MessageType,
|
||||
SessionConfig,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class InferenceRequest:
|
||||
"""Request for ASR inference."""
|
||||
audio: np.ndarray
|
||||
sample_rate: int
|
||||
context_info: Optional[str]
|
||||
request_time: float
|
||||
segment_start_sec: float
|
||||
segment_end_sec: float
|
||||
|
||||
|
||||
class ASRWorker:
|
||||
"""
|
||||
ASR Worker that wraps VibeVoiceASRInference for real-time use.
|
||||
|
||||
Features:
|
||||
- Async interface for WebSocket integration
|
||||
- Streaming output via TextIteratorStreamer
|
||||
- Request queuing for handling concurrent segments
|
||||
- Graceful model loading and error handling
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str = "microsoft/VibeVoice-ASR",
|
||||
device: str = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
attn_implementation: str = "flash_attention_2",
|
||||
):
|
||||
"""
|
||||
Initialize the ASR worker.
|
||||
|
||||
Args:
|
||||
model_path: Path to VibeVoice ASR model
|
||||
device: Device to run inference on
|
||||
dtype: Model data type
|
||||
attn_implementation: Attention implementation
|
||||
"""
|
||||
self.model_path = model_path
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
self.attn_implementation = attn_implementation
|
||||
|
||||
self._inference = None
|
||||
self._is_loaded = False
|
||||
self._load_lock = threading.Lock()
|
||||
|
||||
# Inference queue for serializing requests
|
||||
self._inference_semaphore = asyncio.Semaphore(1)
|
||||
|
||||
def load_model(self) -> bool:
|
||||
"""
|
||||
Load the ASR model.
|
||||
|
||||
Returns:
|
||||
True if model loaded successfully
|
||||
"""
|
||||
with self._load_lock:
|
||||
if self._is_loaded:
|
||||
return True
|
||||
|
||||
try:
|
||||
# Import here to avoid circular imports and allow lazy loading
|
||||
# In Docker, the file is copied as vibevoice_asr_gradio_demo.py
|
||||
try:
|
||||
from vibevoice_asr_gradio_demo import VibeVoiceASRInference
|
||||
except ImportError:
|
||||
from vibevoice_asr_gradio_demo_patched import VibeVoiceASRInference
|
||||
|
||||
print(f"Loading VibeVoice ASR model from {self.model_path}...")
|
||||
self._inference = VibeVoiceASRInference(
|
||||
model_path=self.model_path,
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
attn_implementation=self.attn_implementation,
|
||||
)
|
||||
self._is_loaded = True
|
||||
print("ASR model loaded successfully")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to load ASR model: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
@property
|
||||
def is_loaded(self) -> bool:
|
||||
"""Check if model is loaded."""
|
||||
return self._is_loaded
|
||||
|
||||
async def transcribe_segment(
|
||||
self,
|
||||
audio: np.ndarray,
|
||||
sample_rate: int = 16000,
|
||||
context_info: Optional[str] = None,
|
||||
segment_start_sec: float = 0.0,
|
||||
segment_end_sec: float = 0.0,
|
||||
config: Optional[SessionConfig] = None,
|
||||
on_partial: Optional[Callable[[TranscriptionResult], None]] = None,
|
||||
) -> TranscriptionResult:
|
||||
"""
|
||||
Transcribe an audio segment asynchronously.
|
||||
|
||||
Args:
|
||||
audio: Audio data as float32 array
|
||||
sample_rate: Audio sample rate
|
||||
context_info: Optional context for transcription
|
||||
segment_start_sec: Start time of segment in session
|
||||
segment_end_sec: End time of segment in session
|
||||
config: Session configuration
|
||||
on_partial: Callback for partial results
|
||||
|
||||
Returns:
|
||||
Final transcription result
|
||||
"""
|
||||
if not self._is_loaded:
|
||||
if not self.load_model():
|
||||
return TranscriptionResult(
|
||||
type=MessageType.ERROR,
|
||||
text="",
|
||||
is_final=True,
|
||||
latency_ms=0,
|
||||
)
|
||||
|
||||
config = config or SessionConfig()
|
||||
request_time = time.time()
|
||||
|
||||
# Serialize inference requests
|
||||
async with self._inference_semaphore:
|
||||
return await self._run_inference(
|
||||
audio=audio,
|
||||
sample_rate=sample_rate,
|
||||
context_info=context_info,
|
||||
segment_start_sec=segment_start_sec,
|
||||
segment_end_sec=segment_end_sec,
|
||||
config=config,
|
||||
request_time=request_time,
|
||||
on_partial=on_partial,
|
||||
)
|
||||
|
||||
async def _run_inference(
|
||||
self,
|
||||
audio: np.ndarray,
|
||||
sample_rate: int,
|
||||
context_info: Optional[str],
|
||||
segment_start_sec: float,
|
||||
segment_end_sec: float,
|
||||
config: SessionConfig,
|
||||
request_time: float,
|
||||
on_partial: Optional[Callable[[TranscriptionResult], None]],
|
||||
) -> TranscriptionResult:
|
||||
"""Run the actual inference in a thread pool."""
|
||||
from transformers import TextIteratorStreamer
|
||||
|
||||
# Create streamer for partial results
|
||||
streamer = None
|
||||
if config.return_partial_results and on_partial:
|
||||
streamer = TextIteratorStreamer(
|
||||
self._inference.processor.tokenizer,
|
||||
skip_prompt=True,
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
|
||||
# Result container for thread
|
||||
result_container = {"result": None, "error": None}
|
||||
|
||||
def run_inference():
|
||||
try:
|
||||
# Save audio to temp file (required by current implementation)
|
||||
import tempfile
|
||||
import soundfile as sf
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
|
||||
temp_path = f.name
|
||||
|
||||
# Write audio
|
||||
audio_int16 = (audio * 32768.0).clip(-32768, 32767).astype(np.int16)
|
||||
sf.write(temp_path, audio_int16, sample_rate, subtype='PCM_16')
|
||||
|
||||
try:
|
||||
result = self._inference.transcribe(
|
||||
audio_path=temp_path,
|
||||
max_new_tokens=config.max_new_tokens,
|
||||
temperature=config.temperature,
|
||||
context_info=context_info,
|
||||
streamer=streamer,
|
||||
)
|
||||
result_container["result"] = result
|
||||
finally:
|
||||
# Clean up temp file
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except:
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
result_container["error"] = str(e)
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
# Start inference in background thread
|
||||
inference_thread = threading.Thread(target=run_inference)
|
||||
inference_thread.start()
|
||||
|
||||
# Stream partial results if enabled
|
||||
partial_text = ""
|
||||
if streamer and on_partial:
|
||||
try:
|
||||
for new_text in streamer:
|
||||
partial_text += new_text
|
||||
partial_result = TranscriptionResult(
|
||||
type=MessageType.PARTIAL_RESULT,
|
||||
text=partial_text,
|
||||
is_final=False,
|
||||
latency_ms=(time.time() - request_time) * 1000,
|
||||
)
|
||||
# Call callback (may be async)
|
||||
if asyncio.iscoroutinefunction(on_partial):
|
||||
await on_partial(partial_result)
|
||||
else:
|
||||
on_partial(partial_result)
|
||||
except Exception as e:
|
||||
print(f"Error during streaming: {e}")
|
||||
|
||||
# Wait for completion
|
||||
inference_thread.join()
|
||||
|
||||
latency_ms = (time.time() - request_time) * 1000
|
||||
|
||||
if result_container["error"]:
|
||||
return TranscriptionResult(
|
||||
type=MessageType.ERROR,
|
||||
text=f"Error: {result_container['error']}",
|
||||
is_final=True,
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
|
||||
result = result_container["result"]
|
||||
|
||||
# Convert segments to our format
|
||||
segments = []
|
||||
for seg in result.get("segments", []):
|
||||
# Adjust timestamps relative to session
|
||||
seg_start = seg.get("start_time", 0)
|
||||
seg_end = seg.get("end_time", 0)
|
||||
|
||||
# If segment has relative timestamps, adjust to absolute
|
||||
if isinstance(seg_start, (int, float)) and isinstance(seg_end, (int, float)):
|
||||
adjusted_start = segment_start_sec + seg_start
|
||||
adjusted_end = segment_start_sec + seg_end
|
||||
else:
|
||||
adjusted_start = segment_start_sec
|
||||
adjusted_end = segment_end_sec
|
||||
|
||||
segments.append(TranscriptionSegment(
|
||||
start_time=adjusted_start,
|
||||
end_time=adjusted_end,
|
||||
speaker_id=seg.get("speaker_id", "SPEAKER_00"),
|
||||
text=seg.get("text", ""),
|
||||
))
|
||||
|
||||
return TranscriptionResult(
|
||||
type=MessageType.FINAL_RESULT,
|
||||
text=result.get("raw_text", ""),
|
||||
is_final=True,
|
||||
segments=segments,
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
|
||||
async def transcribe_stream(
|
||||
self,
|
||||
audio: np.ndarray,
|
||||
sample_rate: int = 16000,
|
||||
context_info: Optional[str] = None,
|
||||
segment_start_sec: float = 0.0,
|
||||
segment_end_sec: float = 0.0,
|
||||
config: Optional[SessionConfig] = None,
|
||||
) -> AsyncGenerator[TranscriptionResult, None]:
|
||||
"""
|
||||
Transcribe an audio segment with streaming output.
|
||||
|
||||
Yields partial results followed by final result.
|
||||
|
||||
Args:
|
||||
audio: Audio data
|
||||
sample_rate: Sample rate
|
||||
context_info: Optional context
|
||||
segment_start_sec: Segment start time
|
||||
segment_end_sec: Segment end time
|
||||
config: Session config
|
||||
|
||||
Yields:
|
||||
TranscriptionResult objects (partial and final)
|
||||
"""
|
||||
result_queue: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
async def on_partial(result: TranscriptionResult):
|
||||
await result_queue.put(result)
|
||||
|
||||
# Start transcription task
|
||||
transcribe_task = asyncio.create_task(
|
||||
self.transcribe_segment(
|
||||
audio=audio,
|
||||
sample_rate=sample_rate,
|
||||
context_info=context_info,
|
||||
segment_start_sec=segment_start_sec,
|
||||
segment_end_sec=segment_end_sec,
|
||||
config=config,
|
||||
on_partial=on_partial,
|
||||
)
|
||||
)
|
||||
|
||||
# Yield partial results as they come
|
||||
while not transcribe_task.done():
|
||||
try:
|
||||
result = await asyncio.wait_for(result_queue.get(), timeout=0.1)
|
||||
yield result
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
|
||||
# Drain any remaining partial results
|
||||
while not result_queue.empty():
|
||||
yield await result_queue.get()
|
||||
|
||||
# Yield final result
|
||||
final_result = await transcribe_task
|
||||
yield final_result
|
||||
Reference in New Issue
Block a user