Fix Kokoro defaulting to CPU in worker processes

The DefaultArgs class had xtts_device hardcoded to 'cpu', which meant
all uvicorn worker processes inherited this default instead of using
auto_torch_device() to detect GPU.

Changes:
- Set DefaultArgs.xtts_device to None initially
- Call auto_torch_device() after class definition to set default
- This ensures workers use GPU if available, not hardcoded CPU
- Fixed log message to show actual device being used (not args value)
- Log moved after device calculation for accuracy

This fixes Kokoro loading on CPU even when GPU is available.
This commit is contained in:
Russell Ballestrini 2025-11-10 04:22:53 -05:00
parent ae958d1bb6
commit da8e960b2d

View file

@ -119,7 +119,7 @@ kokoro_lang = None
# Default args for worker processes (will be overridden in __main__)
class DefaultArgs:
xtts_device = 'cpu'
xtts_device = None # Will be set after torch import
use_deepspeed = False
unload_timer = None
log_level = 'INFO'
@ -130,6 +130,14 @@ class DefaultArgs:
args = DefaultArgs()
# Set default device after torch is available (for worker processes)
# Main process will override this in __main__ with argparse
if args.xtts_device is None:
try:
args.xtts_device = auto_torch_device()
except:
args.xtts_device = 'cpu'
# Voice-to-model lookup cache (loaded at startup)
voice_to_model_cache = {}
@ -800,10 +808,10 @@ async def generate_speech(request: GenerateSpeechRequest):
async with kokoro_load_semaphore:
# Double-check after acquiring lock (another request may have loaded it)
if kokoro_pipeline is None or kokoro_lang != lang_code:
logger.info(f"Loading/switching Kokoro pipeline to language '{lang_code}' on {args.xtts_device}")
# Run blocking model initialization in thread pool to avoid blocking event loop
# Use same device as XTTS for GPU acceleration
# Use GPU if available, fallback to CPU only if explicitly disabled
device = args.xtts_device if args.xtts_device != 'none' else 'cpu'
logger.info(f"Loading/switching Kokoro pipeline to language '{lang_code}' on device '{device}'")
kokoro_pipeline = await asyncio.to_thread(kokoro_wrapper, lang_code=lang_code, device=device)
kokoro_lang = lang_code