engineering/uncloseai-speech
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enable GPU acceleration for Kokoro TTS

Browse source 
Kokoro was hardcoded to use CPU, causing very slow generation times
(3+ minutes for long texts). Now Kokoro uses the same device as XTTS
(auto-detected as 'cuda' if available, otherwise 'cpu').

Changes:
- Add device parameter to kokoro_wrapper __init__ (defaults to 'cpu')
- Pass device to KPipeline constructor
- Use args.xtts_device when initializing Kokoro (same as XTTS)
- Add semaphore lock to prevent concurrent Kokoro model loading
- Log which device Kokoro is loading on

Performance improvement: ~60x faster on GPU vs CPU for long texts
This commit is contained in:
Russell Ballestrini 2025-11-10 04:14:46 -05:00
parent 4576afac39
commit ae958d1bb6
1 changed files with 16 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

24
speech.py
View file

@ -277,19 +277,21 @@ class kokoro_wrapper():
Kokoro is a lightweight decoder-only TTS model (82M params)
Output: 24kHz audio
"""
def __init__(self, lang_code='a'):
def __init__(self, lang_code='a', device='cpu'):
self.lang_code = lang_code
self.device = device
self.sample_rate = 24000 # Kokoro outputs 24kHz
logger.info(f"Loading Kokoro TTS pipeline for language '{lang_code}'")
logger.info(f"Loading Kokoro TTS pipeline for language '{lang_code}' on device '{device}'")
try:
from kokoro import KPipeline
import numpy as np
# KPipeline will use default repo_id if not specified
self.pipeline = KPipeline(lang_code=lang_code)
logger.info(f"Successfully loaded Kokoro pipeline for lang={lang_code}")
# Pass device to KPipeline (supports 'cpu' or 'cuda')
self.pipeline = KPipeline(lang_code=lang_code, device=device)
logger.info(f"Successfully loaded Kokoro pipeline for lang={lang_code} on {device}")
except Exception as e:
logger.error(f"Failed to load Kokoro model: {e}")
raise
@ -794,10 +796,16 @@ async def generate_speech(request: GenerateSpeechRequest):
# Load Kokoro pipeline if not already loaded or if language changed
if kokoro_pipeline is None or kokoro_lang != lang_code:
logger.info(f"Loading/switching Kokoro pipeline to language '{lang_code}'")
# Run blocking model initialization in thread pool to avoid blocking event loop
kokoro_pipeline = await asyncio.to_thread(kokoro_wrapper, lang_code=lang_code)
kokoro_lang = lang_code
# Use semaphore to prevent multiple simultaneous model loads
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
device = args.xtts_device if args.xtts_device != 'none' else 'cpu'
kokoro_pipeline = await asyncio.to_thread(kokoro_wrapper, lang_code=lang_code, device=device)
kokoro_lang = lang_code
# Generate audio (also blocking, so run in thread pool)
audio_data = await asyncio.to_thread(kokoro_pipeline.tts, input_text, voice=kokoro_voice, speed=speed)