#!/usr/bin/env python3
"""
Auto-generated setup script for ComfyUI + Qwen-Image-2.1 (Comfy-Org/Qwen-Image-2.1)
Generated by Model Lab.

Installs: ComfyUI, and the Qwen-Image-2.1 diffusion model, its Qwen3-VL text
encoder, and its VAE (int8 quantized versions by default, sized for a 12-16GB
GPU). 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 at least 12GB VRAM for the int8 default, 24GB+ for --full (bf16).
  2. Set container disk to >=40GB (>=60GB if using --full).
  3. Put this file at /workspace/setup.py (upload it, or paste as the pod's
     on-start script). To also get the no-node-graph Gradio UI, put
     qwen_image_2_1_gradio.py in the SAME directory (/workspace) - --serve
     launches it automatically if it's there, and skips it silently if not.
  4. Container start command:
       python /workspace/setup.py --install && python /workspace/setup.py --serve
  5. Expose/forward TCP port 8188 for ComfyUI. The Gradio UI defaults to
     port 8889 - already in the guide's standard "Expose HTTP ports" list
     (8888,4000,8188,8889), so nothing extra to add if you set that up as
     written. Pass --gradio-port to use a different one instead.

Options:
  --full       Download the bf16 (full precision) diffusion model and text
               encoder instead of the int8 defaults. ~14.2GB + 17.5GB instead
               of ~7.3GB + 6.3GB. Needs 24GB+ VRAM.
  --gguf       Download the Q4_K_M GGUF quantized diffusion model instead
               (~4.2GB, via city96's ComfyUI-GGUF custom node). For 6-10GB
               VRAM GPUs. Still downloads the int8 text encoder and VAE.
  --no-gradio  Don't launch the Gradio UI even if qwen_image_2_1_gradio.py
               is present next to this script. ComfyUI still starts as usual.

Downloads are pinned to a specific commit on each source repo (see
COMFY_ORG_REPO_REVISION / ABIRAY_GGUF_REPO_REVISION below) rather than
"resolve/main", and are never mirrored onto a GenLovers-controlled host -
see the comment above those constants for why. If a download 404s, the
upstream repo has likely restructured since these were pinned; re-check it
and bump the revision rather than assuming the URL is simply wrong.

License note: Qwen-Image-2.1's weights are released under the Qwen RESEARCH
LICENSE AGREEMENT - non-commercial use only (research/evaluation). Commercial
use of the model or its outputs requires a separate license from Alibaba/Qwen
(model-business@notice.qwencloud.com). This differs from the original
Qwen-Image and Qwen-Image-Edit models, which ship Apache 2.0. Verify the
license yourself before using anything generated here commercially.

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):
        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", "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

    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"]:
            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()


# Pinned to a specific commit rather than "resolve/main" on both source repos.
# We deliberately do NOT mirror these weights onto a GenLovers-controlled HF
# account (Qwen-Image-2.1's own license is non-commercial research-use only
# and requires carrying its license file + attribution with any redistributed
# copy - re-hosting a full copy is a redistribution act, not a convenience).
# Pinning a commit is the cheap alternative: protects a reader mid-tutorial if
# upstream force-pushes a filename change or restructures the repo, at zero
# storage cost, while still always resolving to the exact bytes verified here.
# Bump these SHAs (and the `bytes` floors below) only after re-verifying the
# new revision still has every file this script expects, at the paths it
# expects them at - a bad bump silently 404s every download.
COMFY_ORG_REPO_REVISION = "9a44dbdb47cefd046be9c0a13476192f34c8db8e"  # Comfy-Org/Qwen-Image-2.1, checked 2026-09-24
ABIRAY_GGUF_REPO_REVISION = "c9dd12108f53974cd1e0abd708df042d6df0ca8d"  # Abiray/Qwen-Image-2.1-GGUF, checked 2026-09-24

# Three variants selectable via --full / --gguf / default (int8). `bytes` is
# the exact file size at the pinned revision above, minus ~1% slack, used only
# to catch a truncated download or an HTML error page saved in place of the
# real file - not a hash-level integrity check.
DIFFUSION_MODEL_VARIANTS = {
    "int8": {
        "url": f"https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/{COMFY_ORG_REPO_REVISION}/diffusion_models/qwen_image_2.1_int8_convrot.safetensors",
        "target_dir": "diffusion_models",
        "filename": "qwen_image_2.1_int8_convrot.safetensors",
        "bytes": 7180000000,  # exact: 7256783064
    },
    "full": {
        "url": f"https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/{COMFY_ORG_REPO_REVISION}/diffusion_models/qwen_image_2.1_bf16.safetensors",
        "target_dir": "diffusion_models",
        "filename": "qwen_image_2.1_bf16.safetensors",
        "bytes": 14080000000,  # exact: 14230280616
    },
    "gguf": {
        "url": f"https://huggingface.co/Abiray/Qwen-Image-2.1-GGUF/resolve/{ABIRAY_GGUF_REPO_REVISION}/qwen_image_2.1_Q4_K_M.gguf",
        "target_dir": "unet",
        "filename": "qwen_image_2.1_Q4_K_M.gguf",
        "bytes": 4150000000,  # exact: 4189343904
    },
}

TEXT_ENCODER_VARIANTS = {
    "int8": {
        "url": f"https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/{COMFY_ORG_REPO_REVISION}/text_encoders/qwen3vl_8b_w4a8.safetensors",
        "target_dir": "text_encoders",
        "filename": "qwen3vl_8b_w4a8.safetensors",
        "bytes": 6250000000,  # exact: 6312105364
    },
    "full": {
        "url": f"https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/{COMFY_ORG_REPO_REVISION}/text_encoders/qwen3vl_8b_bf16.safetensors",
        "target_dir": "text_encoders",
        "filename": "qwen3vl_8b_bf16.safetensors",
        "bytes": 17350000000,  # exact: 17534334616
    },
}

VAE_ASSET = {
    "url": f"https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/{COMFY_ORG_REPO_REVISION}/vae/qwen_image_2.1_vae_bf16.safetensors",
    "target_dir": "vae",
    "filename": "qwen_image_2.1_vae_bf16.safetensors",
    "bytes": 669000000,  # exact: 675509688
}

MANIFEST = {
  "model_ref": "Qwen/Qwen-Image-2.1",
  "slug": "qwen-image-2-1",
  "modality": "image",
  "archetype": "text_to_image_and_edit",
  "container_disk_in_gb": 40,
  "gpu_preferences": [
    "NVIDIA RTX 4090",
    "NVIDIA RTX A5000",
    "NVIDIA L40S",
    "NVIDIA A40",
  ],
  "custom_nodes": [],
}

GGUF_CUSTOM_NODE = "city96/ComfyUI-GGUF"


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:
    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
    if expected_bytes and os.path.getsize(destination) < expected_bytes:
        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:
    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}")
        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)

    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:
        attempts += 1
        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, variant: str) -> None:
    """variant is one of "int8" (default), "full", "gguf". The text encoder
    and VAE are shared across all three; only the diffusion model swaps to a
    GGUF file loaded by a different node when variant == "gguf"."""
    _state(f"DOWNLOAD_PHASE_START variant={variant}")
    _ensure_aria2()
    models_root = os.path.join(comfyui_dir, "models")
    hf_token = os.environ.get("HF_TOKEN", "")

    diffusion_asset = DIFFUSION_MODEL_VARIANTS["gguf" if variant == "gguf" else variant]
    text_encoder_asset = TEXT_ENCODER_VARIANTS["full" if variant == "full" else "int8"]
    assets = [diffusion_asset, text_encoder_asset, VAE_ASSET]

    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}")
                raise

    _state("DOWNLOAD_PHASE_DONE")


def _install_manager_requirements(comfyui_dir: str) -> None:
    """Current ComfyUI ships Manager built into core (manager_requirements.txt
    at the repo root) rather than as a custom_nodes git clone. Only its Python
    deps need installing here; the --enable-manager flag on the main.py launch
    (below, in serve()) is what actually turns the UI on."""
    req_path = os.path.join(comfyui_dir, "manager_requirements.txt")
    if not os.path.exists(req_path):
        _state("WARN manager_requirements.txt not found, skipping (older ComfyUI checkout?)")
        return
    _state("PIP manager requirements")
    run_command(f"{sys.executable} -m pip install -r manager_requirements.txt", cwd=comfyui_dir)


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

    repos = list(MANIFEST.get("custom_nodes", []))
    if variant == "gguf":
        repos.append(GGUF_CUSTOM_NODE)

    for repo in repos:
        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+) 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, and
    'normal' can 403 remote/proxied requests like RunPod's *.proxy.runpod.net.
    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.
    """
    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. 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"
    _state(f"BOOTSTRAP pip fallback via get-pip.py (Python {py_major}.{py_minor})")
    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}")


