launcher.py

from __future__ import annotations

import argparse
import json
import os
import shutil
import signal
import socket
import subprocess
import sys
import threading
import time
import webbrowser
from collections import deque
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
from urllib.request import Request, urlopen


ROOT = Path(__file__).resolve().parent
UI_ROOT = ROOT / "launcher_ui"
CONFIG_PATH = ROOT / "launcher_config.json"

DEFAULT_CONFIG: dict[str, Any] = {
    "launcher_host": "127.0.0.1",
    "launcher_port": 7859,
    "webui_host": "127.0.0.1",
    "webui_port": 7860,
    "model_dir": "checkpoints",
    "version": "2.5",
    "fp16": True,
}

V25_REQUIRED_FILES = (
    "config.yaml",
    "gpt.pth",
    "s2mel.pth",
    "wav2vec2bert_stats.pt",
    "codec.pth",
    "feat1.pt",
    "feat2.pt",
    "multilingual_zh_ja_yue_char_del.tiktoken",
)
V25_REQUIRED_DIRS = (
    "qwen0.6bemo4-merge",
    "hf_cache/w2v-bert-2.0",
    "hf_cache/bigvgan",
)


def load_config(path: Path = CONFIG_PATH) -> dict[str, Any]:
    config = dict(DEFAULT_CONFIG)
    if path.is_file():
        try:
            loaded = json.loads(path.read_text(encoding="utf-8"))
            if isinstance(loaded, dict):
                config.update(loaded)
        except (OSError, ValueError):
            pass
    return config


def port_is_open(host: str, port: int, timeout: float = 0.25) -> bool:
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def collect_model_status(model_dir: Path) -> dict[str, Any]:
    missing_files = [name for name in V25_REQUIRED_FILES if not (model_dir / name).is_file()]
    missing_dirs = [name for name in V25_REQUIRED_DIRS if not (model_dir / name).is_dir()]
    return {
        "ready": not missing_files and not missing_dirs,
        "path": str(model_dir),
        "missing": missing_files + missing_dirs,
        "checked": len(V25_REQUIRED_FILES) + len(V25_REQUIRED_DIRS),
    }


def module_available(python_exe: Path, module: str) -> bool:
    if not python_exe.is_file():
        return False
    command = [
        str(python_exe),
        "-c",
        f"import importlib.util,sys;sys.exit(0 if importlib.util.find_spec({module!r}) else 1)",
    ]
    try:
        return subprocess.run(
            command,
            cwd=ROOT,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            timeout=10,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        ).returncode == 0
    except (OSError, subprocess.TimeoutExpired):
        return False


