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
ONE Qwen3-TTS model per process — CustomVoice (built-in speakers), Base
(voice cloning) or VoiceDesign (described voice). This module sets it up
end-to-end: pip-install the package into its own managed venv
(``app/envs/qwen`` — the app venv and the faster backend's never receive
it). There are no questions to ask — the port lives in
``app/converter/config.py`` (the hub's Settings screen) and the model is
chosen per run on the hub's Generate-audiobooks screen; only one server
runs at a time.

Model weights are not part of the install: each demo lazily fetches its
~4GB snapshot from HuggingFace into the standard hub cache the first time a
server for it starts. This module tracks those three cache directories so
weights can be pre-fetched ("Install") or deleted ("Uninstall") per model —
via the hub's Configure-qwen-tts screen — and so the backend uninstaller can
remove every downloaded weight alongside the package.

It is driven by ``audiobook.py``'s hub but can also be run directly:

Usage:
    python app/backends/qwen.py [--skip-install]
"""

import argparse
import os
import shutil
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, VOICE_MODE_CLONE,
                               VOICE_MODE_CUSTOM, VOICE_MODE_DESIGN)
from ui import taskview, tui

QWEN_PIP_PKG = "qwen-tts"
# The dedicated venv qwen-tts is installed into (never the app env or the
# faster backend's).
QWEN_ENV = envs.QWEN_ENV_DIR
DEFAULT_PORT = 7860

# The models a single demo server can host. A running server identifies
# itself via its probe identity (backends.probe), so "which model is up" is
# always read off the server, never assumed.
MODEL_REPOS = {
    "CustomVoice": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
    "Base": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
    "VoiceDesign": "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
}
# Probe identity -> the model name reported in statuses/menus.
IDENTITY_TO_MODEL = {
    probe.IDENTITY_QWEN_CUSTOM: "CustomVoice",
    probe.IDENTITY_QWEN_CLONE: "Base",
    probe.IDENTITY_QWEN_DESIGN: "VoiceDesign",
}
# The model a fresh managed start (Start/Stop Backend Servers menu) hosts;
# Generate-audiobooks runs needing another model restart the server.
DEFAULT_MODEL = "CustomVoice"

# Built-in CustomVoice speakers. The canonical list lives in
# converter.clients.speakers (shared with the audio.cpp backend's
# Convert-form Speaker picker); a run's speaker is picked per run.
QWEN_SPEAKERS = QWEN3_TTS_SPEAKERS


# -- HuggingFace weight cache ------------------------------------------------
#
# The demo servers pull each model's snapshot into huggingface_hub's default
# hub cache on first start (nothing in this project redirects it). These
# helpers read and delete the same directories ``from_pretrained`` writes,
# honoring the same environment overrides, so per-model (un)installs land
# exactly where a server start would look.

def _hf_cache_dir() -> Path:
    """The HF hub cache dir servers fetch model weights into.

    Resolution mirrors huggingface_hub.constants: HF_HUB_CACHE beats
    HUGGINGFACE_HUB_CACHE beats HF_HOME/hub beats ~/.cache/huggingface/hub.
    """
    override = os.environ.get("HF_HUB_CACHE") or os.environ.get(
        "HUGGINGFACE_HUB_CACHE")
    if override:
        return Path(override)
    home = os.environ.get("HF_HOME")
    if home:
        return Path(home) / "hub"
    return Path.home() / ".cache" / "huggingface" / "hub"


def repo_dir(repo_id: str) -> Path:
    """The cache directory HF keeps REPO_ID's weights in (models--Qwen--…)."""
    return _hf_cache_dir() / ("models--" + repo_id.replace("/", "--"))


def model_repo_dir(model: str) -> Path:
    """The cached-weights directory for a MODEL_REPOS key."""
    return repo_dir(MODEL_REPOS[model])


def _tree_has_file(path: Path) -> bool:
    """True when any file or symlink exists under PATH (recursively)."""
    try:
        for item in path.iterdir():
            # Snapshot files are symlinks into blobs/; count them even when
            # temporarily broken (presence is what the loader checks).
            if item.is_symlink() or item.is_file():
                return True
            if item.is_dir() and _tree_has_file(item):
                return True
    except OSError:
        return False
    return False


def model_installed(model: str) -> bool:
    """True when MODEL's weights look complete in the local HF cache.

    A fetched repo has refs/main plus at least one file under snapshots/;
    anything less counts as not installed. An interrupted download simply
    resumes — via Install, or the next server start for that model.
    """
    directory = model_repo_dir(model)
    if not (directory / "refs" / "main").is_file():
        return False
    return _tree_has_file(directory / "snapshots")


def installed_models() -> List[str]:
    """The MODEL_REPOS keys whose weights are already on disk."""
    return [name for name in MODEL_REPOS if model_installed(name)]


def delete_model_weights(models: Optional[List[str]] = None) -> int:
    """Delete the cached HF weight dirs of MODELS (every model by default).

    Best-effort rmtree of each ``models--Qwen--…`` directory; returns how
    many were present and removed. Only those directories are ever touched —
    the rest of the HF cache may be shared with unrelated tools. Prints its
    progress, which streams into the task view under curses too.
    """
    names = sorted(MODEL_REPOS) if models is None else list(models)
    removed = 0
    for name in names:
        directory = model_repo_dir(name)
        if not directory.is_dir():
            continue
        print(f"[INFO] Removing cached {MODEL_REPOS[name]} weights...")
        shutil.rmtree(directory)
        if directory.exists():
            print(f"[WARNING] Could not fully remove {directory}")
            continue
        removed += 1
    if removed:
        print(f"[OK] Deleted cached weights for {removed} "
              f"{'model' if removed == 1 else 'models'}.")
    return removed


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


def model_for_identity(identity: Optional[str]) -> Optional[str]:
    """The model name a qwen demo answers as (None when not a known identity)."""
    return IDENTITY_TO_MODEL.get(identity)


def model_for_voice_mode(voice_mode: str) -> str:
    """The MODEL_REPOS key a run with VOICE_MODE needs hosted.

    The mirror of ``model_for_identity`` for a planned (not yet running)
    run: instructions design the voice (VoiceDesign), a reference .wav
    clones (Base), otherwise built-in speakers (CustomVoice).
    """
    return {VOICE_MODE_DESIGN: "VoiceDesign",
            VOICE_MODE_CLONE: "Base",
            VOICE_MODE_CUSTOM: "CustomVoice"}[voice_mode]


def desired_identity(model: str) -> str:
    """The probe identity the model's demo answers as (used while booting)."""
    return {
        "CustomVoice": probe.IDENTITY_QWEN_CUSTOM,
        "Base": probe.IDENTITY_QWEN_CLONE,
        "VoiceDesign": probe.IDENTITY_QWEN_DESIGN,
    }[model]


def _wizard(stdscr, args: argparse.Namespace) -> dict:
    """Collect the setup settings without asking anything.

    The qwen backend has no per-install choices: install happens when the
    package is missing (and not skipped by flag), and every other value —
    port — lives in app/converter/config.py (the hub's Settings screen);
    the model and voice are picked per run on the Generate-audiobooks
    screen.
    """
    return {
        "do_install": (not _is_installed()) and not args.skip_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. The pip install streams
    through EMIT and aborts on CANCEL; the list is empty (no-op) when
    there is nothing to install.
    """
    steps: List[taskview.TaskStep] = []
    if not settings["do_install"]:
        return steps

    def install(emit, cancel):
        rc = common.pip_install([QWEN_PIP_PKG], emit=emit, cancel=cancel,
                                env_dir=QWEN_ENV)
        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))
    return steps


def _execute(settings: dict) -> int:
    """Console tail: pip install (no-op when already installed)."""
    return taskview.run_steps_inline(_execute_steps(settings))


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

    There are no questions: settings are computed up front and the install
    runs inside the TUI task view on this same screen — skipped entirely
    (a silent no-op) when nothing needs installing. Returns 0 always —
    the flow cannot be aborted, so Esc/Ctrl-C never short-circuits it.
    """
    args = build_parser().parse_args([])
    settings = _wizard(stdscr, args)
    if not settings["do_install"]:
        return 0
    return taskview.run_steps(stdscr, "Setting up qwen-tts",
                              _execute_steps(settings))


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


def _collect_from_flags(args: argparse.Namespace,
                        parser: argparse.ArgumentParser) -> dict:
    return {
        "do_install": (not _is_installed()) and not args.skip_install,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Set up the Qwen3-TTS demo backend: pip install "
                    "qwen-tts into its managed venv (app/envs/qwen).")
    parser.add_argument("--skip-install", action="store_true",
                        help="Do not pip install qwen-tts")
    return parser


def build_spec(model: str) -> ServerSpec:
    """The single managed ServerSpec hosting MODEL on the configured port.

    Public because callers outside this module (the hub's run preparation
    and the CLI's managed-server bootstrap) need to boot exactly the model
    their run selected, which can differ from the default one.
    """
    url = config.QWEN_API_URL
    return ServerSpec(
        "qwen", url,
        [str(envs.env_script("qwen-tts-demo", QWEN_ENV)), MODEL_REPOS[model],
         "--ip", "127.0.0.1",
         "--port", str(common.port_of(url, DEFAULT_PORT))],
        identity=desired_identity(model))


def _managed_running_model() -> Optional[str]:
    """The model a locally-managed, up-and-running demo answers as.

    None when no pid file exists (this tool never started that server), the
    process is gone, or the probe cannot identify which model it hosts.
    """
    if servers.pid_for("qwen") is None:
        return None
    if not servers.alive("qwen"):
        return None
    return model_for_identity(probe.identify_server(config.QWEN_API_URL))


def detect() -> BackendStatus:
    """Detect whether qwen-tts is installed, plus the launch command.

    One managed spec exists, hosting the default model on the single
    configured port (Generate-audiobooks runs needing another model boot
    it via their own spec). Which model currently answers there is read
    via the probe (local pid alive => check our own URL; otherwise the
    remote URL) so the status names the *running* model even when it
    differs from the default one.
    """
    installed = _is_installed()
    model = DEFAULT_MODEL
    url = config.QWEN_API_URL
    details: List[str] = []
    details.append("pip: installed" if installed else
                   "not installed — run setup to pip install qwen-tts")
    details.append(f"port: {common.port_of(url, DEFAULT_PORT)}")
    details.append(f"default model: {model}")
    specs = [build_spec(model)]
    managed = servers.manages(specs)
    # A locally-managed server names its running model via the probe of the
    # managed URL; a remotely-run demo names it via the remote-URL probe.
    local_models: List[str] = []
    if managed and servers.alive(specs[0].name):
        found = model_for_identity(probe.identify_server(url))
        if found is not None:
            local_models.append(found)
    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 an externally-run qwen demo server at the remote URL.

    Returns ``([model, ...], {spec_name: url})``. The remote URL must answer
    as one of the three demos (see probe.identify_server); when it equals
    the local URL and this tool started that server, it is ignored (already
    reported "[local]").
    """
    remote_models = []
    remote_urls = {}
    url = (config.QWEN_REMOTE_URL or "").strip()
    if not url:
        return remote_models, remote_urls
    if managed and probe.same_endpoint(url, config.QWEN_API_URL):
        return remote_models, remote_urls
    model = model_for_identity(probe.identify_server(url))
    if model is not None:
        remote_urls["qwen"] = url
        remote_models.append(model)
    return remote_models, remote_urls


def _hf_download_prefix() -> Optional[List[str]]:
    """The qwen env's hf CLI argv prefix (None when neither script is present)."""
    for name in ("hf", "huggingface-cli"):
        candidate = envs.env_script(name, QWEN_ENV)
        if candidate.is_file():
            return [str(candidate)]
    return None


def install_model(model: str, *, emit=None, cancel=None) -> int:
    """Pre-download MODEL's weights into the HF cache via the venv's hf CLI.

    Exactly what the first server start does implicitly, made explicit:
    streamed progress through EMIT (percent bars feed the task view),
    CANCEL kills the download process group mid-flight, and re-running
    resumes where a previous attempt left off. Returns the exit code.
    """
    prefix = _hf_download_prefix()
    if prefix is None:
        print("[ERROR] No hf CLI found in the qwen venv; pip install "
              f"{QWEN_PIP_PKG} first")
        return 1
    repo = MODEL_REPOS[model]
    print(f"[INFO] Downloading {repo} into {_hf_cache_dir()}...")
    rc = common.run_console_subprocess(prefix + ["download", repo],
                                       emit=emit, cancel=cancel)
    if rc == 0:
        print(f"[OK] {repo} downloaded.")
    return rc


def uninstall_model(model: str, *, emit=None, cancel=None) -> int:
    """Remove MODEL's cached weights (the inverse of install_model).

    A locally-managed server that currently answers as MODEL is stopped
    first (best-effort) so its weights are not deleted under a live
    process; no server, a down one, or one serving another model leaves
    everything else untouched. CANCEL is honored after that stop phase
    only. Returns the exit code.
    """
    if _managed_running_model() == model:
        servers.stop("qwen")
    if common.cancel_requested(cancel):
        return 130
    delete_model_weights([model])
    return 0


def uninstall(*, emit=None, cancel=None) -> int:
    """Remove the qwen-tts backend entirely: stop its server, pip uninstall,
    then delete every downloaded model.

    qwen-tts is a pip package (``qwen_tts`` + the ``qwen-tts-demo`` script)
    installed into its own managed venv (``app/envs/qwen``, QWEN_ENV), so
    uninstalling it removes the backend — the app venv and the faster
    backend's env are never touched. Any server this tool started is stopped
    first (best-effort). The three models' weight snapshots — multi-GB
    directories lazily fetched into the HuggingFace cache
    (~/.cache/huggingface/hub) — are deleted too, matching the hub's
    confirmation dialog; only those three directories are removed, never the
    shared cache itself.

    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 (stop servers / pip /
    delete models) — a started phase always completes, so pip is never
    killed mid-run. Returns the exit code (130 when cancelled before a
    remaining phase).
    """
    if servers.pid_for("qwen") is not None:
        # 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.
        servers.stop("qwen")
    if common.cancel_requested(cancel):
        return 130
    rc = common.pip_uninstall([QWEN_PIP_PKG], emit=emit, env_dir=QWEN_ENV)
    if rc != 0:
        print(f"[WARNING] pip uninstall failed (exit {rc}); remove "
              f"{QWEN_PIP_PKG} from {QWEN_ENV} manually")
    else:
        print(f"[OK] {QWEN_PIP_PKG} removed.")
    # Weights go even when the pip step failed: the package can be
    # re-installed any time, multi-GB snapshots are what actually cost disk.
    if common.cancel_requested(cancel):
        return 130
    delete_model_weights()
    return rc


def update(*, emit=None, cancel=None) -> int:
    """Update the qwen-tts backend: pip install -U qwen-tts in its venv.

    A managed server that is running is stopped first (best-effort): it
    imports ``qwen_tts`` from the very venv being upgraded, so an in-place
    upgrade under a live process would leave it serving stale code.
    CANCEL is a ``threading.Event`` honored between phases only (stop
    server / pip) — a started phase always completes, so pip is never
    killed mid-run. pip itself is the freshness check: it resolves the
    latest version, upgrades when there is one, and reports "Requirement
    already satisfied" otherwise. Returns the exit code (130 when
    cancelled before a remaining phase).
    """
    if servers.pid_for("qwen") is not None:
        servers.stop("qwen")
    if common.cancel_requested(cancel):
        return 130
    rc = common.pip_install([QWEN_PIP_PKG], emit=emit, cancel=cancel,
                            env_dir=QWEN_ENV, upgrade=True)
    if rc != 0:
        print(f"[WARNING] pip install -U failed (exit {rc}); update "
              f"{QWEN_PIP_PKG} manually")
    else:
        print(f"[OK] {QWEN_PIP_PKG} is up to date (or just upgraded).")
    return rc


def models_screen(stdscr) -> int:
    """Per-model (un)install screen: the hub's Configure-qwen-tts leaf.

    Each of the three models gets exactly one action reflecting disk state:
    Install pre-downloads its weights via the hf CLI (a streamed, resumable,
    cancelable task-view run) and Uninstall deletes them again (stopping a
    managed server that answers as that model first). Install requires the
    pip package; without one a guidance flash replaces the download, since
    pre-fetched weights without a backend to serve them buy nothing. The
    status table and options re-render after every action, so Esc pops back
    to Configure Backends. Always returns 0.
    """
    while True:
        present = installed_models()
        rows = [(name, "installed", "ok") if name in present
                else (name, "not installed", "warn") for name in MODEL_REPOS]
        options = []
        for name in MODEL_REPOS:
            if name in present:
                options.append((f"Uninstall {name}", ("uninstall", name)))
            elif _is_installed():
                options.append((f"Install {name}", ("install", name)))
            else:
                options.append((f"{name} (backend not installed)",
                                ("noop", name)))

        def make_work(action: str, target: str):
            def work(emit, cancel) -> int:
                if action == "install":
                    return install_model(target, emit=emit, cancel=cancel)
                return uninstall_model(target, emit=emit, cancel=cancel)
            return work

        choice = tui.menu(
            stdscr, "Configure qwen-tts", options,
            back_value=tui.Wizard.BACK,
            help_lines=["Models are normally pulled when their server first",
                        "starts; Install pre-downloads one right now."],
            table_title="Model State", table_rows=rows)
        if choice is tui.Wizard.BACK:
            return 0
        action, name = choice
        if action == "noop":
            continue
        if action == "install" and not _is_installed():
            tui.flash(stdscr, "Install the qwen-tts backend first "
                      "(Configure Backends > Install Backend).", "warn")
            continue
        title = (f"Download {MODEL_REPOS[name]}" if action == "install"
                 else f"Delete {name} weights")
        step = taskview.TaskStep(title, make_work(action, name))
        rc = taskview.run_steps(stdscr, title, [step], wait_on_finish=False)
        if rc == 0:
            tui.flash(stdscr,
                      f"{name} downloaded." if action == "install"
                      else f"{name} weights removed.", "ok")
        else:
            verb = "download" if action == "install" else "remove"
            tui.flash(stdscr, f"Could not {verb} {name}.", "err")


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