path: root/backends/faster.py

#!/usr/bin/env python3
"""Set up the faster-qwen3-tts backend for the audiobook generator.

faster-qwen3-tts is an OpenAI-compatible Qwen3-TTS server with CUDA-graph
inference (NVIDIA GPU required). It always uses voice cloning, with the
reference voice configured on the server through a ``voices.json``. This
module sets the whole backend up end-to-end as a TUI: pip-install the
package, clone the repo (for ``examples/openai_server.py``), build a
``voices.json`` from a directory of .wav references (transcribed with
Whisper), sync ``converter/config.py``, and print the launch command. It is
driven by ``audiobook.py``'s hub but can also be run directly with flags.

Usage:
    python -m backends.faster [--wavs WAV_DIR] [--output PATH]
        [--language LANG] [--whisper-model NAME] [--force]
        [--port PORT] [--voice NAME] [--skip-install] [--skip-clone]
"""

import argparse
import importlib.util
import json
import sys
from pathlib import Path
from typing import List, Optional

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import tui
from backends import BackendStatus, ModifyAction
from backends import common
from backends.common import TTS_ROOT, find_wav_files, normalize_dir_arg
from converter import config
from converter.tts import normalize_language, transcribe_reference_audio, \
    whisper_backend_available

FASTER_DIR_NAME = "faster-qwen3-tts"
FASTER_GIT_URL = "https://github.com/andimarafioti/faster-qwen3-tts"
FASTER_PIP_PKG = "faster-qwen3-tts[demo]"
WHISPER_MODELS = ("tiny", "base", "small", "medium", "large-v3")


def _checkout() -> Path:
    return TTS_ROOT / FASTER_DIR_NAME


def _is_installed() -> bool:
    return importlib.util.find_spec("faster_qwen3_tts") is not None


def _is_cloned() -> bool:
    return (_checkout() / "examples" / "openai_server.py").is_file()


def _config_port() -> int:
    import urllib.parse
    try:
        return urllib.parse.urlsplit(config.FASTER_API_URL).port or 8000
    except ValueError:
        return 8000


def build_voices(wav_files: list, language: str, whisper_model: str) -> dict:
    """Transcribe each wav file and build the voices mapping."""
    voices = {}
    for wav_file in wav_files:
        name = wav_file.stem
        print(f"[INFO] Transcribing {wav_file.name} (voice '{name}')...")
        text = transcribe_reference_audio(str(wav_file), model_name=whisper_model)
        if text:
            print(f"[OK] {name}: {text}")
        else:
            print(f"[WARNING] No transcript for '{name}'; the faster backend "
                  "strongly recommends an accurate transcript — consider "
                  "editing voices.json by hand before starting the server")
        voices[name] = {
            "ref_audio": str(wav_file.resolve()),
            "ref_text": text or "",
            "language": language,
        }
    return voices


def _write_voices_json(output_path: Path, wav_dir: Path, language: str,
                       whisper_model: str, force: bool) -> Optional[dict]:
    """Transcribe the wav dir and write voices.json; return the voices dict."""
    wav_files = find_wav_files(wav_dir)
    if not wav_files:
        print(f"[ERROR] No .wav files found in {wav_dir}")
        return None
    if whisper_backend_available() is None:
        print("[WARNING] Neither faster_whisper nor whisper was found, so "
              "transcripts will be empty — install one or edit voices.json "
              "by hand.")
    voices = build_voices(wav_files, language, whisper_model)
    with output_path.open("w", encoding="utf-8") as handle:
        json.dump(voices, handle, indent=4, ensure_ascii=False)
        handle.write("\n")
    print(f"[OK] Wrote {output_path} with {len(voices)} voice(s): "
          f"{', '.join(voices)}")
    return voices


