#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI 一键安装 + 启动
=======================
流程: 克隆 ComfyUI → 建虚拟环境 → 装依赖 → 装 ComfyUI-Manager → 启动服务

用法:
  python install_comfyui.py              # 全自动: 安装并启动
  python install_comfyui.py --no-launch  # 只安装, 不启动
  python install_comfyui.py --novram     # 显存不够被 kill 时用(更激进的卸载)

中断后直接重跑, 已完成步骤自动跳过。

"""
import os
import re
import sys
import time
import shutil
import socket
import subprocess

# ============ 配置 ============
WORKSPACE_CANDIDATES = ["/home/admin/workspace", "/mnt/workspace"]  # 自动探测，第一个存在的胜出
REPO_URL = "https://github.com/comfyanonymous/ComfyUI.git"
PORT = 8188
LOG_DIR = "/tmp/comfyui_install_logs"

# 启动参数集中维护（--novram 时会自动追加到最后）
LAUNCH_ARGS = [
    "--enable-manager",
    "--disable-pinned-memory",   # 省系统内存：不再锁定 ~2 倍模型大小的页
    "--fast-disk",               # 省系统内存：优先从磁盘流式加载而非常驻 RAM
    "--preview-method", "none",  # 关生成中预览，省显存
    "--reserve-vram", "0",       # 默认会预留 1.2GB 显存，设 0 可拿回来
]

# ============ 路径（自动派生，不要手改）============
COMFY_DIR = ""
VENV_DIR = ""


def resolve_paths():
    global COMFY_DIR, VENV_DIR
    ws = next((c for c in WORKSPACE_CANDIDATES if os.path.isdir(c)), None)
    if ws is None:
        ws = WORKSPACE_CANDIDATES[0]
        os.makedirs(ws, exist_ok=True)
    COMFY_DIR = os.path.join(ws, "ComfyUI")
    VENV_DIR = os.path.join(ws, "venv")    # 放工作区，实例重启后仍在
    return ws


def run(cmd, **kw):
    print(f"  $ {' '.join(cmd)}")
    subprocess.run(cmd, check=True, **kw)


def port_open():
    try:
        socket.create_connection(("127.0.0.1", PORT), timeout=2).close()
        return True
    except Exception:
        return False


# ---------- 幂等工具 ----------
def _norm(name):
    """包名规范化：去版本约束与环境标记，统一小写"""
    return re.split(r"[<>=!~\[;@]", name.strip())[0].strip().lower().replace("_", "-")


def installed_set(pip_cmd):
    """当前环境已装包名集合"""
    try:
        r = subprocess.run(pip_cmd + ["list", "--format=freeze"],
                           capture_output=True, text=True, timeout=300)
        if r.returncode != 0:
            return set()
        return {ln.split("==")[0].split(" @")[0].strip()
                for ln in r.stdout.splitlines() if ln.strip()}
    except Exception:
        return set()


def missing_from_requirements(pip_cmd, req_path):
    """返回 (缺失行列表, 总行数)。全部已装则缺失为空"""
    have = installed_set(pip_cmd)
    need, total = [], 0
    try:
        lines = open(req_path, encoding="utf-8").read().splitlines()
    except OSError:
        return [], 0
    for ln in lines:
        s = ln.strip()
        if not s or s.startswith("#") or s.startswith("-"):
            continue                      # 跳过注释与 --index-url 等选项行
        total += 1
        if _norm(s) not in have:
            need.append(s)
    return need, total


def main():
    ws = resolve_paths()
    vpy = os.path.join(VENV_DIR, "bin", "python")
    vpip = [vpy, "-m", "pip"]        # ★ 用解释器调 pip 模块，杜绝装错环境

    # pip 缓存放工作区：重跑命中缓存，省重复下载
    cache_dir = os.path.join(ws, ".cache", "pip")
    os.environ["PIP_CACHE_DIR"] = cache_dir
    os.makedirs(cache_dir, exist_ok=True)

    print("=" * 58)
    print(" ComfyUI 一键安装 + 启动")
    print("=" * 58)
    print(f" 工作区  : {ws}")
    print(f" ComfyUI : {COMFY_DIR}")
    print(f" venv    : {VENV_DIR}")
    print(f" pip缓存 : {cache_dir}")
    print("=" * 58)

    # 1. 克隆（有 main.py 才算装好，避免空壳目录被误判为已完成）
    print("\n[1/4] 克隆 ComfyUI")
    if os.path.isfile(os.path.join(COMFY_DIR, "main.py")):
        print(f"  [跳过] 已存在: {COMFY_DIR}")
    else:
        if os.path.isdir(COMFY_DIR):
            print(f"  [清理] 目录不完整(无 main.py)，重新克隆: {COMFY_DIR}")
            shutil.rmtree(COMFY_DIR)
        run(["git", "clone", "--depth", "1", REPO_URL, COMFY_DIR])
        print("  [OK] 克隆完成")

    # 2. 虚拟环境（--system-site-packages 复用镜像预装的 torch，不重下几个 GB）
    print("\n[2/4] 创建虚拟环境")
    if os.path.exists(vpy):
        print(f"  [跳过] venv 已存在: {VENV_DIR}")
    else:
        run([sys.executable, "-m", "venv", "--system-site-packages", VENV_DIR])
        print("  [OK] 创建完成（复用镜像预装的 torch）")

    # 3. 依赖
    print("\n[3/4] 安装依赖")
    run(vpip + ["install", "-q", "--upgrade", "pip"])

    req = os.path.join(COMFY_DIR, "requirements.txt")
    if os.path.exists(req):
        need, total = missing_from_requirements(vpip, req)
        if total and not need:
            print(f"  [跳过] requirements.txt 的 {total} 个包已全部安装")
        else:
            print(f"  torch 等镜像已装的会自动跳过，不重复下载（缺 {len(need)}/{total}）")
            run(vpip + ["install", "-r", req])
        print("  [OK] 依赖安装完成")
    else:
        print(f"  [警告] 未找到 {req}，跳过依赖安装")

    print("  安装 comfyui-manager（解决工作流节点缺失）")
    if _norm("comfyui-manager") in installed_set(vpip):
        print("  [跳过] comfyui-manager 已安装")
    else:
        run(vpip + ["install", "-q", "comfyui-manager"])
    print("  [OK] comfyui-manager 安装完成")

    print("\n" + "=" * 58)
    print(" 安装完成 ✅")
    print("=" * 58)

    if "--no-launch" in sys.argv:
        print(f" [完成] 未启动。手动启动:")
        print(f"        cd {COMFY_DIR} && {vpy} main.py --listen 0.0.0.0 --port {PORT}")
        return

    # 4. 启动
    print("\n[4/4] 启动服务")
    if port_open():
        print(f"  [跳过] 端口 {PORT} 已在监听（服务可能已运行）")
        return

    args = list(LAUNCH_ARGS)
    if "--novram" in sys.argv:
        args.append("--novram")

    cmd = [vpy, "main.py", "--listen", "0.0.0.0", "--port", str(PORT)] + args
    os.makedirs(LOG_DIR, exist_ok=True)
    log_path = os.path.join(LOG_DIR, "comfyui.log")

    lf = open(log_path, "w", encoding="utf-8")
    proc = subprocess.Popen(cmd, cwd=COMFY_DIR, stdout=lf, stderr=subprocess.STDOUT)
    lf.close()          # 子进程已持有自己的 fd，父进程可以关掉

    print(f"  [信息] 已后台启动 (PID {proc.pid})，日志: {log_path}")
    print(f"  [等待] 等待 {PORT} 就绪（最多 300 秒）")
    for i in range(300):
        if port_open():
            print(f"  [OK] 就绪（耗时 {i} 秒）→ http://127.0.0.1:{PORT}")
            return
        if proc.poll() is not None:      # 进程秒退，不用傻等
            print(f"  [错误] 进程已退出(返回码 {proc.returncode})，查看日志: {log_path}")
            return
        time.sleep(1)

    print(f"  [错误] 300 秒内未就绪，查看日志: {log_path}")


if __name__ == "__main__":
    try:
        main()
    except subprocess.CalledProcessError as e:
        print(f"\n[中止] 命令执行失败: {' '.join(e.cmd)}")
        print("       修复后直接重跑本脚本，已完成步骤会自动跳过")
        sys.exit(1)
    except KeyboardInterrupt:
        print("\n[退出] 已中断")
        sys.exit(130)

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ModelScope 大模型下载工具（带进度条 + 日志 + 断点续传 + 重试）

功能:
- 实时显示每个文件的下载进度条 (tqdm) + 速度/已下载/预计剩余
- 全量日志记录到文件,控制台只显示进度
- 断点续传: 中断后重跑自动从已下载字节继续
- 失败重试 + 退避,应对 "资源解析服务请求失败"
- 完成时做大小校验 (达到 expected 的 95% 视为完成)
"""
import os
import sys
import time
import json
import logging
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    from tqdm import tqdm
except ImportError:
    os.system(f"{sys.executable} -m pip install -q tqdm")
    from tqdm import tqdm