def read_python_version(python_exe: Path) -> str:
    if not python_exe.is_file():
        return "未找到"
    try:
        result = subprocess.run(
            [str(python_exe), "-c", "import platform;print(platform.python_version())"],
            cwd=ROOT,
            capture_output=True,
            text=True,
            timeout=10,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        return result.stdout.strip() or "未知"
    except (OSError, subprocess.TimeoutExpired):
        return "未知"


def read_gpu_info() -> list[dict[str, str]]:
    command = [
        "nvidia-smi",
        "--query-gpu=name,memory.total,driver_version",
        "--format=csv,noheader,nounits",
    ]
    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=10,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
    except (OSError, subprocess.TimeoutExpired):
        return []
    if result.returncode != 0:
        return []
    gpus = []
    for line in result.stdout.splitlines():
        parts = [part.strip() for part in line.split(",")]
        if len(parts) >= 3:
            gpus.append({"name": parts[0], "memory_mb": parts[1], "driver": parts[2]})
    return gpus


def read_runtime_gpu(python_exe: Path, detected_gpus: list[dict[str, str]]) -> dict[str, str] | None:
    if not python_exe.is_file():
        return detected_gpus[0] if detected_gpus else None
    command = [
        str(python_exe),
        "-c",
        (
            "import json,torch;"
            "payload={'name':torch.cuda.get_device_name(0),"
            "'memory_mb':str(round(torch.cuda.get_device_properties(0).total_memory/1048576))}"
            " if torch.cuda.is_available() else None;"
            "print(json.dumps(payload) if payload else '')"
        ),
    ]
    try:
        result = subprocess.run(
            command,
            cwd=ROOT,
            capture_output=True,
            text=True,
            timeout=20,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        runtime_gpu = json.loads(result.stdout.strip()) if result.returncode == 0 and result.stdout.strip() else None
    except (OSError, ValueError, subprocess.TimeoutExpired):
        runtime_gpu = None
    if not runtime_gpu:
        return detected_gpus[0] if detected_gpus else None
    matched = next((gpu for gpu in detected_gpus if gpu["name"] == runtime_gpu["name"]), None)
    runtime_gpu["driver"] = matched["driver"] if matched else "未知"
    return runtime_gpu


class LauncherState:
    def __init__(self, config: dict[str, Any]) -> None:
        self.config = config
        self.python_exe = ROOT / ".venv" / "Scripts" / "python.exe"
        self.model_dir = (ROOT / str(config["model_dir"])).resolve()
        self.process: subprocess.Popen[str] | None = None
        self.phase = "idle"
        self.last_error = ""
        self.started_at: float | None = None
        self.logs: deque[dict[str, Any]] = deque(maxlen=500)
        self.log_sequence = 0
        self.lock = threading.RLock()
        self.worker: threading.Thread | None = None
        self.python_version = read_python_version(self.python_exe)
        self.gpus = read_gpu_info()
        self.runtime_gpu = read_runtime_gpu(self.python_exe, self.gpus)
        self._env_checked_at = 0.0
        self._gradio_ready = False
        self.append_log("启动器已就绪，等待工作台启动。", "system")

    @property
    def webui_url(self) -> str:
        return f"http://{self.config['webui_host']}:{int(self.config['webui_port'])}/"

    def append_log(self, message: str, stream: str = "stdout") -> None:
        clean = message.rstrip("\r\n")
        if not clean:
            return
        with self.lock:
            self.log_sequence += 1
            self.logs.append(
                {
                    "id": self.log_sequence,
                    "time": time.strftime("%H:%M:%S"),
                    "stream": stream,
                    "message": clean,
                }
            )

    def _check_environment(self, force: bool = False) -> bool:
        now = time.monotonic()
        if force or now - self._env_checked_at > 8:
            self._gradio_ready = module_available(self.python_exe, "gradio")
            self._env_checked_at = now
        return self._gradio_ready

    def _process_running(self) -> bool:
        return self.process is not None and self.process.poll() is None

    def status(self) -> dict[str, Any]:
        host = str(self.config["webui_host"])
        port = int(self.config["webui_port"])
        listening = port_is_open(host, port)
        owned = self._process_running()
        with self.lock:
            if owned and listening:
                self.phase = "running"
            elif self.process is not None and self.process.poll() is not None:
                if self.phase not in {"idle", "error", "stopped"}:
                    self.phase = "error"
                    self.last_error = f"工作台已退出，代码 {self.process.returncode}。"
            elif listening and not owned:
                self.phase = "external"

            model = collect_model_status(self.model_dir)
            environment_ready = self.python_exe.is_file() and self._check_environment()
            gpu = self.runtime_gpu
            return {
                "phase": self.phase,
                "error": self.last_error,
                "service": {
                    "listening": listening,
                    "owned": owned,
                    "url": self.webui_url,
                    "port": port,
                    "uptime_seconds": int(time.time() - self.started_at) if self.started_at else 0,
                },
                "model": model,
                "environment": {
                    "ready": environment_ready,
                    "python": self.python_version,
                    "gradio": self._gradio_ready,
                    "python_path": str(self.python_exe),
                },
                "gpu": gpu,
                "gpus": self.gpus,
                "version": str(self.config["version"]),
                "fp16": bool(self.config["fp16"]),
                "log_cursor": self.log_sequence,
            }

    def get_logs(self, after: int = 0) -> list[dict[str, Any]]:
        with self.lock:
            return [entry for entry in self.logs if entry["id"] > after]

    def start(self) -> tuple[bool, str]:
        with self.lock:
            if self.worker and self.worker.is_alive():
                return False, "启动任务正在执行。"
            if self._process_running() or port_is_open(
                str(self.config["webui_host"]), int(self.config["webui_port"])
            ):
                return False, "工作台已经在运行。"
            self.worker = threading.Thread(target=self._start_workflow, daemon=True)
            self.worker.start()
        return True, "已提交启动任务。"

    def _run_setup(self) -> bool:
        uv_exe = shutil.which("uv")
        if not uv_exe:
            self.last_error = "环境缺少 Gradio，且未找到 uv。"
            self.append_log(self.last_error, "error")
            return False
        self.phase = "preparing"
        command = [uv_exe, "sync", "--extra", "webui", "--frozen"]
        self.append_log("首次运行：正在补齐 WebUI 依赖。", "system")
        process = subprocess.Popen(
            command,
            cwd=ROOT,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        )
        assert process.stdout is not None
        for line in process.stdout:
            self.append_log(line)
        code = process.wait()
        if code != 0:
            self.last_error = f"WebUI 依赖安装失败，代码 {code}。"
            self.append_log(self.last_error, "error")
            return False
        self._check_environment(force=True)
        return self._gradio_ready

    def _start_workflow(self) -> None:
        model = collect_model_status(self.model_dir)
        if not model["ready"]:
            self.phase = "error"
            self.last_error = "模型文件不完整：" + "、".join(model["missing"])
            self.append_log(self.last_error, "error")
            return
        if not self.python_exe.is_file():
            self.phase = "error"
            self.last_error = "未找到项目 Python 环境：.venv\\Scripts\\python.exe"
            self.append_log(self.last_error, "error")
            return
        if not self._check_environment(force=True) and not self._run_setup():
            self.phase = "error"
            return

        command = [
            str(self.python_exe),
            "-u",
            str(ROOT / "webui.py"),
            "--version",
            str(self.config["version"]),
            "--model_dir",
            str(self.model_dir),
            "--host",
            str(self.config["webui_host"]),
            "--port",
            str(int(self.config["webui_port"])),
        ]
        if self.config.get("fp16", True):
            command.append("--fp16")

        self.phase = "starting"
        self.last_error = ""
        self.append_log("正在加载 IndexTTS 2.5，首次启动需要一些时间。", "system")
        env = os.environ.copy()
        env["PYTHONUNBUFFERED"] = "1"
        env["INDEXTTS_SKIP_EXAMPLE_DOWNLOAD"] = "1"
        try:
            self.process = subprocess.Popen(
                command,
                cwd=ROOT,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                encoding="utf-8",
                errors="replace",
                creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
            )
        except OSError as exc:
            self.phase = "error"
            self.last_error = f"无法启动工作台：{exc}"
            self.append_log(self.last_error, "error")
            return

        self.started_at = time.time()
        threading.Thread(target=self._read_process_output, daemon=True).start()
        deadline = time.monotonic() + 600
        while time.monotonic() < deadline:
            if self.process.poll() is not None:
                self.phase = "error"
                self.last_error = f"工作台启动失败，代码 {self.process.returncode}。"
                self.append_log(self.last_error, "error")
                return
            if port_is_open(str(self.config["webui_host"]), int(self.config["webui_port"])):
                self.phase = "running"
                self.append_log("IndexTTS 2.5 工作台已可用。", "system")
                return
            time.sleep(1)
        self.phase = "error"
        self.last_error = "工作台加载超过 10 分钟，请查看运行日志。"
        self.append_log(self.last_error, "error")

    def _read_process_output(self) -> None:
        process = self.process
        if process is None or process.stdout is None:
            return
        for line in process.stdout:
            self.append_log(line)

    def stop(self) -> tuple[bool, str]:
        with self.lock:
            process = self.process
            if process is None or process.poll() is not None:
                if port_is_open(str(self.config["webui_host"]), int(self.config["webui_port"])):
                    return False, "端口由外部进程占用，启动器不会终止它。"
                self.phase = "idle"
                return False, "工作台未运行。"
            self.phase = "stopping"
            threading.Thread(target=self._stop_workflow, args=(process,), daemon=True).start()
        return True, "正在停止工作台。"

    def _stop_workflow(self, process: subprocess.Popen[str]) -> None:
        self.append_log("正在停止 IndexTTS 2.5 工作台。", "system")
        process.terminate()
        try:
            process.wait(timeout=15)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=5)
        with self.lock:
            if self.process is process:
                self.process = None
            self.phase = "idle"
            self.started_at = None
        self.append_log("工作台已停止。", "system")

    def restart(self) -> tuple[bool, str]:
        with self.lock:
            if self.worker and self.worker.is_alive():
                return False, "当前任务尚未完成。"
            self.worker = threading.Thread(target=self._restart_workflow, daemon=True)
            self.worker.start()
        return True, "正在重启工作台。"

    def _restart_workflow(self) -> None:
        process = self.process
        if process is not None and process.poll() is None:
            self._stop_workflow(process)
        time.sleep(0.5)
        self._start_workflow()


class LauncherHandler(BaseHTTPRequestHandler):
    state: LauncherState
    server_version = "IndexTTSLauncher/1.0"

    def log_message(self, _format: str, *_args: Any) -> None:
        return

    def _send_json(self, payload: dict[str, Any], status: int = 200) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.send_header("X-Content-Type-Options", "nosniff")
        self.end_headers()
        self.wfile.write(body)

    def _send_file(self, path: Path, content_type: str) -> None:
        try:
            body = path.read_bytes()
        except OSError:
            self.send_error(HTTPStatus.NOT_FOUND)
            return
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-cache")
        self.send_header("X-Content-Type-Options", "nosniff")
        workbench_origin = self.state.webui_url.rstrip("/")
        self.send_header(
            "Content-Security-Policy",
            "default-src 'self'; "
            "img-src 'self' data:; "
            "style-src 'self'; "
            "script-src 'self'; "
            "connect-src 'self'; "
            f"frame-src 'self' {workbench_origin}",
        )
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/api/status":
            self._send_json(self.state.status())
            return
        if parsed.path == "/api/logs":
            try:
                after = int(dict(part.split("=", 1) for part in parsed.query.split("&") if "=" in part).get("after", "0"))
            except ValueError:
                after = 0
            self._send_json({"logs": self.state.get_logs(after)})
            return
        if parsed.path == "/api/health":
            self._send_json({"ok": True})
            return

        static_files = {
            "/": (UI_ROOT / "index.html", "text/html; charset=utf-8"),
            "/index.html": (UI_ROOT / "index.html", "text/html; charset=utf-8"),
            "/styles.css": (UI_ROOT / "styles.css", "text/css; charset=utf-8"),
            "/app.js": (UI_ROOT / "app.js", "text/javascript; charset=utf-8"),
            "/index-icon.png": (ROOT / "assets" / "index_icon.png", "image/png"),
            "/favicon.ico": (ROOT / "assets" / "index_icon.png", "image/png"),
        }
        target = static_files.get(parsed.path)
        if target is None:
            self.send_error(HTTPStatus.NOT_FOUND)
            return
        self._send_file(*target)

    def do_POST(self) -> None:
        actions = {
            "/api/start": self.state.start,
            "/api/stop": self.state.stop,
            "/api/restart": self.state.restart,
        }
        action = actions.get(urlparse(self.path).path)
        if action is None:
            self.send_error(HTTPStatus.NOT_FOUND)
            return
        ok, message = action()
        self._send_json({"ok": ok, "message": message}, HTTPStatus.ACCEPTED if ok else HTTPStatus.CONFLICT)


def build_server(config: dict[str, Any], state: LauncherState) -> ThreadingHTTPServer:
    handler = type("ConfiguredLauncherHandler", (LauncherHandler,), {"state": state})
    return ThreadingHTTPServer(
        (str(config["launcher_host"]), int(config["launcher_port"])),
        handler,
    )


def start_existing_launcher(dashboard_url: str) -> None:
    try:
        request = Request(f"{dashboard_url.rstrip('/')}/api/start", method="POST")
        with urlopen(request, timeout=3):
            pass
    except OSError:
        pass


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="IndexTTS 2.5 local launcher")
    parser.add_argument("--autostart", action="store_true", help="Start the WebUI after opening the dashboard")
    parser.add_argument("--no-browser", action="store_true", help="Do not open the default browser")
    args = parser.parse_args(argv)

    os.chdir(ROOT)
    config = load_config()
    dashboard_url = f"http://{config['launcher_host']}:{int(config['launcher_port'])}/"
    state = LauncherState(config)
    try:
        server = build_server(config, state)
    except OSError:
        if args.autostart:
            start_existing_launcher(dashboard_url)
        if not args.no_browser:
            webbrowser.open(dashboard_url)
        return 0

    if args.autostart:
        state.start()
    if not args.no_browser:
        threading.Timer(0.6, lambda: webbrowser.open(dashboard_url)).start()

    def shutdown(_signum: int, _frame: Any) -> None:
        threading.Thread(target=server.shutdown, daemon=True).start()

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)
    try:
        server.serve_forever(poll_interval=0.5)
    finally:
        state.stop()
        server.server_close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())