def _wizard(stdscr, args: argparse.Namespace) -> Optional[dict]:
    """Linear TUI wizard collecting every faster-setup decision."""
    _GO_BACK = object()

    def confirm(question: str, default: bool = True) -> Optional[bool]:
        res = tui.confirm(stdscr, question, default=default,
                          cancel_value=_GO_BACK)
        return None if res is _GO_BACK else res

    # Step 0: pip install (if not installed and not skipped).
    do_install = False
    if not _is_installed() and not args.skip_install:
        choice = confirm("faster-qwen3-tts is not installed. "
                         "pip install it now?", default=True)
        if choice is None:
            return None
        do_install = choice

    # Step 1: clone (if not cloned and not skipped).
    do_clone = False
    if not _is_cloned() and not args.skip_clone:
        choice = confirm(f"faster-qwen3-tts repo not cloned. Clone it into "
                         f"./{FASTER_DIR_NAME}?", default=True)
        if choice is None:
            return None
        do_clone = choice

    # Step 2: voices.json — wav dir, language, whisper model, output path.
    wav_dir = args.input_dir
    if wav_dir is None:
        wav_dir = tui.browse_directory(
            stdscr, "Select the directory with your .wav voices",
            info=common.wav_dir_info, preview=common.wav_dir_preview,
            start=Path.cwd())
    language = args.language
    if language is None:
        lang_text = tui.line_edit(
            stdscr, "Language", config.LANGUAGE,
            validate=lambda s: None if _try_language(s)
            else "Unknown language (e.g. English, en)",
            help_lines=["Language for every voice, as passed to the TTS "
                        "model (names or short codes accepted)"])
        language = lang_text
    whisper_model = args.whisper_model
    if whisper_model is None:
        whisper_model = tui.menu(
            stdscr, "Whisper model for transcription",
            [(m, m) for m in WHISPER_MODELS],
            default_index=WHISPER_MODELS.index("base"))
    output_path = args.output
    if output_path is None:
        # Default into the cloned checkout; fall back to the wav directory
        # when the checkout is not present (so a flag-only run still works).
        output_path = (_checkout() / "voices.json") if _is_cloned() \
            else (wav_dir / "voices.json")
        if output_path.exists() and not args.force:
            choice = confirm(f"{output_path} already exists. Overwrite?",
                             default=True)
            if choice is None or choice is False:
                # Fall back to a path in the current directory.
                output_path = Path.cwd() / "voices.json"

    # Step 3: port + default voice.
    port = args.port
    if port is None:
        port_text = tui.line_edit(
            stdscr, "Server port", str(_config_port()),
            validate=lambda s: None if (s.isdigit() and 1 <= int(s) <= 65535)
            else "Enter a port number between 1 and 65535")
        port = int(port_text)

    return {
        "do_install": do_install,
        "do_clone": do_clone,
        "wav_dir": wav_dir,
        "language": language,
        "whisper_model": whisper_model,
        "output_path": output_path,
        "port": port,
        "force": args.force,
    }


def _try_language(value: str) -> bool:
    try:
        normalize_language(value)
        return True
    except ValueError:
        return False


def _execute(settings: dict) -> int:
    """Console tail: install, clone, write voices.json, sync, advise."""
    if settings["do_install"]:
        rc = common.pip_install([FASTER_PIP_PKG])
        if rc != 0:
            print(f"[WARNING] pip install failed (exit {rc}); install "
                  f"{FASTER_PIP_PKG} manually")
        else:
            print("[OK] faster-qwen3-tts installed")

    if settings["do_clone"]:
        rc = common.git_clone(FASTER_GIT_URL, _checkout())
        if rc != 0:
            print(f"[WARNING] git clone failed (exit {rc}); clone manually: "
                  f"git clone {FASTER_GIT_URL} {_checkout()}")
        else:
            print(f"[OK] cloned into {_checkout()}")

    voices = _write_voices_json(settings["output_path"], settings["wav_dir"],
                                settings["language"], settings["whisper_model"],
                                settings["force"])
    if voices is None:
        return 1

    # Sync converter/config.py port + default voice.
    port = settings["port"]
    new_url = common.url_with_port(config.FASTER_API_URL, port)
    if new_url != config.FASTER_API_URL:
        if common.update_config_value("FASTER_API_URL", new_url):
            print(f"[OK] Updated FASTER_API_URL to {new_url}")
        else:
            print("[WARNING] Could not update FASTER_API_URL; edit "
                  "converter/config.py by hand")
    default_voice = next(iter(voices))
    if default_voice != config.FASTER_VOICE:
        if common.update_config_value("FASTER_VOICE", default_voice):
            print(f"[OK] Updated FASTER_VOICE to {default_voice}")
        else:
            print("[WARNING] Could not update FASTER_VOICE; edit "
                  "converter/config.py by hand")

    _print_launch_hint(settings["output_path"], port)
    return 0


def _print_launch_hint(voices_path: Path, port: int) -> None:
    print()
    if _is_cloned():
        print("Start the server with:")
        print(f"  python {_checkout()}/examples/openai_server.py "
              f"--voices {voices_path} --port {port}")
    else:
        print("[INFO] Clone faster-qwen3-tts to get examples/openai_server.py,")
        print(f"        then run it with --voices {voices_path} --port {port}")


def run_tui(args: Optional[argparse.Namespace] = None) -> int:
    """Run the faster setup wizard end-to-end."""
    import curses
    if args is None:
        args = build_parser().parse_args([])
    try:
        settings = curses.wrapper(_wizard, args)
    except tui.WizardCancelled:
        print("\n[INFO] Cancelled; nothing was written")
        return 1
    try:
        curses.curs_set(1)
    except curses.error:
        pass
    if settings is None:
        print("[INFO] Aborted")
        return 1
    return _execute(settings)