GRADIO_SCRIPT_NAME = "qwen_image_2_1_gradio.py"


def _gradio_script_path() -> Optional[str]:
    """The Gradio UI is optional and lives as a separate downloadable file
    (per the guide). Look for it next to this script rather than assuming a
    fixed /workspace path, so it works whether both files were uploaded
    together into /workspace or into some other working directory."""
    candidate = os.path.join(os.path.dirname(os.path.abspath(__file__)), GRADIO_SCRIPT_NAME)
    return candidate if os.path.isfile(candidate) else None


def _install_gradio_requirements() -> None:
    _state("PIP gradio UI requirements")
    run_command(f"{sys.executable} -m pip install gradio requests pillow")


def install(variant: str) -> 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.

    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(f"PHASE install start variant={variant}")
        _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)

        _install_manager_requirements(comfyui_dir)
        setup_custom_nodes(comfyui_dir, variant)
        download_assets(comfyui_dir, variant)
        _set_manager_security_weak(comfyui_dir)

        # LoRAs (docs step 6 of the guide) are dropped in manually by the
        # reader, but the folder should exist so Manager's model browser and
        # the workflow's Load LoRA node see it immediately.
        os.makedirs(os.path.join(comfyui_dir, "models", "loras"), exist_ok=True)

        if _gradio_script_path():
            _install_gradio_requirements()

        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 _wait_for_http_ready(url: str, timeout_s: int, label: str) -> bool:
    """Poll a URL until it responds (any status code counts - we're checking
    the port is bound and answering, not that the response is a 200), or
    give up after timeout_s. Used to gate the Gradio launch on ComfyUI
    actually being up, rather than a fixed sleep that's either too short on
    a slow pod or wastes time on a fast one."""
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            urllib.request.urlopen(url, timeout=5)
            return True
        except urllib.error.HTTPError:
            return True  # server answered, even with an error status
        except Exception:
            time.sleep(1)
    _state(f"WARN {label} did not become ready within {timeout_s}s at {url}")
    return False


