c86e323d7a
Quality-first multi-pass audiobook generation from English manuscripts. Includes: - umbraaudio/ core package (pipeline, config, voices, web review UI) - scripts/ CLI utilities + model download script - tts_server/ remote TTS server - specs/ design docs, tests/, pyproject.toml + uv.lock - Qwen3-TTS git submodule (pinned 022e286) with local demo.py UI patch saved at scripts/qwen3_tts_demo_umbra_patch.patch - .env with runtime config/secrets Excluded (see .gitignore): models/ (downloaded via scripts/download_models.sh), output/, logs/, .venv, .umbra runtime data, and book projects (DealHunterShip_*), which are committed separately.
116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
"""CLI for the model coordinator.
|
|
|
|
Ensures only the model needed by the current pipeline stage is resident in
|
|
memory, freeing maximum GPU/RAM for whichever model is about to work.
|
|
|
|
Usage:
|
|
python scripts/model_coordinator.py status
|
|
python scripts/model_coordinator.py switch \\
|
|
voice_design|chapter_analysis|utterance_generation|idle
|
|
python scripts/model_coordinator.py unload-all
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
from umbraaudio.config import load_config
|
|
from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
console = Console()
|
|
|
|
|
|
def _render_status(state: dict) -> None:
|
|
qwen_model = state.get("qwen_model") or "—"
|
|
lm_loaded = state.get("lm_loaded") or []
|
|
lm_str = ", ".join(lm_loaded) if lm_loaded else "—"
|
|
table = Table(title="Model Coordinator Status")
|
|
table.add_column("Server", style="bold cyan")
|
|
table.add_column("State", style="bold green")
|
|
table.add_row("Qwen TTS server", state.get("qwen_server", "unknown"))
|
|
table.add_row("Qwen current model", qwen_model)
|
|
table.add_row(" voicedesign", "loaded" if state.get("voicedesign_loaded") else "unloaded")
|
|
table.add_row(" base", "loaded" if state.get("base_loaded") else "unloaded")
|
|
table.add_row("LM Studio loaded", lm_str)
|
|
console.print(table)
|
|
|
|
|
|
def cmd_status(coord: ModelCoordinator, _args: argparse.Namespace) -> int:
|
|
_render_status(coord.status())
|
|
return 0
|
|
|
|
|
|
def cmd_switch(coord: ModelCoordinator, args: argparse.Namespace) -> int:
|
|
try:
|
|
stage = PipelineStage(args.stage)
|
|
except ValueError:
|
|
console.print(
|
|
f"[red]Unknown stage '{args.stage}'. "
|
|
f"Choose from: {', '.join(s.value for s in PipelineStage)}[/red]"
|
|
)
|
|
return 2
|
|
console.print(f"Switching to stage [bold]{stage.value}[/bold]...")
|
|
result = coord.switch_to(stage)
|
|
if not result.get("ok"):
|
|
console.print(f"[red]Failed: {result.get('error', 'unknown error')}[/red]")
|
|
console.print("Before:", json.dumps(result.get("before", {}), indent=2))
|
|
return 1
|
|
console.print("[green]OK[/green]")
|
|
_render_status(result["after"])
|
|
return 0
|
|
|
|
|
|
def cmd_unload_all(coord: ModelCoordinator, _args: argparse.Namespace) -> int:
|
|
console.print("Unloading all models...")
|
|
result = coord.switch_to(PipelineStage.IDLE)
|
|
if not result.get("ok"):
|
|
console.print(f"[red]Failed: {result.get('error', 'unknown error')}[/red]")
|
|
return 1
|
|
console.print("[green]All models unloaded[/green]")
|
|
_render_status(result["after"])
|
|
return 0
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Switch which models are resident for the current pipeline stage."
|
|
)
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
|
|
p_status = sub.add_parser("status", help="Show which models are currently loaded")
|
|
p_status.set_defaults(func=cmd_status)
|
|
|
|
p_switch = sub.add_parser("switch", help="Switch to a pipeline stage's model set")
|
|
p_switch.add_argument(
|
|
"stage",
|
|
choices=[s.value for s in PipelineStage],
|
|
help="Which stage's models should be resident",
|
|
)
|
|
p_switch.set_defaults(func=cmd_switch)
|
|
|
|
p_unload = sub.add_parser("unload-all", help="Unload all models from both servers")
|
|
p_unload.set_defaults(func=cmd_unload_all)
|
|
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(message)s")
|
|
args = build_parser().parse_args(argv)
|
|
coord = ModelCoordinator(load_config())
|
|
return args.func(coord, args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|