# ============ 配置 (按需修改) ============
BASE = "https://modelscope.cn/models/Comfy-Org/MiniMax-H3/resolve/master"
BASE_DIR = "/home/admin/workspace/ComfyUI/models"

# 任务列表: (相对路径, 期望大小GB)
JOBS = [
    ("diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors", 20.97),
    ("diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors", 20.97),
    ("text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", 17.14),
    ("vae/minimax_h3_video_vae_fp16.safetensors", 5.21),
    ("vae/minimax_h3_audio_vae_fp32.safetensors", 0.61),
  
    ("loras/minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors", 2),
]

WORKERS = 2              # 并发数(大模型建议 2,避免把解析服务打满)
RETRIES = 4             # 每个文件最大重试次数
CHUNK_SIZE = 1 << 20    # 1MB
TIMEOUT = 120           # 单次请求超时(秒)
LOG_DIR = "/tmp/model_download_logs"


# ============ 日志初始化 ============
def setup_logger():
    os.makedirs(LOG_DIR, exist_ok=True)
    log_file = os.path.join(
        LOG_DIR, f"download_{datetime.now():%Y%m%d_%H%M%S}.log"
    )
    logger = logging.getLogger("downloader")
    logger.setLevel(logging.DEBUG)

    # 文件 handler: 记录全部 DEBUG
    fh = logging.FileHandler(log_file, encoding="utf-8")
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
    logger.addHandler(fh)

    # 控制台 handler: 只显示 INFO 及以上(进度条由 tqdm 单独管理)
    ch = logging.StreamHandler(sys.stdout)
    ch.setLevel(logging.INFO)
    ch.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
    logger.addHandler(ch)

    logger.info(f"日志文件: {log_file}")
    return logger


LOG = setup_logger()


# ============ 下载核心 ============
def _stream_download(url, path, exp_bytes, done, desc, pbar):
    """从 done 字节处继续流式下载,更新进度条"""
    headers = {"Range": f"bytes={done}-"} if done else {}
    # 遇到 200(不支持断点) 则从头开始
    import requests
    r = requests.get(url, headers=headers, stream=True, timeout=TIMEOUT)
    if r.status_code == 200 and done:
        LOG.warning(f"{desc} 服务端不支持 Range, 从头下载")
        done = 0
    r.raise_for_status()

    # 若服务端返回 Content-Length, 用它校正总大小
    content_length = r.headers.get("Content-Length")
    if content_length:
        pbar.total = done + int(content_length)
        pbar.refresh()

    mode = "ab" if done else "wb"
    with open(path, mode) as f:
        for chunk in r.iter_content(CHUNK_SIZE):
            if not chunk:
                continue
            f.write(chunk)
            done += len(chunk)
            pbar.update(len(chunk))
    return done


def download(job):
    """下载单个文件,带重试"""
    rel, exp_gb = job
    exp_bytes = int(exp_gb * 1e9)
    url = f"{BASE}/{rel}"
    path = os.path.join(BASE_DIR, rel)
    os.makedirs(os.path.dirname(path), exist_ok=True)

    # 已完成则跳过
    if os.path.exists(path) and os.path.getsize(path) >= exp_bytes * 0.95:
        LOG.info(f"[跳过] {rel} (已存在且达到预期大小)")
        return f"[跳过] {rel}"

    desc = os.path.basename(rel)
    for attempt in range(1, RETRIES + 1):
        current = os.path.getsize(path) if os.path.exists(path) else 0
        with tqdm(
            total=exp_bytes,
            initial=current,
            unit="B",
            unit_scale=True,
            desc=f"{desc:40s}",
            miniters=1,
            leave=True,
        ) as pbar:
            try:
                _stream_download(url, path, exp_bytes, current, desc, pbar)
                final = os.path.getsize(path)
                speed_mb = final / (time.time() - pbar.start_t) / 1e6 if False else 0
                LOG.info(
                    f"[完成] {rel}: {final/1e9:.2f}/{exp_gb:.2f} GB "
                    f"({final/exp_bytes*100:.1f}%)"
                )
                return f"[完成] {rel}: {final/1e9:.2f} GB"
            except Exception as e:
                LOG.error(f"[重试 {attempt}/{RETRIES}] {rel}: {e}")
                time.sleep(2 ** attempt)  # 指数退避

    return f"[失败] {rel}: 超过最大重试次数"


# ============ 入口 ============
def main():
    LOG.info(f"目标目录: {BASE_DIR}")
    LOG.info(f"任务数: {len(JOBS)}, 并发: {WORKERS}")
    os.makedirs(BASE_DIR, exist_ok=True)

    results = []
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futures = {ex.submit(download, j): j for j in JOBS}
        for fut in as_completed(futures):
            results.append(fut.result())

    # 汇总
    LOG.info("=" * 50)
    LOG.info("下载结果汇总:")
    for msg in results:
        LOG.info(f"  {msg}")

    # 最终核对
    LOG.info("=" * 50)
    LOG.info("最终文件核对:")
    all_ok = True
    for rel, exp_gb in JOBS:
        p = os.path.join(BASE_DIR, rel)
        if os.path.exists(p) and os.path.getsize(p) >= int(exp_gb * 1e9) * 0.95:
            LOG.info(f"  [OK]   {rel}: {os.path.getsize(p)/1e9:.2f}/{exp_gb:.2f} GB")
        else:
            all_ok = False
            size = os.path.getsize(p) if os.path.exists(p) else 0
            LOG.warning(f"  [缺失] {rel}: {size/1e9:.2f}/{exp_gb:.2f} GB")

    if all_ok:
        LOG.info("全部文件下载完成 ✅")
    else:
        LOG.warning("部分文件缺失,可重跑脚本自动续传")
        sys.exit(1)


if __name__ == "__main__":
    main()

pip3 install  pyngrok

python3 - <<'PY'        # ← 开头：把下面内容喂给 python3
from pyngrok import ngrok
ngrok.set_auth_token("你的token")
t = ngrok.connect("127.0.0.1:8188", bind_tls=True)
print("公网地址:", t.public_url)
while True: time.sleep(60)
PY                      # ← 结尾：到此为止

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ngrok 内网穿透 - 命令行工具（修正版）

