path: root/app/backends/qwen.py

#!/usr/bin/env python3
"""Set up the Qwen3-TTS demo backend for the audiobook generator.

qwen-tts is a pip package providing the ``qwen-tts-demo`` server, which
hosts the Qwen3-TTS CustomVoice (built-in speakers) and Base (voice
cloning) models on separate ports. This module sets it up end-to-end as a
TUI: pip-install the package, configure the two ports and the built-in
speaker in ``app/converter/config.py``, and print the launch commands. It is
driven by ``audiobook.py``'s hub but can also be run directly with flags.

Usage:
    python app/backends/qwen.py [--port-custom PORT] [--port-clone PORT]
        [--speaker NAME] [--skip-install]
"""

import argparse
import sys
from pathlib import Path
from typing import List, Optional

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

from backends import (
    BackendStatus,
    ServerSpec,
    common,
    envs,
    format_launch_hint,
    probe,
    servers,
)
from converter import config
from ui import tui

QWEN_PIP_PKG = "qwen-tts"
QWEN_CUSTOMVOICE_MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
QWEN_BASE_MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
DEFAULT_CUSTOM_PORT = 7860
DEFAULT_CLONE_PORT = 7861

# Built-in CustomVoice speakers (see app/converter/config.py SPEAKER).
QWEN_SPEAKERS = ("Vivian", "Serena", "Uncle_Fu", "Dylan", "Eric", "Ryan",
                 "Aiden", "Ono_Anna", "Sohee")


def _is_installed() -> bool:
    if envs.env_script("qwen-tts-demo").is_file():
        return True
    return envs.module_available("qwen_tts")


