#!/usr/bin/env python3
"""
Auto-generated setup script for ComfyUI + Krea 2 (Genlovers/Krea-2-tuto)
Generated by Model Lab.

Installs: ComfyUI, ComfyUI-Manager, the Krea 2 turbo fp8 checkpoint, its
text encoder + VAE, and 5 style LoRAs. Then serves ComfyUI on :8188.

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 >=24GB VRAM (checkpoint alone is ~13GB fp8).
  2. Set container disk to >=60GB (models total ~19GB, plus ComfyUI + venv).
  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
  5. Expose/forward TCP port 8188 (RunPod: add an HTTP port mapping to 8188).
  6. Optional: set env var HF_TOKEN if Genlovers/Krea-2-tuto ever becomes
     gated/private (it's public today, so this is not required).
  If pip/ensurepip is missing on the base image, install() bootstraps pip
  automatically via a Python-version-pinned get-pip.py — no manual fix needed.

Python version: current ComfyUI requires Python >=3.9 (its
comfyui-frontend-package dependency publishes no wheel for 3.8, which some
RunPod base images still ship). Before anything else runs, this script
checks its own interpreter version; if it's too old, it looks for an
already-installed newer python3.X, or apt-installs one, then re-execs
itself under it — no manual Python upgrade or template swap needed.

Downloads: at most 2 assets download at once (DOWNLOAD_CONCURRENCY). Each
asset's aria2c run already opens up to 16 connections on its own, so running
all N assets fully concurrently would open up to 16*N connections on one
pod — enough contention that aria2c's own mid-flight retries could rewrite
a destination file out from under the size check right after it looked
complete, surfacing as false "size mismatch: gone" failures on downloads
that actually succeeded. Capping to 2 at a time fixed that while still
overlapping the small LoRAs with the one huge checkpoint download. Every
asset's expected byte count is pulled exactly from the Hub API (not
rounded), so the size check can't reject a genuinely complete file either.

RunPod proxy access: ComfyUI is launched with --enable-cors-header. Without
it, ComfyUI's own request validation sees that RunPod's proxy forwards a
Host header that doesn't match the browser's Origin, treats that as a
forged request, and 403s every browser request to the *.proxy.runpod.net
URL — even though the port is reachable and ComfyUI's own logs show nothing
wrong (ComfyUI issue #4865). This is a different mechanism from Manager's
own security_level setting (see below); both are needed for full proxy
access, and this one blocks reaching the UI at all, not just Manager routes.

ComfyUI-Manager security: newer Manager versions (V3.x+) moved their
config.ini from user/default/ComfyUI-Manager/ to user/__manager/, and their
startup migration actively RAISES security_level back to 'normal' if it
only finds the legacy path — which can further restrict Manager's own
routes. This script writes security_level = weak to BOTH paths so it sticks
regardless of Manager version.

CPU vs GPU pods: install() always runs fully (deps, custom nodes, and every
model download) regardless of hardware, since none of that needs a GPU — so
you can pre-stage everything on a cheap CPU pod first. But if no NVIDIA GPU
is detected when --serve runs, ComfyUI is NOT launched; instead the script
stays up and serves a plain-text message on :8188 telling you to switch to
a GPU pod (same disk/volume) and re-run with --serve. This avoids either a
hard crash or an unusably slow CPU inference attempt.

Re-running this script: every port bind (the CPU-only notice, the failure
diagnostic, and ComfyUI itself) first kills any stale process still holding
that port from an earlier attempt, so re-invoking after a failure or a
Ctrl-C doesn't need a manual `kill` first. Downloads are also idempotent —
a re-run skips any asset already present at its full expected size and only
retries the ones that are missing or incomplete.

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 errno
import http.server
import json
import os
import re
import shutil
import subprocess
import sys
import threading
import time
import urllib.request
import urllib.parse
import urllib.error
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Iterable, Optional, Tuple

_STATE_LOCK = threading.Lock()

MIN_PYTHON = (3, 9)  # comfyui-frontend-package (a ComfyUI requirements.txt
                      # pin) publishes no wheel for Python <3.9; older pods
                      # ("RunPod PyTorch" images have shipped 3.8) need a
                      # newer interpreter installed before anything else runs.
_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, ComfyUI 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):
        # Already tried re-execing once this process tree and we're STILL on
        # an old interpreter — apt-install must have failed. Don't loop
        # forever; surface a clear, actionable error instead.
        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"
    )

    # 1) Maybe a newer interpreter is already on the box (common on Ubuntu/
    #    Debian base images that ship several python3.X alongside the
    #    default python3 symlink) — try the most likely names before apt.
    candidates = ["python3.13", "python3.12", "python3.11", "python3.10", "python3.9"]
    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

    # 2) Not present -> apt-install one. Try newest-first; take whichever
    #    the base distro's default repos actually have (no PPA dependency,
    #    since arbitrary pod images may have no network path to add one).
    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", "python3.9"]:
            # distutils was removed from the stdlib (and its apt package) in
            # 3.12+, so a combined install with "-distutils" 404s on that
            # version; retry with just the interpreter + venv module if the
            # first attempt fails (_ensure_pip() bootstraps pip separately
            # regardless, so distutils itself isn't a hard requirement here).
            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.9/3.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()


MANIFEST = {
  "model_ref": "Genlovers/Krea-2-tuto",
  "slug": "krea-2",
  "modality": "image",
  "archetype": "unet_split_flux",
  "container_disk_in_gb": 60,
  "gpu_preferences": [
    "NVIDIA RTX A4500",
    "NVIDIA RTX A5000",
    "NVIDIA A40",
    "NVIDIA GeForce RTX 3090"
  ],
  "custom_nodes": [
    "Comfy-Org/ComfyUI-Manager"
  ],
  "assets": [
    {
      "url": "https://huggingface.co/Genlovers/Krea-2-tuto/resolve/main/krea2TurboFP8_krea2TURBO.safetensors",
      "target_dir": "diffusion_models",
      "filename": "krea2TurboFP8_krea2TURBO.safetensors",
      "bytes": 12900096996
    },
    {
      "url": "https://huggingface.co/Comfy-Org/Krea-2/resolve/main/text_encoders/qwen3vl_4b_fp8_scaled.safetensors",
      "target_dir": "text_encoders",
      "filename": "qwen3vl_4b_fp8_scaled.safetensors",
      "bytes": 5242467968
    },
    {
      "url": "https://huggingface.co/Comfy-Org/Krea-2/resolve/main/vae/qwen_image_vae.safetensors",
      "target_dir": "vae",
      "filename": "qwen_image_vae.safetensors",
      "bytes": 253806246
    },
    {
      "url": "https://huggingface.co/Genlovers/Krea-2-tuto/resolve/main/KNPV3_1.safetensors",
      "target_dir": "loras",
      "filename": "KNPV3_1.safetensors",
      "bytes": 228587712
    },
    {
      "url": "https://huggingface.co/Genlovers/Krea-2-tuto/resolve/main/bloomgirls-ultrarealism-krea2_4k.safetensors",
      "target_dir": "loras",
      "filename": "bloomgirls-ultrarealism-krea2_4k.safetensors",
      "bytes": 228587800
    },
    {
      "url": "https://huggingface.co/Genlovers/Krea-2-tuto/resolve/main/cutifier_krea2.safetensors",
      "target_dir": "loras",
      "filename": "cutifier_krea2.safetensors",
      "bytes": 228587728
    },
    {
      "url": "https://huggingface.co/Genlovers/Krea-2-tuto/resolve/main/realism_engine_krea2_v3.1.safetensors",
      "target_dir": "loras",
      "filename": "realism_engine_krea2_v3.1.safetensors",
      "bytes": 1562410296
    },
    {
      "url": "https://huggingface.co/Genlovers/Krea-2-tuto/resolve/main/snofs_krea_v1_3D.safetensors",
      "target_dir": "loras",
      "filename": "snofs_krea_v1_3D.safetensors",
      "bytes": 1562410336
    }
  ]
}


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/Civitai tokens embedded in download URLs 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) -> None:
    safe_cmd = _redact(cmd)
    print(f"\n$ {safe_cmd}")
    _state(f"RUN {safe_cmd}")
    subprocess.run(cmd, cwd=cwd, check=True, shell=True)


def _dl_with_curl(url: str, destination: str, expected_bytes: int) -> bool:
    # -C - resumes a partial file; --fail makes HTTP errors non-zero.
    # NOTE: deliberately no --retry-all-errors — that flag needs curl >=7.71
    # and errors out ("is unknown") on the older curl some slim base images
    # ship. --retry already covers transient/connection failures without it.
    try:
        run_command(
            f'curl -L --fail --retry 5 --retry-delay 3 '
            f'-C - -o "{destination}" "{url}"'
        )
        return _verify(destination, expected_bytes)
    except Exception as e:
        _state(f"WARN curl failed for {os.path.basename(destination)}: {e}")
        return False


def _dl_with_wget(url: str, destination: str, expected_bytes: int) -> bool:
    try:
        run_command(f'wget -c -t 5 --waitretry=3 -O "{destination}" "{url}"')
        return _verify(destination, expected_bytes)
    except Exception as e:
        _state(f"WARN wget failed for {os.path.basename(destination)}: {e}")
        return False


def _verify(destination: str, expected_bytes: int) -> bool:
    if not os.path.exists(destination) or os.path.getsize(destination) == 0:
        return False
    # Expected size is a floor; smaller-is-suspect (error page) but allow
    # unverified assets (expected_bytes==0) to pass on non-empty.
    if expected_bytes and os.path.getsize(destination) < expected_bytes:
        # Delete partial/error payload so next attempt restarts clean.
        try:
            os.remove(destination)
        except Exception:
            pass
        _state(f"WARN size mismatch for {os.path.basename(destination)}: "
               f"{os.path.getsize(destination) if os.path.exists(destination) else 'gone'} (want >={expected_bytes})")
        return False
    return True


def _dl_with_urllib(url: str, destination: str, expected_bytes: int) -> bool:
    tmp = destination + ".urltmp"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req, timeout=120) as resp, open(tmp, "wb") as f:
            shutil.copyfileobj(resp, f)
        if _verify(tmp, expected_bytes):
            os.replace(tmp, destination)
            return True
        try:
            os.remove(tmp)
        except Exception:
            pass
        return False
    except urllib.error.HTTPError as e:
        _state(f"WARN urllib HTTP {e.code} for {_redact(url)}")
        return False
    except Exception as e:
        _state(f"WARN urllib failed for {_redact(url)}: {e}")
        return False


def _dl_with_aria2(url: str, destination: str, expected_bytes: int) -> bool:
    # aria2c splits a single URL across N connections (-x) and N segments (-s),
    # which is far faster than curl on the big safetensors files. Prefer it.
    if not shutil.which("aria2c"):
        return False
    try:
        run_command(
            f'aria2c -x16 -s16 -k1M --file-allocation=none --auto-file-renaming=false '
            f'--max-tries=5 --retry-wait=3 --continue=true '
            f'-d "{os.path.dirname(destination)}" -o "{os.path.basename(destination)}" "{url}"'
        )
        return _verify(destination, expected_bytes)
    except Exception as e:
        _state(f"WARN aria2c failed for {os.path.basename(destination)}: {e}")
        # aria2 may leave .aria2 control files; clean them for a clean retry.
        try:
            for suffix in (".aria2",):
                ctl = destination + suffix
                if os.path.exists(ctl):
                    os.remove(ctl)
        except Exception:
            pass
        return False


def download_url_file(url: str, destination: str, expected_bytes: int = 0) -> None:
    os.makedirs(os.path.dirname(destination), exist_ok=True)

    # Already present and valid -> skip (idempotent across pod restarts).
    if _verify(destination, expected_bytes):
        _state(f"OK already present: {os.path.basename(destination)}")
        return

    filename = os.path.basename(destination)
    _state(f"DOWNLOAD {filename} (expected >= {expected_bytes} bytes)")

    attempts = 0
    while attempts < 2:  # one full pass + one recovery pass
        attempts += 1
        # Fastest first.
        if _dl_with_aria2(url, destination, expected_bytes):
            _state(f"OK downloaded: {filename}")
            return
        if shutil.which("curl") and _dl_with_curl(url, destination, expected_bytes):
            _state(f"OK downloaded: {filename}")
            return
        if shutil.which("wget") and _dl_with_wget(url, destination, expected_bytes):
            _state(f"OK downloaded: {filename}")
            return
        if _dl_with_urllib(url, destination, expected_bytes):
            _state(f"OK downloaded: {filename}")
            return
        _state(f"RETRY pass {attempts} for {filename}")

    raise RuntimeError(
        f"Download failed after retries: {filename} from {url}"
    )


def _ensure_aria2() -> None:
    """Install aria2 (multi-connection downloader) if missing. Without it we
    fall back to single-stream curl, which maxes out far slower on big files."""
    if shutil.which("aria2c"):
        return
    _state("APT install aria2 (for fast multi-connection downloads)")
    try:
        run_command("apt-get update -qq && apt-get install -y -qq aria2")
    except Exception as e:
        _state(f"WARN aria2 install failed, will fall back to curl: {e}")


def _download_one(asset: dict, models_root: str, hf_token: str) -> None:
    url = asset["url"]
    expected_bytes = int(asset.get("bytes") or 0)
    if hf_token and "huggingface.co" in url and "token=" not in url:
        separator = "&" if "?" in url else "?"
        url = f"{url}{separator}token={hf_token}"
    target_dir = os.path.join(models_root, asset["target_dir"])
    dest = os.path.join(target_dir, asset["filename"])
    download_url_file(url, dest, expected_bytes)


DOWNLOAD_CONCURRENCY = 2  # max simultaneous asset downloads


def download_assets(comfyui_dir: str) -> None:
    _state("DOWNLOAD_PHASE_START")
    _ensure_aria2()
    models_root = os.path.join(comfyui_dir, "models")
    hf_token = os.environ.get("HF_TOKEN", "")

    assets = MANIFEST.get("assets", [])
    if len(assets) <= 1:
        for asset in assets:
            _download_one(asset, models_root, hf_token)
    else:
        # Each asset's aria2c run opens up to -x16 connections on its own;
        # running all assets fully concurrently (max_workers=len(assets))
        # multiplies that across every file at once (8 files x 16 = 128
        # connections on one pod), which starves bandwidth/CPU badly enough
        # that aria2c's own mid-flight retries can rewrite a destination file
        # out from under the size check right after it looked complete,
        # surfacing as spurious "size mismatch: gone" failures on otherwise-
        # good downloads. Capping simultaneous files keeps total connection
        # count sane while still overlapping the small LoRAs with the one
        # huge checkpoint.
        workers = min(DOWNLOAD_CONCURRENCY, len(assets))
        _state(f"DOWNLOADING {len(assets)} assets, {workers} at a time")
        with ThreadPoolExecutor(max_workers=workers) as ex:
            futures = {ex.submit(_download_one, a, models_root, hf_token): a for a in assets}
            for fut in as_completed(futures):
                asset = futures[fut]
                try:
                    fut.result()
                except Exception as e:
                    _state(f"ERROR asset {asset.get('filename')}: {e}")

    _state("DOWNLOAD_PHASE_DONE")


def setup_custom_nodes(comfyui_dir: str) -> None:
    custom_nodes_dir = os.path.join(comfyui_dir, "custom_nodes")
    os.makedirs(custom_nodes_dir, exist_ok=True)

    for repo in MANIFEST.get("custom_nodes", []):
        repo_name = repo.split("/")[-1]
        repo_path = os.path.join(custom_nodes_dir, repo_name)
        if os.path.isdir(repo_path):
            _state(f"OK custom node already cloned: {repo_name}")
        else:
            _state(f"CLONE custom node {repo}")
            run_command(f"git clone https://github.com/{repo}.git", cwd=custom_nodes_dir)

        req_txt = os.path.join(repo_path, "requirements.txt")
        if os.path.exists(req_txt):
            _state(f"PIP custom node requirements: {repo_name}")
            run_command(f"{sys.executable} -m pip install -r requirements.txt", cwd=repo_path)


def _set_manager_security_weak(comfyui_dir: str) -> None:
    """Write ComfyUI-Manager's security_level = weak to every config path it
    might read. Newer Manager versions (V3.x+, confirmed against V3.41)
    moved config.ini from user/default/ComfyUI-Manager/ to user/__manager/,
    and their startup migration actively RAISES anything below 'normal' back
    up — so a legacy-path-only config gets silently overridden on every
    boot ("Security level adjusted: weak -> normal"), and 'normal' can 403
    remote/proxied requests like RunPod's *.proxy.runpod.net, even though
    ComfyUI itself is listening and healthy. Writing the new path directly
    is what actually sticks; the legacy path is kept too for older Manager
    versions that never migrated.
    """
    _state("CONFIG manager security weak")
    config_paths = [
        os.path.join(comfyui_dir, "user", "default", "ComfyUI-Manager", "config.ini"),
        os.path.join(comfyui_dir, "user", "__manager", "config.ini"),
    ]
    for config_path in config_paths:
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        with open(config_path, "w", encoding="utf-8", newline="\n") as f:
            f.write("[default]\nsecurity_level = weak\n")


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.
    Returns (has_gpu, detail) where detail is a human-readable reason,
    logged so a user can tell at a glance why a run went CPU-only.
    """
    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 _ensure_pip() -> None:
    """Some minimal RunPod templates ship a python3 with no pip module and no
    ensurepip at all ('No module named pip' / 'No module named ensurepip') —
    both can be stripped from a slim base image. Fall back to get-pip.py,
    version-pinned to the running interpreter: the current get-pip.py refuses
    to run on Python <3.10 and tells you to use the /pip/<major.minor>/
    variant instead, so build that URL rather than hardcoding the generic one."""
    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"
    _state(f"BOOTSTRAP pip fallback via get-pip.py (Python {py_major}.{py_minor})")
    # Version-pinned filename + forced re-download: download_url_file()'s
    # idempotency check ("already present, skip") is right for multi-GB model
    # assets but wrong for this bootstrap script — a stale /tmp/get-pip.py
    # fetched via the generic (non-version-pinned) URL on an earlier failed
    # attempt would otherwise be silently reused and fail the same way again.
    get_pip_path = f"/tmp/get-pip-{py_major}.{py_minor}.py"
    if os.path.exists(get_pip_path):
        os.remove(get_pip_path)
    download_url_file(get_pip_url, get_pip_path)
    run_command(f"{sys.executable} {get_pip_path}")