修复点:
1. 移除 region 参数 (新版 HTTPv2 隧道不支持 tunnel 级 region 字段,会报 400)
2. connect() 只用 addr / bind_tls / auth 等合法关键字
3. 先探测本地端口是否就绪,再建隧道 (避免建完隧道 502)
4. token 优先从环境变量读取,避免硬编码泄露
5. 依赖缺失自动 pip install
6. 支持 connect / status / disconnect / kill 四种动作
"""
import os
import sys
import time
import socket
import subprocess
import argparse


# ============ 自动安装依赖 ============
def ensure_pyngrok():
    try:
        import pyngrok  # noqa
    except ImportError:
        print("[...] 未检测到 pyngrok,正在安装...")
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "-q", "pyngrok"]
        )


# ============ 端口探测 ============
def port_open(host, port, timeout=2.0):
    """检测本地端口是否可接受连接"""
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except Exception:
        return False


def http_ready(url, timeout=5):
    """简单 HTTP 探测,确认服务真的在响应"""
    try:
        import urllib.request
        req = urllib.request.Request(url, method="HEAD")
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status < 500
    except Exception:
        return False


# ============ ngrok 封装 ============
def configure_token(token):
    """写入最小配置,不包含 region 字段"""
    from pyngrok import ngrok
    if token:
        ngrok.set_auth_token(token)
    # 强制生成一个干净的配置(不写 region)
    cfg_dir = os.path.expanduser("~/.config/ngrok")
    os.makedirs(cfg_dir, exist_ok=True)
    cfg_path = os.path.join(cfg_dir, "ngrok.yml")
    with open(cfg_path, "w", encoding="utf-8") as f:
        f.write(f"version: 3\nagent:\n  authtoken: {token or 'YOUR_TOKEN'}\n")
    print(f"[OK] authtoken 已配置 -> {cfg_path}")


def build_kwargs(port, bind_tls, basic_auth):
    """只使用合法的关键字参数,绝不传 region"""
    kwargs = {
        "addr": port,          # 关键字传端口,最稳妥
        "bind_tls": bind_tls,
    }
    if basic_auth:
        kwargs["auth"] = basic_auth  # "user:pass"
    return kwargs


def action_connect(port, token, bind_tls, basic_auth, wait):
    from pyngrok import ngrok

    # 1) 配置 token
    configure_token(token)

    # 2) 端口自检
    print(f"[...] 检测本地服务 http://127.0.0.1:{port}")
    if not port_open("127.0.0.1", port):
        print(f"[错误] 端口 {port} 未监听,请先启动 ComfyUI 等服务")
        print(f"       可尝试: curl -sS http://127.0.0.1:{port}/system_stats")
        sys.exit(1)

    if not http_ready(f"http://127.0.0.1:{port}"):
        print(f"[警告] 端口开放但 HTTP 未就绪,隧道可能暂时 502")

    # 3) 建立隧道 (不带 region)
    print(f"[...] 正在建立隧道: http://127.0.0.1:{port}")
    kwargs = build_kwargs(port, bind_tls, basic_auth)
    tunnel = ngrok.connect(**kwargs)
    url = tunnel.public_url

    print("=" * 60)
    print(f"  公网地址 : {url}")
    print(f"  本地端口 : http://127.0.0.1:{port}")
    if basic_auth:
        print(f"  访问密码 : {basic_auth}")
    print("=" * 60)

    # 4) 阻塞保持,等待 Ctrl+C
    if wait:
        print("[提示] 隧道已在后台运行,按 Ctrl+C 退出")
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print("\n[...] 正在断开隧道...")
            ngrok.disconnect(url)
            ngrok.kill()
            print("[OK] 已退出")
    return url


def action_status():
    from pyngrok import ngrok
    tunnels = ngrok.get_tunnels()
    if not tunnels:
        print("当前没有活跃的隧道")
        return
    print("当前活跃隧道:")
    for t in tunnels:
        print(f"  - {t.public_url}  ->  {t.config.get('addr', '?')}")


def action_disconnect():
    from pyngrok import ngrok
    for t in ngrok.get_tunnels():
        ngrok.disconnect(t.public_url)
    print("[OK] 已断开所有隧道")


def action_kill():
    from pyngrok import ngrok
    ngrok.kill()
    print("[OK] ngrok 进程已终止")


# ============ 命令行入口 ============
def parse_args():
    parser = argparse.ArgumentParser(
        description="ngrok 内网穿透命令行工具 (修正 region 报错版)"
    )
    parser.add_argument(
        "action",
        nargs="?",
        default="connect",
        choices=["connect", "status", "disconnect", "kill"],
        help="动作 (默认: connect)",
    )
    parser.add_argument("--port", type=int, default=8188, help="本地端口 (默认 8188)")
    parser.add_argument(
        "--token",
        type=str,
        default=os.environ.get("NGROK_AUTHTOKEN", ""),
        help="ngrok authtoken (默认读环境变量 NGROK_AUTHTOKEN)",
    )
    parser.add_argument(
        "--basic-auth",
        type=str,
        default="",
        help="访问密码,格式 user:pass (可选)",
    )
    parser.add_argument(
        "--no-tls",
        action="store_true",
        help="关闭强制 HTTPS (默认开启 bind_tls)",
    )
    parser.add_argument(
        "--no-wait",
        action="store_true",
        help="建隧道后立即返回,不阻塞 (用于脚本串联)",
    )
    return parser.parse_args()


def main():
    ensure_pyngrok()
    args = parse_args()

    if not args.token and args.action == "connect":
        print("[警告] 未设置 NGROK_AUTHTOKEN,请先:")
        print("        export NGROK_AUTHTOKEN='你的token'")
        print("        或 --token 你的token")

    bind_tls = not args.no_tls

    if args.action == "connect":
        action_connect(args.port, args.token, bind_tls, args.basic_auth, not args.no_wait)
    elif args.action == "status":
        action_status()
    elif args.action == "disconnect":
        action_disconnect()
    elif args.action == "kill":
        action_kill()


if __name__ == "__main__":
    main()

    
# 常用命令    
# python3 ngrok.py  --token 你的token                            # 建隧道（默认 8188）
# python3 ngrok.py --port 7860 --token xxx      # 换端口
# python3 ngrok.py --token xxx --basic-auth admin:123   # 加访问密码
# python3 ngrok.py status                       # 看当前隧道
# python3 ngrok.py disconnect                   # 断开
# python3 ngrok.py kill                         # 杀进程

export NGROK_AUTHTOKEN='你的token'
python3 ngrok.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI + MiniMax-H3 + ngrok 一键部署脚本
==========================================
三合一:
  1. ComfyUI 安装部署 (git clone + venv + 依赖 + ComfyUI-Manager)
  2. MiniMax-H3 模型下载 (进度条 + 日志 + 断点续传 + 重试)
  3. ngrok 内网穿透 (启动服务后自动建隧道并打印公网链接)

★ 核心设计: 工作区持久化
  实例重启后系统目录(如 /usr/local/lib/python3.x/site-packages)会被重置,
  但工作区(如 /mnt/workspace)会保留。因此:
    - venv 建在工作区内  → 重启后依赖还在, 不用重装
    - 模型下载到工作区   → 重启后 60GB 模型还在, 不用重下
    - pip 缓存放工作区   → 万一要重装, 缓存命中省下载

★ 路径自动探测:
  按顺序探测候选工作区, 第一个存在的胜出; 都不存在则用第一个并创建。
  你不用关心实例到底是 /mnt/workspace 还是 /home/admin/workspace。

★ 并行:
  ComfyUI 克隆 / venv 创建+装 pyngrok / 模型下载 三路并行。
  pip 操作统一加线程锁, 避免多任务同时装包把环境写坏。

★ 幂等:
  - ComfyUI 已存在(有 main.py) → 跳过克隆
  - requirements.txt 里的包已全部安装 → 跳过整条 pip 命令
  - comfyui-manager / pyngrok / requests / tqdm 已装 → 跳过
  - 模型文件已达预期大小 → 跳过
  - torch 能正常调用 CUDA → 跳过对齐

用法:
  python setup_all.py                # 全自动: 安装 → 下模型 → 启服务 → 建隧道 → 打印链接
  python setup_all.py --no-tunnel    # 不建隧道 (用平台自带代理时)
  python setup_all.py --no-launch    # 只安装+下模型, 不启动服务
  python setup_all.py --no-align     # 跳过 torch/CUDA 版本对齐（已手动处理时用）
  python setup_all.py --novram       # 追加 --novram（OOM 兜底，默认不加）

中断后直接重跑, 所有已完成的步骤会自动跳过。
"""
import os
import re
import sys
import time
import socket
import logging
import argparse
import subprocess
import threading
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

# ==========================================================================
# 配置区（一般只需改 NGROK_TOKEN）
# ==========================================================================

# ---- 工作区候选（自动探测，按顺序，第一个存在的胜出）----
WORKSPACE_CANDIDATES = [
    "/home/admin/workspace",
    "/mnt/workspace",
]

# ---- 派生路径（不要手改，自动跟随工作区）----
#   COMFY_DIR  = <workspace>/ComfyUI
#   MODELS_DIR = <workspace>/ComfyUI/models
#   VENV_DIR   = <workspace>/venv
COMFY_DIR = ""
MODELS_DIR = ""
VENV_DIR = ""

# ---- venv 开关 ----
# True  = 在工作区建独立环境（推荐）
#         且带 --system-site-packages，可复用【镜像预装】的 torch 等大件，
#         避免重新下载几个 GB；ComfyUI 特有依赖装进 venv（工作区，重启保留）。
# False = 直接用系统 Python（不推荐：装到系统目录的会随重启丢失）
USE_VENV = True

# ---- ngrok ----
NGROK_TOKEN = ""      # ← 填你的 ngrok authtoken；留空则不建隧道（改用平台自带代理）
NGROK_REGION = ""     # 可选: us/eu/ap/jp/au/sa/in；免费版一般不支持，留空即可

# ---- 服务端口 ----
PORT = 8188

# ---- torch / CUDA 对齐 ----
# 宿主驱动锁定 CUDA 12.8，而 PyPI 默认 wheel 是 cu130 → 跨大版本不兼容，
# 会报 "NVIDIA driver on your system is too old (found version 12080)"。
# 脚本会先检测 torch.cuda.is_available()，为 False 时才重装对齐。
# DSW 里不能升级驱动，只能改 PyTorch 这一侧。
TORCH_INDEX_URL = os.environ.get(
    "TORCH_INDEX_URL", "https://download.pytorch.org/whl/cu126"
)

# ---- 启动参数（追加到固定命令之后）----
# 固定部分已包含: --disable-pinned-memory --fast-disk --preview-method none --reserve-vram 0
EXTRA_ARGS = []

# ---- 模型下载配置 ----
BASE_URL = "https://modelscope.cn/models/Comfy-Org/MiniMax-H3/resolve/master"

# 任务列表: (相对路径, 期望大小GB)
#   expected_gb > 0 : 达到该值 95% 视为完成（精确校验）
#   expected_gb = 0 : 不校验大小，只要下载完成且文件非空（作者未公布精确大小时用）
#
# ⚠️ 已修正的原始错误：
#   1) LoRA 文件名被截断: minimax_h3_fl2v_turbo_4step_v1.safetensors
#      → 真实名:          minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors
#   2) LoRA 大小填 2GB，实际约 0.74GB → 永远判缺失
#   3) nvfp4 大小填 17.14GB，实际约 15.69GB → 永远判缺失
#
# 💡 合计约 63GB。若确定不做 R2V，可删掉 ref2va 那行，省 21GB。
JOBS = [
    ("diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors", 20.97),
    ("diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors", 20.97),
    ("text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", 0),
    ("vae/minimax_h3_video_vae_fp16.safetensors", 5.21),
    ("vae/minimax_h3_audio_vae_fp32.safetensors", 0.61),
    ("loras/minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors", 0),
]

WORKERS = 2           # 模型下载并发数（大模型建议 2，避免把解析服务打满）
RETRIES = 4           # 每个文件最大重试次数
CHUNK_SIZE = 1 << 20  # 1MB
TIMEOUT = 120         # 单次请求超时(秒)
LOG_DIR = "/tmp/comfyui_setup_logs"

# 所有 pip 操作共用的锁（pip 非线程安全，并发装同一环境会写坏）
PIP_LOCK = threading.Lock()

# ==========================================================================
# 路径探测
# ==========================================================================
def resolve_paths():
    """探测工作区并派生所有路径（保证单源，不会错位）"""
    global COMFY_DIR, MODELS_DIR, VENV_DIR

    ws = None
    for cand in WORKSPACE_CANDIDATES:
        if os.path.isdir(cand):
            ws = cand
            break
    if ws is None:
        # 都不存在：用第一个并创建
        ws = WORKSPACE_CANDIDATES[0]
        os.makedirs(ws, exist_ok=True)

    COMFY_DIR = os.path.join(ws, "ComfyUI")
    MODELS_DIR = os.path.join(COMFY_DIR, "models")   # 派生，绝不手写第二遍
    VENV_DIR = os.path.join(ws, "venv")              # 工作区内，重启保留
    return ws


# ==========================================================================
# 日志
# ==========================================================================
def setup_logger():
    os.makedirs(LOG_DIR, exist_ok=True)
    log_file = os.path.join(LOG_DIR, f"setup_{datetime.now():%Y%m%d_%H%M%S}.log")
    logger = logging.getLogger("setup")
    logger.setLevel(logging.DEBUG)

    fh = logging.FileHandler(log_file, encoding="utf-8")
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
    logger.addHandler(fh)

    ch = logging.StreamHandler(sys.stdout)
    ch.setLevel(logging.INFO)
    ch.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
    logger.addHandler(ch)

    logger.info(f"日志文件: {log_file}")
    return logger


LOG = setup_logger()


# ==========================================================================
# 通用工具
# ==========================================================================
def run_cmd(cmd, **kw):
    """执行命令，失败抛异常"""
    LOG.debug(f"$ {' '.join(cmd)}")
    return subprocess.run(cmd, check=True, **kw)


def port_open(host="127.0.0.1", port=PORT, timeout=2.0):
    """检测端口是否在监听"""
    try:
        s = socket.create_connection((host, port), timeout=timeout)
        s.close()
        return True
    except Exception:
        return False


def venv_python():
    """venv 解释器路径；未启用 venv 则用当前解释器"""
    if USE_VENV:
        return os.path.join(VENV_DIR, "bin", "python")
    return sys.executable


def venv_pip():
    """★ 统一返回列表形式的 pip 命令。

    用 <venv>/bin/python -m pip，而不是裸的 <venv>/bin/pip 脚本：
    裸 pip 脚本的 shebang 写死了创建它时的解释器，可能与 venv 实际解释器
    不一致，导致"装到全局、venv 里 import 不到、版本号一直不变"的诡异现象。
    python -m pip 保证装到当前解释器对应的 site-packages，不会串环境。
    """
    if USE_VENV:
        return [os.path.join(VENV_DIR, "bin", "python"), "-m", "pip"]
    return [sys.executable, "-m", "pip"]


# ==========================================================================
# 包管理（幂等：已装则跳过）
# ==========================================================================
def _norm(name):
    """包名规范化：去版本约束/环境标记/extras，统一小写与连字符"""
    return re.split(r"[<>=!~\[;@]", name.strip())[0].strip().lower().replace("_", "-")


def installed_set(pip_cmd):
    """返回当前环境已装包名集合"""
    try:
        r = subprocess.run(list(pip_cmd) + ["list", "--format=freeze"],
                           capture_output=True, text=True, timeout=300)
        if r.returncode != 0:
            return set()
        out = set()
        for line in r.stdout.splitlines():
            name = line.split("==")[0].split(" @")[0].strip()
            if name:
                out.add(_norm(name))
        return out
    except Exception as e:                                    # noqa: BLE001
        LOG.debug(f"pip list 失败: {e}")
        return set()


def ensure_packages(pip_cmd, packages, label=""):
    """只装缺失的包；全部已装则一次 pip 命令都不跑"""
    with PIP_LOCK:
        installed = installed_set(pip_cmd)
        missing = [p for p in packages if _norm(p) not in installed]
        if not missing:
            LOG.info(f"[跳过] 已安装: {', '.join(packages)}")
            return True
        LOG.info(f"[安装] {label or ''}{', '.join(missing)}")
        run_cmd(list(pip_cmd) + ["install", "-q"] + missing)
        LOG.info(f"[OK] {', '.join(missing)} 安装完成")
    return True


def install_requirements(pip_cmd, req_path):
    """预检 requirements.txt：全部已装则跳过整条命令，省去依赖解析时间"""
    with PIP_LOCK:
        if not os.path.exists(req_path):
            LOG.warning(f"[警告] 未找到 {req_path}，跳过依赖安装")
            return False

        installed = installed_set(pip_cmd)
        needed, total = [], 0
        for line in open(req_path, encoding="utf-8"):
            s = line.strip()
            if not s or s.startswith("#") or s.startswith("-"):
                continue                      # 跳过注释与 --index-url 等选项行
            total += 1
            if _norm(s) not in installed:
                needed.append(s)

        if total and not needed:
            LOG.info(f"[跳过] requirements.txt 的 {total} 个包已全部安装")
            return True

        LOG.info(f"[安装] requirements.txt（缺 {len(needed)}/{total} 个，首次较慢）")
        run_cmd(list(pip_cmd) + ["install", "-q", "-r", req_path])
        LOG.info("[OK] 依赖安装完成")
    return True


def torch_status(vpy):
    """返回 '版本 cuda版本 是否可用' """
    code = ("import torch;print(torch.__version__, torch.version.cuda, "
            "torch.cuda.is_available())")
    try:
        r = subprocess.run([vpy, "-c", code], capture_output=True,
                           text=True, timeout=300)
        return r.stdout.strip()
    except Exception as e:                                    # noqa: BLE001
        return f"<检测失败: {e}>"


def align_torch(pip_cmd):
    """torch 与宿主驱动 CUDA 版本对齐（能正常用 CUDA 则跳过）"""
    vpy = venv_python()
    out = torch_status(vpy)
    LOG.info(f"[检测] torch: {out or '<无输出>'}")

    if out.endswith("True"):
        LOG.info("[跳过] torch 可正常调用 CUDA，无需对齐")
        return True

    LOG.warning("[对齐] torch 与驱动 CUDA 不匹配，重装匹配版本")
    LOG.warning(f"       源: {TORCH_INDEX_URL}")
    with PIP_LOCK:
        run_cmd(list(pip_cmd) + ["install", "--force-reinstall",
                                 "torch", "torchvision", "torchaudio",
                                 "--index-url", TORCH_INDEX_URL])
    LOG.info(f"[检测] 对齐后: {torch_status(vpy)}")
    return True


# ==========================================================================
# 任务一：ComfyUI 克隆 + 依赖（与 venv、模型下载并行）
# ==========================================================================
def task_clone_and_install(venv_ready):
    LOG.info("=" * 60)
    LOG.info("[安装] ComfyUI 部署开始")
    LOG.info("=" * 60)

    # 1. 克隆（有 main.py 即视为已装，无论是否 git 仓库）
    main_py = os.path.join(COMFY_DIR, "main.py")
    if os.path.isfile(main_py):
        LOG.info(f"[跳过] ComfyUI 已存在: {COMFY_DIR}")
    elif os.path.isdir(COMFY_DIR):
        LOG.info(f"[跳过] 目录已存在(非 git 仓库): {COMFY_DIR}")
    else:
        LOG.info(f"[克隆] ComfyUI → {COMFY_DIR}")
        run_cmd(["git", "clone",
                 "https://github.com/comfyanonymous/ComfyUI.git", COMFY_DIR])
        LOG.info("[OK] 克隆完成")

    # 2. 等 venv 建好（venv 与克隆并行，这里通常几秒内就放行）
    if USE_VENV and not venv_ready.wait(timeout=600):
        LOG.error("[错误] 等待 venv 超时")
        return False

    # 3. 依赖
    install_requirements(venv_pip(), os.path.join(COMFY_DIR, "requirements.txt"))
    LOG.info("[安装] ComfyUI 部署完成 ✅")
    return True


# ==========================================================================
# 任务二：venv + pyngrok（与克隆、模型下载并行）
# ==========================================================================
def task_venv_and_pyngrok(venv_ready):
    try:
        if not USE_VENV:
            LOG.info("[跳过] 未启用 venv，使用系统 Python")
            return True

        vpy = os.path.join(VENV_DIR, "bin", "python")
        if os.path.exists(vpy):
            LOG.info(f"[跳过] venv 已存在: {VENV_DIR}")
        else:
            LOG.info(f"[创建] venv → {VENV_DIR}")
            # ★ 关键：加 --system-site-packages
            #   让 venv 能看见【镜像预装】的包（torch 等大件）。
            #   镜像层的包重启后仍在，可安全复用 —— 省下重新下载几个 GB 的 torch。
            #
            #   requirements.txt 里 torch/torchsde/torchvision/torchaudio 都是裸名
            #   （无版本约束），pip 会判定 already satisfied 直接跳过，零下载。
            #   ⚠️ 但这也意味着：镜像若预装的是 cu130 而驱动是 12.8，
            #      torch 会一直不可用 —— 所以阶段 B 有一道 torch 对齐兜底。
            #
            #   而 ComfyUI 特有的新依赖会装进 venv 自己的 site-packages
            #   （位于工作区，重启保留），不会污染系统目录。
            #
            #   注意：只有【镜像预装】的包可这样复用；
            #        你自己 pip 装到系统目录的包仍会随重启丢失。
            run_cmd([sys.executable, "-m", "venv",
                     "--system-site-packages", VENV_DIR])
            LOG.info("[OK] venv 创建完成（复用镜像预装包模式）")
    finally:
        # 无论成功失败都放行，避免其余任务死锁
        venv_ready.set()

    # pyngrok 装进 venv（工作区持久），而不是系统目录
    ensure_packages(venv_pip(), ["pyngrok"], label="pyngrok: ")
    return True