def _config_port(url: str, fallback: int) -> int:
    import urllib.parse
    try:
        return urllib.parse.urlsplit(url).port or fallback
    except ValueError:
        return fallback


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

    Driven by ``tui.Wizard`` as a stack of screen closures: each screen
    shows one widget and returns the next screen, ``Wizard.BACK`` (Esc/q —
    pop to the previous screen), or the settings dict. Steps whose value is
    already provided by a flag (``--port-custom``, ``--port-clone``,
    ``--speaker``, ``--skip-install``) are folded into the ``_after_*``
    guards and never become screens, so Esc always lands on the previous
    real screen. Esc on the first screen aborts the wizard.
    """
    _GO_BACK = object()
    s: dict = {}

    def _after_install():
        if args.port_custom is None:
            return screen_custom_port
        s["custom_port"] = args.port_custom
        return _after_custom_port()

    def screen_custom_port():
        port_text = tui.line_edit(
            stdscr, "CustomVoice (built-in speaker) port",
            str(_config_port(config.QWEN_API_URL, DEFAULT_CUSTOM_PORT)),
            validate=lambda s: None if (s.isdigit() and 1 <= int(s) <= 65535)
            else "Enter a port number between 1 and 65535",
            help_lines=["The port for qwen-tts-demo CustomVoice (speaker mode)"],
            back_value=_GO_BACK)
        if port_text is _GO_BACK:
            return tui.Wizard.BACK
        s["custom_port"] = int(port_text)
        return _after_custom_port()

    def _after_custom_port():
        if args.port_clone is None:
            return screen_clone_port
        s["clone_port"] = args.port_clone
        return _after_clone_port()

    def screen_clone_port():
        port_text = tui.line_edit(
            stdscr, "Base (voice clone) port",
            str(_config_port(config.CLONE_API_URL, DEFAULT_CLONE_PORT)),
            validate=lambda s: None if (s.isdigit() and 1 <= int(s) <= 65535)
            else "Enter a port number between 1 and 65535",
            help_lines=["The port for qwen-tts-demo Base (voice cloning)"],
            back_value=_GO_BACK)
        if port_text is _GO_BACK:
            return tui.Wizard.BACK
        s["clone_port"] = int(port_text)
        return _after_clone_port()

    def _after_clone_port():
        if args.speaker is None:
            return screen_speaker
        s["speaker"] = args.speaker
        return _finalize()

    def screen_speaker():
        speaker = tui.menu(
            stdscr, "Built-in CustomVoice speaker",
            [(s, s) for s in QWEN_SPEAKERS],
            default_index=max(0, QWEN_SPEAKERS.index(config.SPEAKER)
                              if config.SPEAKER in QWEN_SPEAKERS else 0),
            help_lines=["Used by audiobook.py --backend qwen without --clone"],
            back_value=_GO_BACK)
        if speaker is _GO_BACK:
            return tui.Wizard.BACK
        s["speaker"] = speaker
        return _finalize()

    def _finalize() -> dict:
        return {
            "do_install": s.get("do_install", False),
            "custom_port": s["custom_port"],
            "clone_port": s["clone_port"],
            "speaker": s["speaker"],
        }

    # pip install happens without asking: when the package is missing (and
    # not skipped by flag), the wizard just does it and moves to the next
    # screen.
    s["do_install"] = (not _is_installed()) and not args.skip_install
    return tui.Wizard().run(_after_install())


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

    custom_url = common.url_with_port(config.QWEN_API_URL, settings["custom_port"])
    if custom_url != config.QWEN_API_URL:
        if common.update_config_value("QWEN_API_URL", custom_url):
            print(f"[OK] Updated QWEN_API_URL to {custom_url}")
        else:
            print("[WARNING] Could not update QWEN_API_URL; edit "
                  "app/converter/config.py by hand")
    clone_url = common.url_with_port(config.CLONE_API_URL, settings["clone_port"])
    if clone_url != config.CLONE_API_URL:
        if common.update_config_value("CLONE_API_URL", clone_url):
            print(f"[OK] Updated CLONE_API_URL to {clone_url}")
        else:
            print("[WARNING] Could not update CLONE_API_URL; edit "
                  "app/converter/config.py by hand")
    if settings["speaker"] != config.SPEAKER:
        if common.update_config_value("SPEAKER", settings["speaker"]):
            print(f"[OK] Updated SPEAKER to {settings['speaker']}")
        else:
            print("[WARNING] Could not update SPEAKER; edit "
                  "app/converter/config.py by hand")

    _print_launch_hint(settings["custom_port"], settings["clone_port"])
    return 0


def _print_launch_hint(custom_port: int, clone_port: int) -> None:
    demo = envs.env_script("qwen-tts-demo")
    print()
    print("Start the servers (in separate terminals), or use the hub's")
    print("'Server' menu / let a conversion start one automatically:")
    print(f"  {demo} {QWEN_CUSTOMVOICE_MODEL} --ip 127.0.0.1 "
          f"--port {custom_port}")
    print(f"  {demo} {QWEN_BASE_MODEL} --ip 127.0.0.1 "
          f"--port {clone_port}")
    print("Then run: python audiobook.py --backend qwen")


def setup_screen(stdscr) -> int:
    """Run the setup wizard on an existing curses screen (the hub's).

    The hub drives this as one screen of its own ``tui.Wizard`` stack, so
    Esc on the wizard's first screen simply returns here and the hub pops
    back to the menu that launched it. The console tail (pip install /
    config sync) runs under ``tui.suspend`` so the hub's curses session
    stays intact. Returns 0 on completion, 1 when the user aborted.
    """
    args = build_parser().parse_args([])
    settings = _wizard(stdscr, args)
    if settings is None:
        return 1
    with tui.suspend(stdscr):
        return _execute(settings)


def run_tui(args: Optional[argparse.Namespace] = None) -> int:
    """Run the qwen 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) -> dict:
    return {
        "do_install": (not _is_installed()) and not args.skip_install,
        "custom_port": args.port_custom if args.port_custom is not None
        else _config_port(config.QWEN_API_URL, DEFAULT_CUSTOM_PORT),
        "clone_port": args.port_clone if args.port_clone is not None
        else _config_port(config.CLONE_API_URL, DEFAULT_CLONE_PORT),
        "speaker": args.speaker or config.SPEAKER,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Set up the Qwen3-TTS demo backend: pip install, "
                    "configure ports/speaker, and print launch commands.")
    parser.add_argument("--port-custom", type=int, default=None,
                        help="CustomVoice (speaker) port (default: "
                             f"{DEFAULT_CUSTOM_PORT})")
    parser.add_argument("--port-clone", type=int, default=None,
                        help="Base (voice clone) port (default: "
                             f"{DEFAULT_CLONE_PORT})")
    parser.add_argument("--speaker", type=str, default=None,
                        choices=QWEN_SPEAKERS,
                        help="Built-in CustomVoice speaker (default: "
                             f"{config.SPEAKER})")
    parser.add_argument("--skip-install", action="store_true",
                        help="Do not pip install qwen-tts")
    return parser