def install() -> bool:
    """Runs the full install (deps, custom nodes, model downloads) regardless
    of hardware — none of that is GPU-bound and a CPU-only pod can still do
    all of it, which lets a user pre-stage a CPU pod's downloads before
    switching to a GPU pod (no re-downloading the ~19GB of models there).

    Returns True if a GPU was detected (safe to `serve()` next), False if
    this pod is CPU-only (caller should not attempt to launch ComfyUI).
    """
    comfyui_dir = "/workspace/ComfyUI"
    try:
        _state("PHASE install start")
        _ensure_pip()
        if not os.path.isdir(comfyui_dir):
            _state("CLONE ComfyUI")
            run_command("git clone https://github.com/comfyanonymous/ComfyUI.git /workspace/ComfyUI")

        _state("PIP ComfyUI requirements")
        run_command(f"{sys.executable} -m pip install -r requirements.txt", cwd=comfyui_dir)

        setup_custom_nodes(comfyui_dir)
        download_assets(comfyui_dir)
        _set_manager_security_weak(comfyui_dir)

        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) -> None:
    """--enable-cors-header is required when ComfyUI sits behind RunPod's
    proxy: the proxy forwards requests with a Host header that doesn't match
    the browser's Origin, and ComfyUI's own request validation treats that
    mismatch as a forged request and 403s it — even though the port is
    reachable and the server is otherwise healthy (ComfyUI issue #4865).
    Without this flag, ComfyUI's own logs look completely clean while every
    browser request to the RunPod proxy URL still gets rejected.
    """
    comfyui_dir = "/workspace/ComfyUI"
    _state("PHASE serve")
    _free_port_if_stale(port)
    print(f"[SETUP] Launching ComfyUI on {host}:{port} (--enable-cors-header for RunPod proxy)...")
    try:
        os.chdir(comfyui_dir)
    except Exception as e:
        _state(f"FATAL chdir {comfyui_dir} failed: {e}")
        raise
    os.execv(sys.executable, [sys.executable, "main.py", "--listen", host, "--port", str(port), "--enable-cors-header"])


