add --engines CLI flag to allowlist TTS backends
Operators can now restrict speech.py to a subset of TTS engines via
--engines (or the SPEECH_ENABLED_ENGINES env var). Disabled engines:
- hidden from /v1/voices
- short-circuited at the TTS request handler with a clean
BadRequestError ('Model X is not enabled on this server')
- skipped at app.register_model() time so they don't appear in
/v1/models
Default (--engines unset) preserves current behavior: every engine whose
Python deps are importable is enabled.
Motivation: running speech.py on a GPU shared with an LLM server (e.g.
llama-qwen on the same 4090) used to require commenting out
register_model lines AND auditing every caller to never hit
tts-1-qwen. Both forms of discipline broke in practice — speech.py was
seen squatting 10 GiB of VRAM for 3 days because the qwen branch loaded
on a stray request. Single allowlist closes that hole.
Usage:
python speech.py --engines f5,piper # lean: F5 + Piper only
python speech.py --engines tts-1-f5 # equivalent (full IDs OK)
python speech.py # default: all engines
Workers (uvicorn -W N) re-import this module, so the allowlist lives in
SPEECH_ENABLED_ENGINES env var (set by __main__ before uvicorn.run);
each worker re-parses at module load. Module-level set ENABLED_ENGINES
is populated by _parse_engines_env() — None = no restriction.
Validation routes through is_engine_available(model_id), which now ANDs
two gates: (a) operator allowlist, (b) Python-deps importable. Same
helper drives /v1/voices filtering, request-handler short-circuit, and
app.register_model() loop in __main__.
This commit is contained in:
parent
d4ff4e8da4
commit
a0d9e9a966
1 changed files with 87 additions and 11 deletions
98
speech.py
98
speech.py
|
|
@ -67,11 +67,55 @@ except ImportError:
|
|||
torch = None
|
||||
F5TTS = None
|
||||
|
||||
# Engine allowlist controlled via the --engines CLI flag. Lives in the
|
||||
# SPEECH_ENABLED_ENGINES env var because uvicorn workers re-import this
|
||||
# module — a global set populated in __main__ wouldn't propagate to
|
||||
# worker processes. __main__ sets the env var before uvicorn.run; workers
|
||||
# parse it at module load via _parse_engines_env() below.
|
||||
ENGINE_SHORT_TO_MODEL = {
|
||||
'piper': 'tts-1',
|
||||
'xtts': 'tts-1-hd',
|
||||
'silero': 'tts-1-silero',
|
||||
'kokoro': 'tts-1-kokoro',
|
||||
'qwen': 'tts-1-qwen',
|
||||
'f5': 'tts-1-f5',
|
||||
}
|
||||
ALL_MODEL_IDS = set(ENGINE_SHORT_TO_MODEL.values())
|
||||
|
||||
def _parse_engines_env():
|
||||
raw = os.environ.get('SPEECH_ENABLED_ENGINES', '').strip()
|
||||
if not raw:
|
||||
return None
|
||||
allowed = set()
|
||||
for e in (s.strip() for s in raw.split(',')):
|
||||
if not e:
|
||||
continue
|
||||
if e in ENGINE_SHORT_TO_MODEL:
|
||||
allowed.add(ENGINE_SHORT_TO_MODEL[e])
|
||||
elif e in ALL_MODEL_IDS:
|
||||
allowed.add(e)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"SPEECH_ENABLED_ENGINES: unknown engine '{e}' "
|
||||
f"(valid short names: {', '.join(sorted(ENGINE_SHORT_TO_MODEL))}, "
|
||||
f"or full tts-1-* ids)"
|
||||
)
|
||||
return allowed
|
||||
|
||||
ENABLED_ENGINES = _parse_engines_env() # None = all enabled; set = allowlist
|
||||
|
||||
def is_engine_available(model_id):
|
||||
"""Whether the TTS engine for a given model_id is loadable in this process.
|
||||
Used to filter advertised models/voices so we don't 503 on requests for
|
||||
engines whose Python deps weren't installed. tts-1/tts-1-hd/tts-1-silero/
|
||||
tts-1-kokoro are assumed available — they surface their own load errors."""
|
||||
Two gates:
|
||||
1. operator allowlist (--engines / SPEECH_ENABLED_ENGINES): if set
|
||||
and model_id not in it -> unavailable.
|
||||
2. Python deps for engines that conditionally import (qwen, f5): if
|
||||
the import failed at module load -> unavailable.
|
||||
Used by /v1/voices to filter advertised models, by the TTS request
|
||||
handler to short-circuit disabled engines with a clean BadRequestError,
|
||||
and by __main__ to decide which models to app.register_model()."""
|
||||
if ENABLED_ENGINES is not None and model_id not in ENABLED_ENGINES:
|
||||
return False
|
||||
if model_id == 'tts-1-qwen':
|
||||
return QWEN_TTS_AVAILABLE
|
||||
if model_id == 'tts-1-f5':
|
||||
|
|
@ -858,6 +902,15 @@ async def generate_speech(request: GenerateSpeechRequest):
|
|||
else:
|
||||
raise BadRequestError(f"Voice '{voice}' not found in any model. Please specify a model.", param='voice')
|
||||
|
||||
# Operator allowlist check — if --engines was set, disabled engines
|
||||
# short-circuit here with a 4xx (vs trying to load + hitting deeper
|
||||
# 503s downstream).
|
||||
if not is_engine_available(model):
|
||||
raise BadRequestError(
|
||||
f"Model '{model}' is not enabled on this server "
|
||||
f"(operator restricted via --engines)",
|
||||
param='model')
|
||||
|
||||
if (request.timestamps or request.sse) and model != 'tts-1-f5':
|
||||
raise BadRequestError("timestamps and sse modes are currently supported only for model 'tts-1-f5'", param='timestamps')
|
||||
|
||||
|
|
@ -1482,25 +1535,48 @@ if __name__ == "__main__":
|
|||
parser.add_argument('-P', '--port', action='store', default=8000, type=int, help="Server tcp port")
|
||||
parser.add_argument('-H', '--host', action='store', default='0.0.0.0', help="Host to listen on, Ex. 0.0.0.0")
|
||||
parser.add_argument('-L', '--log-level', default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set the log level")
|
||||
parser.add_argument('--engines', action='store', default=None,
|
||||
help="Comma-separated allowlist of TTS engines to enable. "
|
||||
"Short names: f5, qwen, piper, xtts, silero, kokoro "
|
||||
"(or full model IDs like tts-1-f5). "
|
||||
"Default: all engines enabled. "
|
||||
"Example: --engines f5,piper keeps speech lean enough to "
|
||||
"share GPU with an LLM server.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
logger.remove()
|
||||
logger.add(sink=sys.stderr, level=args.log_level)
|
||||
|
||||
# Propagate --engines to uvicorn workers via env var (workers re-import
|
||||
# this module + read SPEECH_ENABLED_ENGINES into ENABLED_ENGINES).
|
||||
if args.engines:
|
||||
os.environ['SPEECH_ENABLED_ENGINES'] = args.engines
|
||||
try:
|
||||
ENABLED_ENGINES = _parse_engines_env()
|
||||
except ValueError as e:
|
||||
logger.error(str(e))
|
||||
sys.exit(2)
|
||||
logger.info(f"--engines: restricting to {sorted(ENABLED_ENGINES)}")
|
||||
else:
|
||||
logger.info("All engines enabled (no --engines restriction)")
|
||||
|
||||
if args.preload and not XTTS_AVAILABLE:
|
||||
logger.error("Cannot preload XTTS model - XTTS dependencies not available")
|
||||
elif args.preload:
|
||||
xtts = xtts_wrapper(args.preload, device=args.xtts_device, unload_timer=args.unload_timer)
|
||||
|
||||
# Register Qwen + F5 by default (other models disabled)
|
||||
# app.register_model('tts-1-qwen') # disabled — frees VRAM for llama-qwen LLM
|
||||
app.register_model('tts-1-f5')
|
||||
# To enable other models, uncomment below:
|
||||
# app.register_model('tts-1')
|
||||
# app.register_model('tts-1-hd')
|
||||
# app.register_model('tts-1-silero')
|
||||
# app.register_model('tts-1-kokoro')
|
||||
# Register every model whose engine is_engine_available — combines the
|
||||
# operator --engines allowlist with the Python-deps check. Anything
|
||||
# missing here disappears from /v1/models AND short-circuits with a
|
||||
# BadRequestError at the TTS request handler (see is_engine_available
|
||||
# callsite up top).
|
||||
for model_id in sorted(ALL_MODEL_IDS):
|
||||
if is_engine_available(model_id):
|
||||
app.register_model(model_id)
|
||||
logger.info(f"Registered model: {model_id}")
|
||||
else:
|
||||
logger.info(f"Skipped model: {model_id}")
|
||||
|
||||
# Use multiple workers for true concurrency (each worker = separate process with own GIL)
|
||||
# This prevents thread pool exhaustion and allows concurrent model loading
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue