#!/usr/bin/env python3
"""
Auto-generated setup script for AuK (tencent/AuK) via its official Gradio UI
Generated by Model Lab.

Installs: the official Tencent-Hunyuan/AuK repo with the [gradio] extra
(pip install -e ".[gradio]"), then downloads the AuK Base + AuK-Flash
checkpoints and the Qwen2.5-Omni-3B text encoder AuK depends on for
understanding the instruction/text side of every request. Then serves
auk-gradio on :7860, Tencent's own web UI: a plain web form (text in,
reference audio in, instruction in, audio out) with no node graph to wire up.

RunPod usage — no other setup needed on the pod:
  1. Pick a pod template with Python 3 + CUDA (e.g. "RunPod PyTorch") and a
     GPU with BF16 support. Per Tencent's own README (measured on an
     NVIDIA A800-SXM4-80GB, bf16, current as of PR #19 merged 2026-09-16):
     peak VRAM is ~18GB without --cpu_offload, ~10.5GB with it enabled.
  2. Set container disk to >=60GB (checkpoints total ~25.5GB with both AuK
     variants installed, plus the cloned repo, pip cache, and Python
     environment).
  3. Put this file at /workspace/setup.py (upload it, or paste as the pod's
     on-start script).
  4. Container start command:
       python /workspace/setup.py --install && python /workspace/setup.py --serve
     Add --cpu_offload to the --serve call on a GPU under 20GB VRAM.
  5. Expose/forward TCP port 7860 (RunPod: add an HTTP port mapping to 7860).
     auk-gradio binds 0.0.0.0:7860 by default and needs no reverse-proxy CORS
     flag behind RunPod's proxy.

License note: AuK's own weights are MIT, but the bundled Qwen2.5-Omni-3B
text encoder ships under Alibaba's Qwen Research license, not MIT — check
that license before any commercial use of the encoder weights.

This script is self-contained and self-healing:
  - Each step writes a heartbeat to STATE_FILE so an external supervisor
    can tell whether it is still progressing vs. wedged.
  - Downloads verify HTTP status + expected size; failures are loud and
    resumable via curl --continue-at, so a transient error retries instead
    of silently writing an error page.
  - Any hard failure exits non-zero with a clear message; it never hangs.
"""

import argparse
import http.server
import json
import os
import re
import shutil
import subprocess
import sys
import threading
import time
import urllib.request
import urllib.error
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Optional, Tuple

_STATE_LOCK = threading.Lock()

MIN_PYTHON = (3, 10)  # AuK's own pyproject pins Python >=3.10.
_REEXEC_GUARD_ENV = "MODELLAB_SETUP_REEXECED"


def _print_and_state_bootstrap(msg: str) -> None:
    """Bare print+state usable before _state()'s real definition below runs
    (this fires at import time, ahead of the rest of the module body)."""
    print(f"[SETUP] {msg}", flush=True)
    try:
        with open("/workspace/setup_state.txt", "a", encoding="utf-8", newline="\n") as f:
            f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}\n")
    except Exception:
        pass


def _ensure_modern_python() -> None:
    """If running under Python <MIN_PYTHON, find or apt-install a newer
    interpreter and re-exec this same script under it. Must run before
    anything else — every later step (pip installs, auk-gradio itself)
    assumes sys.executable is modern enough, so this can't be deferred into
    install().
    """
    if sys.version_info[:2] >= MIN_PYTHON:
        return

    current = f"{sys.version_info[0]}.{sys.version_info[1]}"
    min_str = f"{MIN_PYTHON[0]}.{MIN_PYTHON[1]}"

    if os.environ.get(_REEXEC_GUARD_ENV):
        raise RuntimeError(
            f"Re-exec already attempted but still running Python {current} "
            f"(need >={min_str}). Automatic Python upgrade failed on this "
            f"base image — install Python {min_str}+ manually and re-run "
            f"this script with that interpreter."
        )

    _print_and_state_bootstrap(
        f"PYTHON_CHECK {current} < required {min_str}, looking for a newer interpreter"
    )

    candidates = ["python3.13", "python3.12", "python3.11", "python3.10"]
    newer_python = None
    for name in candidates:
        path = shutil.which(name)
        if not path:
            continue
        try:
            out = subprocess.run([path, "-c", "import sys; print('%d.%d' % sys.version_info[:2])"],
                                  capture_output=True, text=True, timeout=15)
            if out.returncode != 0:
                continue
            major_str, _, minor_str = out.stdout.strip().partition(".")
            if (int(major_str), int(minor_str)) >= MIN_PYTHON:
                newer_python = path
                break
        except Exception:
            continue

    if not newer_python:
        _print_and_state_bootstrap("PYTHON_BOOTSTRAP no newer interpreter found, apt-installing one")
        try:
            subprocess.run("apt-get update -qq", shell=True, check=True)
        except Exception as e:
            raise RuntimeError(f"apt-get update failed while bootstrapping Python: {e}")

        for pkg in ["python3.11", "python3.10", "python3.12"]:
            installed = False
            for extras in (f"{pkg}-venv {pkg}-distutils", f"{pkg}-venv", ""):
                cmd = f"apt-get install -y -qq {pkg} {extras}".strip()
                try:
                    subprocess.run(cmd, shell=True, check=True)
                    installed = True
                    break
                except Exception as e:
                    _print_and_state_bootstrap(f"WARN apt-get install '{cmd}' failed: {e}")
            if not installed:
                continue
            path = shutil.which(pkg) or f"/usr/bin/{pkg}"
            if os.path.exists(path):
                newer_python = path
                break

    if not newer_python:
        raise RuntimeError(
            f"Could not find or install a Python >={min_str} interpreter on this "
            f"pod (apt has none of python3.10/3.11/3.12 available). Rebuild "
            f"the pod from a template with a newer Python preinstalled."
        )

    _print_and_state_bootstrap(f"PYTHON_REEXEC switching to {newer_python}")
    os.environ[_REEXEC_GUARD_ENV] = "1"
    os.execv(newer_python, [newer_python] + sys.argv)


_ensure_modern_python()


# Every asset lives directly under WORKDIR/ckpts/<repo-name>/ — the exact
# layout Tencent's own README documents for `hf download ... --local-dir`
# (a full repo snapshot, not a single-file pick), since auk-gradio itself
# looks for auk_base.safetensors / vae.safetensors / the Qwen encoder at
# those repo-shaped paths, not a flat models/ tree the way ComfyUI does.
MANIFEST = {
  "model_ref": "tencent/AuK",
  "slug": "auk-gradio",
  "modality": "audio",
  "archetype": "speech_edit_tts_gradio",
  "container_disk_in_gb": 60,
  "gpu_preferences": [
    "NVIDIA RTX A5000",
    "NVIDIA A40",
    "NVIDIA RTX A6000",
    "NVIDIA L40S"
  ],
  "repos": [
    {"hf_repo": "tencent/AuK", "local_dir": "ckpts/AuK", "required": True},
    {"hf_repo": "tencent/AuK-Flash", "local_dir": "ckpts/AuK-Flash", "required": False},
    {"hf_repo": "Qwen/Qwen2.5-Omni-3B", "local_dir": "ckpts/Qwen2.5-Omni-3B", "required": True},
  ],
}


WORKDIR = "/workspace/AuK"
STATE_FILE = "/workspace/setup_state.txt"
HEARTBEAT_FILE = "/workspace/setup_heartbeat.txt"


def _state(msg: str) -> None:
    """Write a durable progress marker for debuggability / supervision."""
    with _STATE_LOCK:
        try:
            with open(STATE_FILE, "a", encoding="utf-8", newline="\n") as f:
                f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}\n")
        except Exception:
            pass
        _heartbeat()
        print(f"[SETUP] {msg}", flush=True)


def _heartbeat() -> None:
    try:
        with open(HEARTBEAT_FILE, "w", encoding="utf-8") as f:
            f.write(str(int(time.time())))
    except Exception:
        pass


_TOKEN_RE = re.compile(r"(token=)[^&\s\"']+")


def _redact(cmd: str) -> str:
    """Mask HF tokens embedded in commands before they hit stdout or
    STATE_FILE — RunPod's console log viewer is plaintext and world-visible
    to anyone with pod access."""
    return _TOKEN_RE.sub(r"\1***REDACTED***", cmd)


def run_command(cmd: str, cwd: Optional[str] = None, env: Optional[dict] = None) -> None:
    safe_cmd = _redact(cmd)
    print(f"\n$ {safe_cmd}")
    _state(f"RUN {safe_cmd}")
    full_env = {**os.environ, **env} if env else None
    subprocess.run(cmd, cwd=cwd, check=True, shell=True, env=full_env)


def _ensure_pip() -> None:
    """Some minimal RunPod templates ship a python3 with no pip module and no
    ensurepip at all. Fall back to get-pip.py, version-pinned to the running
    interpreter."""
    try:
        subprocess.run([sys.executable, "-m", "pip", "--version"], check=True, capture_output=True)
        return
    except Exception:
        pass
    _state("BOOTSTRAP pip missing, installing via ensurepip")
    try:
        run_command(f"{sys.executable} -m ensurepip --upgrade")
        return
    except Exception as e:
        _state(f"WARN ensurepip failed: {e}")
    py_major, py_minor = sys.version_info[0], sys.version_info[1]
    get_pip_url = f"https://bootstrap.pypa.io/pip/{py_major}.{py_minor}/get-pip.py"
    get_pip_path = f"/tmp/get-pip-{py_major}.{py_minor}.py"
    if os.path.exists(get_pip_path):
        os.remove(get_pip_path)
    req = urllib.request.Request(get_pip_url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req, timeout=120) as resp, open(get_pip_path, "wb") as f:
        shutil.copyfileobj(resp, f)
    run_command(f"{sys.executable} {get_pip_path}")


def _ensure_hf_cli() -> None:
    """`hf download` (the huggingface_hub CLI) is what Tencent's own README
    uses to pull full repo snapshots into ckpts/<repo>/ — installed here
    rather than assumed present on the base image."""
    try:
        subprocess.run(["hf", "--version"], check=True, capture_output=True)
        return
    except Exception:
        pass
    _state("PIP install huggingface_hub[cli] for `hf download`")
    run_command(f'{sys.executable} -m pip install -U "huggingface_hub[cli]"')


def _ensure_hf_transfer() -> None:
    """RunPod's PyTorch template pre-sets HF_HUB_ENABLE_HF_TRANSFER=1, which
    makes `hf download` require the hf_transfer accelerator package — not
    installed by default, so every download fails outright (a hard error,
    not a fallback to the slower path) until this installs it."""
    if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "") not in ("1", "true", "True"):
        return
    try:
        subprocess.run(
            [sys.executable, "-c", "import hf_transfer"],
            check=True, capture_output=True,
        )
        return
    except Exception:
        pass
    _state("PIP install hf_transfer (HF_HUB_ENABLE_HF_TRANSFER=1 is set on this pod)")
    run_command(f'{sys.executable} -m pip install -U hf_transfer')


def _download_repo(repo: dict, hf_token: str) -> None:
    """Download a full HF repo snapshot with `hf download`, retrying via a
    plain re-invocation (the CLI's own resume logic handles partial files on
    retry, unlike the single-file aria2/curl approach the ComfyUI script
    uses for individual .safetensors)."""
    hf_repo = repo["hf_repo"]
    local_dir = os.path.join(WORKDIR, repo["local_dir"])
    marker = os.path.join(local_dir, ".modellab_download_complete")

    if os.path.exists(marker):
        _state(f"OK already present: {hf_repo}")
        return

    os.makedirs(local_dir, exist_ok=True)
    env = {"HF_TOKEN": hf_token} if hf_token else None

    attempts = 0
    last_error: Optional[Exception] = None
    while attempts < 3:
        attempts += 1
        try:
            run_command(f'hf download {hf_repo} --local-dir "{local_dir}"', env=env)
            with open(marker, "w", encoding="utf-8") as f:
                f.write(time.strftime("%Y-%m-%d %H:%M:%S"))
            _state(f"OK downloaded: {hf_repo}")
            return
        except Exception as e:
            last_error = e
            _state(f"WARN hf download attempt {attempts} failed for {hf_repo}: {e}")

    raise RuntimeError(f"Download failed after retries: {hf_repo} ({last_error})")


def download_assets() -> None:
    _state("DOWNLOAD_PHASE_START")
    hf_token = os.environ.get("HF_TOKEN", "")
    repos = MANIFEST.get("repos", [])

    # Two at a time: `hf download` already parallelizes file transfers within
    # one repo, so running all repos fully concurrently would stack that
    # multiple times over on one pod's bandwidth for no real gain, the same
    # reasoning setup_auk_comfyui.py applies to its own asset downloads.
    workers = min(2, len(repos))
    _state(f"DOWNLOADING {len(repos)} repos, {workers} at a time")
    with ThreadPoolExecutor(max_workers=workers) as ex:
        futures = {ex.submit(_download_repo, r, hf_token): r for r in repos}
        for fut in as_completed(futures):
            repo = futures[fut]
            try:
                fut.result()
            except Exception as e:
                if repo.get("required", True):
                    raise
                _state(f"WARN optional repo {repo['hf_repo']} failed, continuing: {e}")

    _state("DOWNLOAD_PHASE_DONE")


def _detect_gpu() -> Tuple[bool, str]:
    """Detect whether this pod actually has a usable NVIDIA GPU attached.

    RunPod's "CPU" pod types have no GPU at all, but a "GPU" pod can still
    come up without one visible yet (driver still initializing) or with a
    misconfigured template, so this checks the real signal (nvidia-smi
    successfully listing a device) rather than trusting the pod type label.
    """
    if not shutil.which("nvidia-smi"):
        return False, "nvidia-smi not found on PATH (no NVIDIA driver/toolkit present)"
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=30,
        )
    except Exception as e:
        return False, f"nvidia-smi failed to run: {e}"
    if result.returncode != 0:
        return False, f"nvidia-smi exited {result.returncode}: {result.stderr.strip()[:200]}"
    gpu_names = [line.strip() for line in result.stdout.splitlines() if line.strip()]
    if not gpu_names:
        return False, "nvidia-smi ran but reported no GPUs"
    return True, ", ".join(gpu_names)