def _serve_cpu_only_notice(host: str, port: int, gpu_detail: str) -> None:
    """Setup succeeded (ComfyUI cloned, models downloaded) but this pod has no
    GPU, so launching ComfyUI would either fail outright or be unusably slow
    for a ~12GB fp8 diffusion 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: ComfyUI + all models are downloaded and ready.\n"
        "This pod has no GPU attached "
        f"({gpu_detail}), so ComfyUI 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 models 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 (uptime resets, ports stay mapped, no signal
    ever reaches the caller). Instead, bind the ComfyUI 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):
        # A prior failed run's diagnostic server (or ComfyUI itself) may still
        # hold this port if the process was killed/restarted rather than
        # exiting cleanly; without SO_REUSEADDR the bind below raises
        # "Address already in use" and this fallback path crashes too.
        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 (killed by a crash, a manual Ctrl-C,
    or the pod's container command being re-run) can leave its diagnostic
    server or ComfyUI itself still running and holding the port. SO_REUSEADDR
    only papers over TIME_WAIT sockets, not a genuinely live listener, so
    without this a re-run keeps failing with "Address already in use" even
    though there's nothing wrong with the new attempt. Best-effort: tries
    fuser then lsof (whichever is on PATH), and never raises — if neither
    tool exists, the caller's own bind will surface the real error.
    """
    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 ComfyUI server")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8188)
    args = parser.parse_args()

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

    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:
        # Always re-check live rather than trust install()'s earlier result:
        # `--serve` can be invoked standalone in a fresh process (e.g. after
        # install already ran once before), or the pod could have been
        # resized from CPU to GPU between the two calls.
        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)


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