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,
    setup,
)
from converter import config
from converter.clients import QWEN3_TTS_SPEAKERS
from ui import taskview, 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). The
# canonical list lives in converter.clients.speakers (shared with the
# audio.cpp backend's Convert-form Speaker picker).
QWEN_SPEAKERS = QWEN3_TTS_SPEAKERS


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_steps(settings: dict) -> List[taskview.TaskStep]:
    """Build the ordered setup steps for the in-TUI task view.

    The same work ``_execute`` runs on the console, split into named steps so
    the view can show per-step state and progress. The pip install streams
    through EMIT and aborts on CANCEL; print()-based steps are captured by
    the view's stdout redirect.
    """
    steps: List[taskview.TaskStep] = []

    if settings["do_install"]:
        def install(emit, cancel):
            rc = common.pip_install([QWEN_PIP_PKG], emit=emit, cancel=cancel)
            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")
            return rc
        steps.append(taskview.TaskStep(f"Install {QWEN_PIP_PKG}", install))

    def sync(emit, cancel):
        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")
        return 0
    steps.append(taskview.TaskStep("Sync config & ports", sync))

    return steps


def _execute(settings: dict) -> int:
    """Console tail: install, sync config, advise."""
    return taskview.run_steps_inline(_execute_steps(settings))


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

    See backends.setup.screen_flow for the shared flow. Returns 0 on
    completion, 1 when the user aborted.
    """
    return setup.screen_flow(stdscr, wizard=_wizard,
                             steps_of=_execute_steps,
                             title="Setting up qwen-tts",
                             parser_factory=build_parser)


def run_tui(args: Optional[argparse.Namespace] = None) -> int:
    """Run the qwen setup wizard end-to-end."""
    if args is None:
        args = build_parser().parse_args([])
    return setup.tui_flow(_wizard, _execute, args=args,
                          aborted_message="[INFO] Aborted")


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(*, emit=None, cancel=None) -> 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).

    With EMIT given (the in-TUI task view) pip runs piped, streaming into
    EMIT, so its output never touches the terminal behind curses. CANCEL is
    a ``threading.Event`` honored between phases only (after the servers
    have been stopped, before pip starts) — a started phase always completes,
    so pip is never killed mid-run. Returns the exit code (130 when
    cancelled before pip ran).
    """
    for name in ("qwen-custom", "qwen-clone"):
        # Only stop when a pid file exists: without one this tool never
        # started the server, so the "not started by this tool" notice
        # would be uninstall-time noise.
        if servers.pid_for(name) is not None:
            servers.stop(name)
    if common.cancel_requested(cancel):
        return 130
    rc = common.pip_uninstall([QWEN_PIP_PKG], emit=emit)
    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 rc


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

    if setup.interactive():
        return run_tui(args)

    settings = _collect_from_flags(args, parser)
    return _execute(settings)


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