def _launch_comfyui_process(host: str, port: int) -> subprocess.Popen:
    """--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 (ComfyUI issue #4865).

    --enable-manager turns on ComfyUI-Manager's UI, needed for "Install
    Missing Custom Nodes" if the reader loads the official workflow before
    this script's own custom-node install has run.

    Runs as a real subprocess (not os.execv) so this script's own process can
    stay alive afterward to also launch and supervise the optional Gradio UI.
    """
    comfyui_dir = "/workspace/ComfyUI"
    log_path = "/workspace/comfyui_stdout.log"
    _state(f"LAUNCH ComfyUI on {host}:{port}, logging to {log_path}")
    log_file = open(log_path, "a", encoding="utf-8")
    return subprocess.Popen(
        [sys.executable, "main.py", "--listen", host, "--port", str(port), "--enable-cors-header", "--enable-manager"],
        cwd=comfyui_dir,
        stdout=log_file,
        stderr=subprocess.STDOUT,
    )


def _launch_gradio_process(comfy_url: str, variant: str, port: int) -> Optional[subprocess.Popen]:
    """Launches qwen_image_2_1_gradio.py if it's present next to this script,
    passing --unet/--clip overrides so it matches whichever weight tier
    --install actually downloaded (the Gradio script's own defaults are the
    int8 filenames). Returns None (not an error) if the file isn't there -
    the Gradio UI is optional and ComfyUI alone is a complete setup."""
    script_path = _gradio_script_path()
    if not script_path:
        _state(f"SKIP Gradio UI: {GRADIO_SCRIPT_NAME} not found next to this script")
        return None

    log_path = "/workspace/gradio_stdout.log"
    cmd = [sys.executable, script_path, "--comfy-url", comfy_url, "--port", str(port)]
    if variant == "full":
        cmd += ["--unet", DIFFUSION_MODEL_VARIANTS["full"]["filename"], "--clip", TEXT_ENCODER_VARIANTS["full"]["filename"]]
    elif variant == "gguf":
        cmd += ["--unet", DIFFUSION_MODEL_VARIANTS["gguf"]["filename"]]  # text encoder stays int8 in gguf mode

    _state(f"LAUNCH Gradio UI on port {port}, logging to {log_path}")
    log_file = open(log_path, "a", encoding="utf-8")
    return subprocess.Popen(cmd, stdout=log_file, stderr=subprocess.STDOUT)


