mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-02 16:45:15 -07:00
fix(backend): close model load races
Kokoro and LuxTTS load_model had no lock, so two concurrent requests could both observe an unloaded model and double-load; they now use the same double-checked asyncio.Lock pattern as the Chatterbox backends. get_stt_backend gets the threading.Lock treatment the TTS and LLM factories already had. Includes the ruff-era typing cleanup for backends/__init__.
This commit is contained in:
@@ -17,15 +17,12 @@ Languages supported (via misaki G2P):
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from . import TTSBackend
|
||||
from .base import (
|
||||
get_torch_device,
|
||||
combine_voice_prompts as _combine_voice_prompts,
|
||||
get_torch_device,
|
||||
model_load_progress,
|
||||
)
|
||||
|
||||
@@ -122,8 +119,9 @@ class KokoroTTSBackend:
|
||||
def __init__(self):
|
||||
self._model = None
|
||||
self._pipelines: dict = {} # lang_code -> KPipeline
|
||||
self._device: Optional[str] = None
|
||||
self._device: str | None = None
|
||||
self.model_size = "default"
|
||||
self._model_load_lock = asyncio.Lock()
|
||||
|
||||
def _get_device(self) -> str:
|
||||
"""Select device. Kokoro supports CUDA and CPU. MPS needs fallback env var."""
|
||||
@@ -157,7 +155,10 @@ class KokoroTTSBackend:
|
||||
"""Load the Kokoro model."""
|
||||
if self._model is not None:
|
||||
return
|
||||
await asyncio.to_thread(self._load_model_sync)
|
||||
async with self._model_load_lock:
|
||||
if self._model is not None:
|
||||
return
|
||||
await asyncio.to_thread(self._load_model_sync)
|
||||
|
||||
def _load_model_sync(self):
|
||||
"""Synchronous model loading."""
|
||||
@@ -239,8 +240,8 @@ class KokoroTTSBackend:
|
||||
text: str,
|
||||
voice_prompt: dict,
|
||||
language: str = "en",
|
||||
seed: Optional[int] = None,
|
||||
instruct: Optional[str] = None,
|
||||
seed: int | None = None,
|
||||
instruct: str | None = None,
|
||||
) -> tuple[np.ndarray, int]:
|
||||
"""
|
||||
Generate audio from text using Kokoro.
|
||||
|
||||
Reference in New Issue
Block a user