# ==========================================================================
# 任务三：模型下载
# ==========================================================================
def _stream_download(url, path, done, desc, pbar, requests):
    """从 done 字节处续传"""
    headers = {"Range": f"bytes={done}-"} if done else {}
    r = requests.get(url, headers=headers, stream=True, timeout=TIMEOUT)
    if r.status_code == 200 and done:
        LOG.warning(f"{desc} 服务端不支持 Range，从头下载")
        done = 0
    r.raise_for_status()

    cl = r.headers.get("Content-Length")
    if cl:
        pbar.total = done + int(cl)
        pbar.refresh()

    mode = "ab" if done else "wb"
    with open(path, mode) as f:
        for chunk in r.iter_content(CHUNK_SIZE):
            if not chunk:
                continue
            f.write(chunk)
            done += len(chunk)
            pbar.update(len(chunk))
    return done


def download_one(job):
    """下载单个文件（带重试 + 断点续传）"""
    import requests
    from tqdm import tqdm

    rel, exp_gb = job
    exp_bytes = int(exp_gb * 1e9)
    url = f"{BASE_URL}/{rel}"
    path = os.path.join(MODELS_DIR, rel)
    os.makedirs(os.path.dirname(path), exist_ok=True)

    # 完成判定
    #  ★ 修正：exp_gb=0 时必须要求文件非空。
    #    原写法 `size >= 0*0.95` 恒为真 → 0 字节空文件也会被判为"已完成"
    if os.path.exists(path):
        size = os.path.getsize(path)
        if exp_bytes > 0 and size >= exp_bytes * 0.95:
            LOG.info(f"[跳过] {rel}（已存在且达到预期大小）")
            return f"[跳过] {rel}"
        if exp_bytes == 0 and size > 0:
            LOG.info(f"[跳过] {rel}（已存在，未设大小校验）")
            return f"[跳过] {rel}"

    desc = os.path.basename(rel)
    for attempt in range(1, RETRIES + 1):
        current = os.path.getsize(path) if os.path.exists(path) else 0
        with tqdm(
            total=exp_bytes if exp_bytes > 0 else None,
            initial=current,
            unit="B", unit_scale=True,
            desc=f"{desc:44s}",
            miniters=1, leave=True,
        ) as pbar:
            try:
                _stream_download(url, path, current, desc, pbar, requests)
                final = os.path.getsize(path)
                pct = f"{final/exp_bytes*100:.1f}%" if exp_bytes > 0 else "未校验"
                LOG.info(f"[完成] {rel}: {final/1e9:.2f} GB ({pct})")
                return f"[完成] {rel}: {final/1e9:.2f} GB"
            except Exception as e:                            # noqa: BLE001
                LOG.error(f"[重试 {attempt}/{RETRIES}] {rel}: {e}")
                if attempt < RETRIES:
                    time.sleep(min(2 ** attempt, 30))

    return f"[失败] {rel}: 超过最大重试次数"


def task_models(venv_ready):
    LOG.info("=" * 60)
    LOG.info("[下载] 模型下载开始")
    LOG.info("=" * 60)
    os.makedirs(MODELS_DIR, exist_ok=True)

    # 等 venv，用 venv 的 pip 装下载器（装进工作区，重启保留）
    if USE_VENV and not venv_ready.wait(timeout=600):
        LOG.error("[错误] 等待 venv 超时")
        return False

    ensure_packages(venv_pip(), ["requests", "tqdm"], label="下载器: ")

    results = []
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        for fut in ex.map(download_one, JOBS):
            results.append(fut)

    # 核对
    LOG.info("=" * 60)
    LOG.info("最终文件核对:")
    all_ok = True
    for rel, exp_gb in JOBS:
        exp_bytes = int(exp_gb * 1e9)
        p = os.path.join(MODELS_DIR, rel)
        if os.path.exists(p):
            size = os.path.getsize(p)
            ok = (size >= exp_bytes * 0.95) if exp_bytes > 0 else (size > 0)
            if ok:
                LOG.info(f"  [OK]   {rel}: {size/1e9:.2f} GB")
            else:
                all_ok = False
                LOG.warning(f"  [缺失] {rel}: {size/1e9:.2f} GB")
        else:
            all_ok = False
            LOG.warning(f"  [缺失] {rel}: 文件不存在")

    if all_ok:
        LOG.info("全部模型下载完成 ✅")
    else:
        LOG.warning("部分模型缺失，重跑本脚本可自动续传")
    return all_ok


# ==========================================================================
# 启动服务 + ngrok 隧道
# ==========================================================================
def launch_comfyui(extra_args):
    """后台启动 ComfyUI"""
    LOG.info("=" * 60)
    LOG.info("[启动] ComfyUI")
    LOG.info("=" * 60)
    log_path = os.path.join(LOG_DIR, "comfyui.log")
    
    if port_open():
        LOG.info(f"[跳过] 端口 {PORT} 已在监听（服务可能已运行）")
        LOG.info(f"可查看日志: {log_path}")
        return None

    vpy = venv_python()
    cmd = [
        vpy, "main.py",
        "--listen", "0.0.0.0",
        "--port", str(PORT),
        "--enable-manager",
        # ↓ 内存优化：这两个是主力的"用磁盘/时间换 RAM"，
        #   实测可把系统内存占用从 80%+ 压到 30% 左右
        "--disable-pinned-memory",
        "--fast-disk",
        # ↓ 关掉生成中预览，省显存
        "--preview-method", "none",
        # ↓ 默认会把 400MB+800MB 显存预留出去；设 0 可白捡约 1.2GB
        "--reserve-vram", "0",
    ] + list(extra_args)

    LOG.info(f"$ {' '.join(cmd)}")

    
    lf = open(log_path, "w", encoding="utf-8")
    proc = subprocess.Popen(cmd, cwd=COMFY_DIR, stdout=lf, stderr=subprocess.STDOUT)
    LOG.info(f"[信息] ComfyUI 已后台启动 (PID {proc.pid})，日志: {log_path}")

    # 等待就绪
    LOG.info(f"[等待] 等待 {PORT} 就绪（最多 300 秒）")
    for i in range(300):
        if port_open():
            LOG.info(f"[OK] ComfyUI 就绪（耗时 {i} 秒）")
            return proc
        time.sleep(1)

    LOG.error(f"[错误] ComfyUI 300 秒内未就绪，请查看日志: {log_path}")
    return proc


def build_tunnel():
    """建立 ngrok 隧道（token 为空则跳过）"""
    if not NGROK_TOKEN:
        LOG.info("[跳过] 未配置 NGROK_TOKEN，不建隧道")
        LOG.info("       如需公网访问，可用平台自带代理，或在脚本顶部填 NGROK_TOKEN")
        return None

    LOG.info("=" * 60)
    LOG.info("[隧道] 建立 ngrok 隧道")
    LOG.info("=" * 60)

    try:
        from pyngrok import ngrok
    except ImportError:
        ensure_packages(venv_pip(), ["pyngrok"], label="pyngrok: ")
        from pyngrok import ngrok

    ngrok.set_auth_token(NGROK_TOKEN)

    # addr 强制 IPv4，规避 ngrok 把 localhost 解析成 [::1] 导致 connection refused
    # 不传 region，规避新版 http 隧道 "field region not found" 400 报错
    addr = f"127.0.0.1:{PORT}"
    tunnel = ngrok.connect(addr=addr, bind_tls=True)
    url = tunnel.public_url

    LOG.info("")
    LOG.info("=" * 60)
    LOG.info(f"  公网链接 : {url}")
    LOG.info(f"  本地端口 : {addr}")
    LOG.info("=" * 60)
    LOG.info("")
    return url


# ==========================================================================
# 入口
# ==========================================================================
def main():
    parser = argparse.ArgumentParser(
        description="ComfyUI + MiniMax-H3 + ngrok 一键部署")
    parser.add_argument("--no-tunnel", action="store_true", help="不建 ngrok 隧道")
    parser.add_argument("--no-launch", action="store_true",
                        help="不启动 ComfyUI（只安装+下载）")
    parser.add_argument("--no-align", action="store_true",
                        help="跳过 torch/CUDA 版本对齐")
    parser.add_argument("--novram", action="store_true",
                        help="追加 --novram（OOM 兜底；注意它会关掉 DynamicVRAM）")
    args = parser.parse_args()

    ws = resolve_paths()

    LOG.info("=" * 60)
    LOG.info(" ComfyUI + MiniMax-H3 一键部署")
    LOG.info("=" * 60)
    LOG.info(f" 工作区    : {ws}")
    LOG.info(f" ComfyUI   : {COMFY_DIR}")
    LOG.info(f" 模型目录  : {MODELS_DIR}")
    LOG.info(f" venv      : {VENV_DIR if USE_VENV else '(未启用，使用系统 Python)'}")
    LOG.info(f" torch源   : {TORCH_INDEX_URL}")
    LOG.info(f" 端口      : {PORT}")
    LOG.info("=" * 60)

    # pip 缓存放工作区：万一要重装，缓存命中可省大量下载
    os.environ["PIP_CACHE_DIR"] = os.path.join(ws, ".cache", "pip")
    os.makedirs(os.environ["PIP_CACHE_DIR"], exist_ok=True)
    LOG.info(f"[信息] pip 缓存目录: {os.environ['PIP_CACHE_DIR']}")

    # ---- 阶段 A：三路并行 ----
    LOG.info("")
    LOG.info("[并行] 克隆ComfyUI ‖ 建venv+装pyngrok ‖ 下载模型")
    venv_ready = threading.Event()
    with ThreadPoolExecutor(max_workers=3) as ex:
        futures = {
            "ComfyUI 安装": ex.submit(task_clone_and_install, venv_ready),
            "venv+pyngrok": ex.submit(task_venv_and_pyngrok, venv_ready),
            "模型下载":     ex.submit(task_models, venv_ready),
        }
        for name, f in futures.items():
            try:
                f.result()
            except Exception as e:                            # noqa: BLE001
                LOG.error(f"[错误] {name} 失败: {e}")

    if not os.path.isfile(os.path.join(COMFY_DIR, "main.py")):
        LOG.error("[中止] ComfyUI 异常，未找到 main.py")
        return

    # ---- 阶段 B：启动前必备 ----
    pip_cmd = venv_pip()

    # comfyui-manager（工作流节点缺失时靠它补）
    ensure_packages(pip_cmd, ["comfyui-manager"], label="comfyui-manager: ")

    # torch/CUDA 对齐（能正常用 CUDA 则自动跳过）
    if not args.no_align:
        align_torch(pip_cmd)

    LOG.info("")
    LOG.info("=" * 60)
    LOG.info(" 安装完成 ✅")
    LOG.info("=" * 60)

    if args.no_launch:
        vpy = venv_python()
        LOG.info("[完成] --no-launch 模式，不启动服务")
        LOG.info(f"       手动启动: cd {COMFY_DIR} && {vpy} main.py "
                 "--listen 0.0.0.0 --port " + str(PORT))
        return

    # ---- 阶段 C：启动 + 隧道 ----
    extra = list(EXTRA_ARGS)
    if args.novram:
        extra.append("--novram")

    launch_comfyui(extra)

    if not args.no_tunnel:
        build_tunnel()

    LOG.info("[提示] 服务运行中，Ctrl+C 退出")
    try:
        while True:
            time.sleep(60)
    except KeyboardInterrupt:
        LOG.info("[退出] 已停止")


if __name__ == "__main__":
    main()


# python3 setup_all.py               # 全自动
# python3 setup_all.py --no-tunnel   # 不进行穿透
# python3 setup_all.py --no-launch   # 只安装和下载，不启动服务
# python3 setup_all.py --no-align    # 跳过 torch/CUDA 对齐
# python3 setup_all.py --novram      # 被 OOM kill 时追加更激进的卸载模式