def serve(host: str, port: int, variant: str, launch_gradio: bool, gradio_port: int) -> None:
    """Starts ComfyUI, waits for it to be reachable, then starts the Gradio
    UI on top of it if requested and available. Both run as subprocesses of
    this one process, which then blocks supervising them: if either exits
    unexpectedly it is restarted, and a signal to this process (RunPod
    stopping the container) is forwarded to both children before exiting -
    without that, stopping the pod would leave an orphaned ComfyUI or Gradio
    process holding its port on the next start.
    """
    _state("PHASE serve")
    _free_port_if_stale(port)
    if launch_gradio:
        _free_port_if_stale(gradio_port)

    comfy_process = _launch_comfyui_process(host, port)
    comfy_url = f"http://127.0.0.1:{port}"
    print(f"[SETUP] Waiting for ComfyUI to come up on {comfy_url} ...")
    _wait_for_http_ready(comfy_url, timeout_s=180, label="ComfyUI")

    gradio_process = None
    if launch_gradio:
        gradio_process = _launch_gradio_process(comfy_url, variant, gradio_port)
        if gradio_process is not None:
            print(f"[SETUP] Gradio UI starting on port {gradio_port} (see /workspace/gradio_stdout.log)")

    print(f"[SETUP] ComfyUI is up on port {port}" + (f", Gradio UI on port {gradio_port}" if gradio_process else "") + ". Ctrl+C or a pod stop shuts both down cleanly.")

    def _shutdown(*_args) -> None:
        _state("SHUTDOWN signal received, terminating child processes")
        for proc in (comfy_process, gradio_process):
            if proc is not None and proc.poll() is None:
                proc.terminate()
        sys.exit(0)

    import signal
    signal.signal(signal.SIGTERM, _shutdown)
    signal.signal(signal.SIGINT, _shutdown)

    try:
        while True:
            time.sleep(5)
            comfy_exit_code = comfy_process.poll()
            if comfy_exit_code is not None:
                _state(f"WARN ComfyUI exited (code {comfy_exit_code}), restarting")
                comfy_process = _launch_comfyui_process(host, port)
            if gradio_process is not None:
                gradio_exit_code = gradio_process.poll()
                if gradio_exit_code is not None:
                    _state(f"WARN Gradio UI exited (code {gradio_exit_code}), restarting")
                    gradio_process = _launch_gradio_process(comfy_url, variant, gradio_port)
    except KeyboardInterrupt:
        _shutdown()


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.
    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. 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):
        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. Best-effort: tries
    fuser then lsof (whichever is on PATH), 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 ComfyUI server")
    parser.add_argument("--full", action="store_true", help="Download bf16 (full precision) weights instead of int8")
    parser.add_argument("--gguf", action="store_true", help="Download the Q4_K_M GGUF quantized model for 6-10GB VRAM GPUs")
    parser.add_argument("--no-gradio", action="store_true", help="Don't launch the Gradio UI even if qwen_image_2_1_gradio.py is present")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8188)
    parser.add_argument("--gradio-port", type=int, default=8889)
    args = parser.parse_args()

    if args.full and args.gguf:
        print("[SETUP] --full and --gguf are mutually exclusive", file=sys.stderr)
        sys.exit(2)
    variant = "full" if args.full else ("gguf" if args.gguf else "int8")

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

    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(variant)
        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, variant, launch_gradio=not args.no_gradio, gradio_port=args.gradio_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)