def _collect_from_flags(args: argparse.Namespace,
                        parser: argparse.ArgumentParser) -> Optional[dict]:
    """Build the settings dict from flags for a non-interactive run."""
    if args.input_dir is None:
        parser.error("--wavs is required in a non-interactive run (or run "
                     "without flags for the TUI wizard)")
    if not args.input_dir.is_dir():
        parser.error(f"WAV directory not found: {args.input_dir}")
    try:
        language = normalize_language(args.language or config.LANGUAGE)
    except ValueError as exc:
        parser.error(str(exc))
    output_path = args.output if args.output is not None \
        else ((_checkout() / "voices.json") if _is_cloned()
              else (args.input_dir / "voices.json"))
    if output_path.exists() and not args.force:
        print("[INFO] Aborted; existing voices.json kept")
        return None
    return {
        "do_install": (not _is_installed()) and not args.skip_install,
        "do_clone": (not _is_cloned()) and not args.skip_clone,
        "wav_dir": args.input_dir,
        "language": language,
        "whisper_model": args.whisper_model or "base",
        "output_path": output_path,
        "port": args.port if args.port is not None else _config_port(),
        "force": args.force,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Set up the faster-qwen3-tts backend: pip install, clone, "
                    "build voices.json, and sync converter/config.py.")
    parser.add_argument("input_dir", type=normalize_dir_arg, nargs="?",
                        default=None, metavar="WAV_DIR",
                        help="Directory with .wav reference files (required in "
                             "a non-interactive run; browsed for in the TUI)")
    parser.add_argument("--output", type=Path, default=None,
                        help="Output path for voices.json (default: "
                             "./faster-qwen3-tts/voices.json, or "
                             "WAV_DIR/voices.json when not cloned)")
    parser.add_argument("--language", type=str, default=None,
                        help="Language for all voices (default: English; "
                             "names and short codes accepted)")
    parser.add_argument("--whisper-model", type=str, default=None,
                        choices=WHISPER_MODELS,
                        help="Whisper model size for transcription "
                             "(default: base)")
    parser.add_argument("--force", action="store_true",
                        help="Overwrite an existing voices.json without "
                             "prompting")
    parser.add_argument("--port", type=int, default=None,
                        help="Server port to record in converter/config.py "
                             "(default: the port in FASTER_API_URL)")
    parser.add_argument("--skip-install", action="store_true",
                        help="Do not pip install faster-qwen3-tts[demo]")
    parser.add_argument("--skip-clone", action="store_true",
                        help="Do not clone the faster-qwen3-tts repo")
    return parser


def detect() -> BackendStatus:
    """Detect how far faster-qwen3-tts is set up, plus the launch command."""
    installed = _is_installed()
    cloned = _is_cloned()
    voices_json = _checkout() / "voices.json"
    configured = installed and cloned and voices_json.exists()
    details: List[str] = []
    details.append("pip: installed" if installed else
                   "not installed — run setup to pip install")
    details.append(f"checkout: {_checkout()}" if cloned else
                   f"not cloned — run setup to clone ./{FASTER_DIR_NAME}")
    details.append(f"voices: {voices_json}" if voices_json.exists() else
                   "no voices.json — run setup to create one")
    launch = ""
    if cloned and voices_json.exists():
        launch = (f"python {_checkout()}/examples/openai_server.py "
                  f"--voices {voices_json} --port {_config_port()}")
    return BackendStatus("faster", "faster-qwen3-tts",
                         installed=installed and cloned,
                         configured=configured, details=details,
                         launch_hint=launch)


def _run_voices_only_tui() -> int:
    """Rebuild voices.json via the TUI (the "modify" action).

    Runs the same wizard but skips the pip/clone prerequisites so it goes
    straight to picking the .wav directory and writing voices.json.
    """
    args = build_parser().parse_args([])
    args.skip_install = True
    args.skip_clone = True
    return run_tui(args)


modify_actions: List[ModifyAction] = [
    ModifyAction("Rebuild voices.json", _run_voices_only_tui),
    ModifyAction("Reconfigure faster-qwen3-tts", run_tui),
]


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if _interactive():
        return run_tui(args)

    settings = _collect_from_flags(args, parser)
    if settings is None:
        return 1
    return _execute(settings)


def _interactive() -> bool:
    try:
        import curses  # noqa: F401
    except ImportError:
        return False
    try:
        return sys.stdin.isatty() and sys.stdout.isatty()
    except (AttributeError, ValueError):
        return False


if __name__ == "__main__":
    sys.exit(main())