def detect() -> BackendStatus:
    """Detect whether qwen-tts is installed, plus the launch commands."""
    installed = _is_installed()
    custom_port = _config_port(config.QWEN_API_URL, DEFAULT_CUSTOM_PORT)
    clone_port = _config_port(config.CLONE_API_URL, DEFAULT_CLONE_PORT)
    details: List[str] = []
    details.append("pip: installed" if installed else
                   "not installed — run setup to pip install qwen-tts")
    details.append(f"CustomVoice port: {custom_port}")
    details.append(f"Base (clone) port: {clone_port}")
    details.append(f"speaker: {config.SPEAKER}")
    demo = str(envs.env_script("qwen-tts-demo"))
    specs = [
        ServerSpec("qwen-custom", config.QWEN_API_URL,
                   [demo, QWEN_CUSTOMVOICE_MODEL, "--ip", "127.0.0.1",
                    "--port", str(custom_port)],
                   identity=probe.IDENTITY_QWEN_CUSTOM),
        ServerSpec("qwen-clone", config.CLONE_API_URL,
                   [demo, QWEN_BASE_MODEL, "--ip", "127.0.0.1",
                    "--port", str(clone_port)],
                   identity=probe.IDENTITY_QWEN_CLONE),
    ]
    managed = servers.manages(specs)
    # Which local servers this tool started (pid alive) name the running
    # models; a remotely-run demo names them via the probe instead.
    local_models = [name for name, spec in
                    (("Base", specs[1]), ("CustomVoice", specs[0]))
                    if servers.alive(spec.name)]
    remote_models, remote_urls = _detect_remote(managed)
    running_models = list(dict.fromkeys(local_models + remote_models))
    return BackendStatus("qwen", "qwen-tts",
                         installed=installed, configured=installed,
                         running=managed or bool(remote_urls),
                         details=details,
                         launch_hint=format_launch_hint(specs),
                         servers=specs,
                         managed=managed,
                         remote=bool(remote_urls),
                         remote_urls=remote_urls,
                         remote_models=remote_models,
                         running_models=running_models)


def _detect_remote(managed: bool = False):
    """Detect externally-run qwen demo servers at the remote URLs.

    Returns ``([model, ...], {spec_name: url})``. Each remote URL (CustomVoice
    and Base) is probed independently and must answer as the matching demo
    (see probe.identify_server); a remote URL equal to the local URL for a
    server this tool started is ignored (already reported "[local]").
    """
    remote_models = []
    remote_urls = {}
    for spec_name, url, local_url, identity in (
            ("qwen-clone", config.CLONE_REMOTE_URL, config.CLONE_API_URL,
             probe.IDENTITY_QWEN_CLONE),
            ("qwen-custom", config.QWEN_REMOTE_URL, config.QWEN_API_URL,
             probe.IDENTITY_QWEN_CUSTOM)):
        url = (url or "").strip()
        if not url:
            continue
        if managed and probe.same_endpoint(url, local_url):
            continue
        if probe.identify_server(url) == identity:
            remote_urls[spec_name] = url
            remote_models.append(
                "Base" if spec_name == "qwen-clone" else "CustomVoice")
    return remote_models, remote_urls


def uninstall() -> int:
    """Remove the qwen-tts backend entirely: stop its servers, pip uninstall.

    qwen-tts is a pip package (``qwen_tts`` + the ``qwen-tts-demo`` script)
    installed into the managed venv, so uninstalling it removes the backend.
    Any server this tool started is stopped first (best-effort). Returns the
    exit code.
    """
    servers.stop("qwen-custom")
    servers.stop("qwen-clone")
    rc = common.pip_uninstall([QWEN_PIP_PKG])
    if rc != 0:
        print(f"[WARNING] pip uninstall failed (exit {rc}); remove "
              f"{QWEN_PIP_PKG} from the managed venv manually")
    else:
        print(f"[OK] {QWEN_PIP_PKG} removed.")
    return 0


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

    if _interactive():
        return run_tui(args)

    settings = _collect_from_flags(args, parser)
    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())