def _get_compute_capability() -> Optional[float]:
    """Read the attached GPU's CUDA compute capability (e.g. 12.0 for
    Blackwell) via nvidia-smi. Returns None if it can't be determined
    (no GPU yet, driver still initializing, or an unparseable value)."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=30,
        )
    except Exception:
        return None
    if result.returncode != 0:
        return None
    lines = [line.strip() for line in result.stdout.splitlines() if line.strip()]
    if not lines:
        return None
    try:
        return float(lines[0])
    except ValueError:
        return None


def _ensure_blackwell_torch() -> None:
    """AuK's own pyproject.toml pins torch>=2.7,<2.8 and its comment says
    users may preinstall a platform-specific CUDA build first so pip sees
    the pin already satisfied and doesn't touch it. That matters on
    Blackwell (RTX 50-series / RTX PRO Blackwell, compute capability
    >=10.0): the default PyPI torch wheels for the 2.7 line have no sm_100+
    kernels, so AuK's plain `pip install` either installs or downgrades to
    a build that detects the GPU (nvidia-smi, has_gpu=True) but fails every
    real op with "CUDA error: no kernel image is available for execution on
    the device." Preinstalling the cu128 wheel index's torch==2.7.0 here,
    before AuK's own install runs, keeps a Blackwell-capable build in place
    since it already satisfies AuK's <2.8 pin and pip won't replace it.
    """
    cap = _get_compute_capability()
    if cap is None or cap < 10.0:
        return
    _state(f"BLACKWELL_GPU detected (compute capability {cap}); preinstalling a cu128 torch build")
    run_command(
        f'{sys.executable} -m pip install "torch==2.7.0" "torchaudio==2.7.0" "torchvision==0.22.0" '
        f'--index-url https://download.pytorch.org/whl/cu128'
    )


_PE_CHECKBOX_OLD = 'use_pe = gr.Checkbox(\n                        value=True,'
_PE_CHECKBOX_NEW = 'use_pe = gr.Checkbox(\n                        value=False,'


def _disable_prompt_enhancer_default() -> None:
    """auk-gradio's "Use Prompt Enhancer" checkbox defaults to checked
    (infer_gradio.py, use_pe = gr.Checkbox(value=True, ...)), and PE needs
    an OpenAI-compatible LLM (LLM_API_KEY / LLM_BASE_URL / LLM_MODEL_NAME)
    this setup script never configures — Tencent Cloud TokenHub by default,
    per AuK's own README. There's no CLI flag or env var to change the
    checkbox's default from outside, so a first-time user hits "Prompt
    Enhancer failed: missing OpenAI-compatible LLM config" on their very
    first Generate click unless they know to uncheck it first. This patches
    that one default in the cloned repo before AuK installs, so the box
    starts unchecked; anyone who does set up an LLM can still check it
    manually per generation.
    """
    target = os.path.join(WORKDIR, "src", "auk", "infer", "infer_gradio.py")
    if not os.path.isfile(target):
        _state(f"WARN skip PE-default patch, file not found: {target}")
        return
    with open(target, "r", encoding="utf-8") as f:
        content = f.read()
    if _PE_CHECKBOX_NEW in content:
        _state("OK Prompt Enhancer default already patched")
        return
    if _PE_CHECKBOX_OLD not in content:
        _state("WARN skip PE-default patch, expected checkbox source not found (upstream may have changed)")
        return
    content = content.replace(_PE_CHECKBOX_OLD, _PE_CHECKBOX_NEW, 1)
    with open(target, "w", encoding="utf-8") as f:
        f.write(content)
    _state("PATCHED Prompt Enhancer checkbox now defaults to unchecked")


def install() -> bool:
    """Runs the full install (clone, deps, model downloads) regardless of
    hardware — none of that is GPU-bound, so a CPU-only pod can pre-stage
    everything before switching to a GPU pod.

    Returns True if a GPU was detected (safe to `serve()` next), False if
    this pod is CPU-only.
    """
    try:
        _state("PHASE install start")
        _ensure_pip()
        _ensure_hf_cli()
        _ensure_hf_transfer()
        _ensure_blackwell_torch()

        if not os.path.isdir(WORKDIR):
            _state("CLONE AuK")
            run_command(f'git clone https://github.com/Tencent-Hunyuan/AuK.git "{WORKDIR}"')

        _disable_prompt_enhancer_default()

        _state("PIP install AuK with the gradio extra")
        run_command(f'{sys.executable} -m pip install -e ".[gradio]"', cwd=WORKDIR)

        # qwen-omni-utils (an AuK dependency) imports audioread directly in
        # its audio-processing path, but only declares librosa, which treats
        # audioread as an optional fallback backend and doesn't always pull
        # it in. Without this, generation fails at runtime with
        # "ModuleNotFoundError: No module named 'audioread'" on the first
        # real request, well after install reports success.
        _state("PIP install audioread (qwen-omni-utils runtime dependency AuK doesn't declare)")
        run_command(f"{sys.executable} -m pip install -U audioread")

        download_assets()

        has_gpu, gpu_detail = _detect_gpu()
        if has_gpu:
            _state(f"GPU_CHECK ok: {gpu_detail}")
        else:
            _state(f"GPU_CHECK none: {gpu_detail}")

        _state("PHASE install done")
        return has_gpu
    except Exception:
        _state("PHASE install FAILED")
        raise


def serve(host: str, port: int, cpu_offload: bool = False) -> None:
    """auk-gradio binds 0.0.0.0:7860 by default and needs no reverse-proxy
    CORS flag the way ComfyUI does behind RunPod's proxy (ComfyUI issue
    #4865) — Gradio's own dev server doesn't apply the same Origin/Host
    validation ComfyUI's does.

    cpu_offload forwards straight to auk-gradio's own --cpu_offload flag,
    which moves the text encoder and DiT to system RAM when idle. Per
    Tencent's README (measured on an A800-SXM4-80GB, bf16, current as of
    PR #19 merged 2026-09-16), this trades inference speed for roughly
    7-8GB less resident VRAM: ~18GB peak without it, ~10.5GB with it.
    """
    _state("PHASE serve")
    _free_port_if_stale(port)
    base_ckpt = os.path.join(WORKDIR, "ckpts", "AuK", "auk_base.safetensors")
    flash_ckpt = os.path.join(WORKDIR, "ckpts", "AuK-Flash", "auk_flash.safetensors")
    args = [
        "auk-gradio",
        "--host", host,
        "--port", str(port),
        "--base_ckpt", base_ckpt,
    ]
    if os.path.exists(flash_ckpt):
        args += ["--flash_ckpt", flash_ckpt]
    if cpu_offload:
        args += ["--cpu_offload"]
    print(f"[SETUP] Launching auk-gradio on {host}:{port} (cpu_offload={cpu_offload})...")
    try:
        os.chdir(WORKDIR)
    except Exception as e:
        _state(f"FATAL chdir {WORKDIR} failed: {e}")
        raise
    os.execvp("auk-gradio", args)


def _serve_cpu_only_notice(host: str, port: int, gpu_detail: str) -> None:
    """Setup succeeded (weights downloaded) but this pod has no GPU, so
    launching auk-gradio would be unusably slow for a multi-GB bf16 model.
    Stay up and serve a clear instruction instead of exiting — RunPod
    restarts an exited container, which would silently re-run the whole
    install loop with nothing to show for it.
    """
    script_name = os.path.basename(sys.argv[0]) or "setup.py"
    message = (
        "Setup complete: AuK + its checkpoints are downloaded and ready.\n"
        "This pod has no GPU attached "
        f"({gpu_detail}), so auk-gradio was not started.\n\n"
        "Next step: stop this pod and switch to a GPU pod type (or resize the "
        "existing pod to a GPU tier), keeping the same network volume/disk so "
        "the downloaded weights carry over. Then run:\n"
        f"  python {script_name} --serve\n"
    )

    class Handler(http.server.BaseHTTPRequestHandler):
        def do_GET(self):
            if self.path.startswith("/setup_state"):
                try:
                    with open(STATE_FILE, "r", encoding="utf-8") as f:
                        body = f.read().encode("utf-8")
                except Exception:
                    body = b"(no state file)"
                self.send_response(200)
                self.send_header("Content-Type", "text/plain")
                self.end_headers()
                self.wfile.write(body)
            else:
                self.send_response(200)
                self.send_header("Content-Type", "text/plain")
                self.end_headers()
                self.wfile.write(message.encode("utf-8"))

        def log_message(self, format, *args):
            pass

    class ReusableHTTPServer(http.server.HTTPServer):
        allow_reuse_address = True

    print(f"\n{'=' * 70}\n{message}{'=' * 70}\n", flush=True)
    _free_port_if_stale(port)
    _state(f"CPU_ONLY_NOTICE binding {host}:{port}, awaiting GPU pod")
    try:
        ReusableHTTPServer((host, port), Handler).serve_forever()
    except OSError as e:
        _state(f"FATAL cpu-only notice server could not bind {host}:{port}: {e}")
        raise


def _serve_error_diagnostic(host: str, port: int, error_text: str) -> None:
    """If install() fails, don't let the process exit — RunPod restarts the
    container when its main process exits, which turns one real error into an
    endless, opaque crash-loop. Instead, bind the gradio port ourselves and
    serve the failure so it's inspectable remotely (GET /setup_state for the
    full log) until the external supervisor's inactivity timeout tears the
    pod down.
    """
    class Handler(http.server.BaseHTTPRequestHandler):
        def do_GET(self):
            if self.path.startswith("/setup_state"):
                try:
                    with open(STATE_FILE, "r", encoding="utf-8") as f:
                        body = f.read().encode("utf-8")
                except Exception:
                    body = b"(no state file)"
                self.send_response(200)
                self.send_header("Content-Type", "text/plain")
                self.end_headers()
                self.wfile.write(body)
            else:
                self.send_response(503)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(json.dumps({"error": "install failed", "detail": error_text}).encode("utf-8"))

        def log_message(self, format, *args):
            pass

    class ReusableHTTPServer(http.server.HTTPServer):
        allow_reuse_address = True

    _free_port_if_stale(port)
    _state(f"DIAGNOSTIC_SERVER binding {host}:{port} after install failure")
    try:
        ReusableHTTPServer((host, port), Handler).serve_forever()
    except OSError as e:
        _state(f"FATAL diagnostic server could not bind {host}:{port}: {e}")
        raise


def _free_port_if_stale(port: int) -> None:
    """Kill any process already bound to `port` before we try to bind it.

    A prior invocation of this script can leave its diagnostic server or
    auk-gradio itself still running and holding the port. Best-effort: tries
    fuser then lsof, and never raises.
    """
    pid_str = ""
    if shutil.which("fuser"):
        result = subprocess.run(
            ["fuser", f"{port}/tcp"], capture_output=True, text=True
        )
        pid_str = result.stdout.strip()
    elif shutil.which("lsof"):
        result = subprocess.run(
            ["lsof", "-t", f"-i:{port}"], capture_output=True, text=True
        )
        pid_str = result.stdout.strip().replace("\n", " ")

    if not pid_str:
        return

    my_pid = os.getpid()
    pids = [p for p in pid_str.split() if p.isdigit() and int(p) != my_pid]
    if not pids:
        return

    _state(f"PORT_CLEANUP killing stale process(es) on port {port}: {' '.join(pids)}")
    for pid in pids:
        try:
            os.kill(int(pid), 9)
        except Exception as e:
            _state(f"WARN could not kill pid {pid}: {e}")
    time.sleep(1)  # give the OS a moment to release the socket


def main() -> None:
    parser = argparse.ArgumentParser(description="Model Lab Setup Script")
    parser.add_argument("--install", action="store_true", help="Run full installation")
    parser.add_argument("--serve", action="store_true", help="Launch auk-gradio")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument(
        "--cpu_offload", action="store_true",
        help="Forward auk-gradio's own --cpu_offload flag: moves the text "
             "encoder and DiT to system RAM when idle, trading inference "
             "speed for roughly 7-8GB less resident VRAM (~18GB peak "
             "without it, ~10.5GB with it, per Tencent's own README).",
    )
    args = parser.parse_args()

    _state(f"ARGS install={args.install} serve={args.serve} cpu_offload={args.cpu_offload}")

    do_install = args.install or (not args.install and not args.serve)
    do_serve = args.serve or (not args.install and not args.serve)

    if do_install:
        try:
            install()
        except Exception as e:
            _state(f"PHASE install FAILED, staying alive for diagnosis: {e}")
            _serve_error_diagnostic(args.host, args.port, str(e))
            return  # unreachable (serve_forever blocks); kept for clarity

    if do_serve:
        has_gpu, gpu_detail = _detect_gpu()
        if not has_gpu:
            _serve_cpu_only_notice(args.host, args.port, gpu_detail)
            return  # unreachable (serve_forever blocks); kept for clarity
        serve(args.host, args.port, cpu_offload=args.cpu_offload)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        _state(f"FATAL {e}")
        print(f"[SETUP] FATAL: {e}", file=sys.stderr)
        sys.exit(1)