mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-02 16:45:15 -07:00
add model selection and expanded language support to /transcribe endpoint
Closes #233
This commit is contained in:
@@ -134,6 +134,7 @@ class STTBackend(Protocol):
|
||||
self,
|
||||
audio_path: str,
|
||||
language: Optional[str] = None,
|
||||
model_size: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Transcribe audio to text.
|
||||
|
||||
@@ -345,18 +345,20 @@ class MLXSTTBackend:
|
||||
self,
|
||||
audio_path: str,
|
||||
language: Optional[str] = None,
|
||||
model_size: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Transcribe audio to text.
|
||||
|
||||
Args:
|
||||
audio_path: Path to audio file
|
||||
language: Optional language hint (en or zh)
|
||||
language: Optional language hint
|
||||
model_size: Optional model size override
|
||||
|
||||
Returns:
|
||||
Transcribed text
|
||||
"""
|
||||
await self.load_model_async(None)
|
||||
await self.load_model_async(model_size)
|
||||
|
||||
def _transcribe_sync():
|
||||
"""Run synchronous transcription in thread pool."""
|
||||
|
||||
@@ -306,18 +306,20 @@ class PyTorchSTTBackend:
|
||||
self,
|
||||
audio_path: str,
|
||||
language: Optional[str] = None,
|
||||
model_size: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Transcribe audio to text.
|
||||
|
||||
Args:
|
||||
audio_path: Path to audio file
|
||||
language: Optional language hint (en or zh)
|
||||
language: Optional language hint
|
||||
model_size: Optional model size override
|
||||
|
||||
Returns:
|
||||
Transcribed text
|
||||
"""
|
||||
await self.load_model_async(None)
|
||||
await self.load_model_async(model_size)
|
||||
|
||||
def _transcribe_sync():
|
||||
"""Run synchronous transcription in thread pool."""
|
||||
|
||||
+2
-1
@@ -149,7 +149,8 @@ class HistoryListResponse(BaseModel):
|
||||
class TranscriptionRequest(BaseModel):
|
||||
"""Request model for audio transcription."""
|
||||
|
||||
language: Optional[str] = Field(None, pattern="^(en|zh)$")
|
||||
language: Optional[str] = Field(None, pattern="^(en|zh|ja|ko|de|fr|ru|pt|es|it)$")
|
||||
model: Optional[str] = Field(None, pattern="^(base|small|medium|large|turbo)$")
|
||||
|
||||
|
||||
class TranscriptionResponse(BaseModel):
|
||||
|
||||
@@ -20,6 +20,7 @@ UPLOAD_CHUNK_SIZE = 1024 * 1024 # 1MB
|
||||
async def transcribe_audio(
|
||||
file: UploadFile = File(...),
|
||||
language: str | None = Form(None),
|
||||
model: str | None = Form(None),
|
||||
):
|
||||
"""Transcribe audio file to text."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
||||
@@ -29,14 +30,23 @@ async def transcribe_audio(
|
||||
|
||||
try:
|
||||
from ..utils.audio import load_audio
|
||||
from ..backends import WHISPER_HF_REPOS
|
||||
|
||||
audio, sr = await asyncio.to_thread(load_audio, tmp_path)
|
||||
duration = len(audio) / sr
|
||||
|
||||
whisper_model = transcribe.get_whisper_model()
|
||||
model_size = whisper_model.model_size
|
||||
model_size = model if model else whisper_model.model_size
|
||||
|
||||
if not whisper_model.is_loaded() and not whisper_model._is_model_cached(model_size):
|
||||
valid_sizes = list(WHISPER_HF_REPOS.keys())
|
||||
if model_size not in valid_sizes:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Invalid model size '{model_size}'. Must be one of: {', '.join(valid_sizes)}",
|
||||
)
|
||||
|
||||
already_loaded = whisper_model.is_loaded() and whisper_model.model_size == model_size
|
||||
if not already_loaded and not whisper_model._is_model_cached(model_size):
|
||||
progress_model_name = f"whisper-{model_size}"
|
||||
task_manager = get_task_manager()
|
||||
|
||||
@@ -59,7 +69,7 @@ async def transcribe_audio(
|
||||
},
|
||||
)
|
||||
|
||||
text = await whisper_model.transcribe(tmp_path, language)
|
||||
text = await whisper_model.transcribe(tmp_path, language, model_size)
|
||||
|
||||
return models.TranscriptionResponse(
|
||||
text=text,
|
||||
|
||||
Reference in New Issue
Block a user