path: root/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 ``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 -m backends.qwen [--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,
    ConfigureAction,
    ServerSpec,
    common,
    envs,
    format_launch_hint,
)
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 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."""
    _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("qwen-tts is not installed. pip install it now?",
                         default=True)
        if choice is None:
            return None
        do_install = choice

    # Step 1: ports.
    custom_port = args.port_custom
    if custom_port is None:
        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)"])
        custom_port = int(port_text)
    clone_port = args.port_clone
    if clone_port is None:
        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)"])
        clone_port = int(port_text)

    # Step 2: built-in speaker.
    speaker = args.speaker
    if speaker is None:
        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"])

    return {
        "do_install": do_install,
        "custom_port": custom_port,
        "clone_port": clone_port,
        "speaker": speaker,
    }


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 "
                  "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 "
                  "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 "
                  "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 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)
    # Running when either server is up — CustomVoice (speaker mode) or Base
    # (voice clone) each suffice for a conversion on their own.
    running = (common.server_running(config.QWEN_API_URL)
               or common.server_running(config.CLONE_API_URL))
    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"))
    servers = [
        ServerSpec("qwen-custom", config.QWEN_API_URL,
                   [demo, QWEN_CUSTOMVOICE_MODEL, "--ip", "127.0.0.1",
                    "--port", str(custom_port)]),
        ServerSpec("qwen-clone", config.CLONE_API_URL,
                   [demo, QWEN_BASE_MODEL, "--ip", "127.0.0.1",
                    "--port", str(clone_port)]),
    ]
    return BackendStatus("qwen", "qwen-tts",
                         installed=installed, configured=installed,
                         running=running, details=details,
                         launch_hint=format_launch_hint(servers),
                         servers=servers)


configure_actions: List[ConfigureAction] = [
    ConfigureAction("Reconfigure qwen-tts (ports/speaker)", run_tui),
]


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())