# ============================================================
# Cell 1
# 环境准备 + 模型 Dataset 挂载/解压 + 安装依赖 + 启动 ComfyUI 本地 API
# 不生图，不加载参考脸
# ============================================================

import os
import sys
import json
import time
import shutil
import socket
import zipfile
import subprocess
import threading
import textwrap
import urllib.request
from pathlib import Path


# ============================================================
# 0. 基础路径和固定参数
# ============================================================

ROOT = Path("/kaggle/working")
INPUT_ROOT = Path("/kaggle/input")

COMFY_DIR = ROOT / "ComfyUI"
MODELS_DIR = COMFY_DIR / "models"

CHECKPOINT_DIR = MODELS_DIR / "checkpoints"
IPADAPTER_DIR = MODELS_DIR / "ipadapter"
LORAS_DIR = MODELS_DIR / "loras"
CLIP_VISION_DIR = MODELS_DIR / "clip_vision"
UNET_DIR = MODELS_DIR / "unet"
CLIP_DIR = MODELS_DIR / "clip"
VAE_DIR = MODELS_DIR / "vae"
PULID_DIR = MODELS_DIR / "pulid"

CUSTOM_NODES_DIR = COMFY_DIR / "custom_nodes"
INPUT_DIR = COMFY_DIR / "input"
OUTPUT_DIR = COMFY_DIR / "output"

PORT = int(os.environ.get("COMFY_PORT", "8188"))
COMFY_URL = f"http://127.0.0.1:{PORT}"
MULTI_GPU_WORKERS = os.environ.get("KAGGLE_MULTI_GPU_WORKERS", "auto").strip().lower()
MAX_COMFY_INSTANCES = int(os.environ.get("KAGGLE_MAX_COMFY_INSTANCES", "2"))

# 如需完全固定环境版本，把下面两个值改成你验证过的 commit。
# 默认 None 表示使用 git clone 当时的最新版本。
COMFYUI_COMMIT = None
IPADAPTER_PLUS_COMMIT = None

RUNTIME_VERSIONS_PATH = ROOT / "runtime_versions.json"
COMFY_INSTANCES_PATH = ROOT / "comfy_instances.json"


# ============================================================
# 工具函数
# ============================================================

def show_step(title):
    print("\n" + "=" * 90)
    print(title)
    print("=" * 90)


def run_cmd(cmd, cwd=None, check=True):
    print("\n" + "=" * 90)
    print("[CMD]", cmd)
    if cwd is not None:
        print("[CWD]", cwd)
    print("=" * 90)

    env = os.environ.copy()
    env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
    env["PYTHONUNBUFFERED"] = "1"
    env["MAX_JOBS"] = "2"

    p = subprocess.Popen(
        cmd,
        shell=True,
        cwd=str(cwd) if cwd is not None else None,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    for line in p.stdout:
        print(line, end="")

    p.wait()
    print("\n[RETURN CODE]", p.returncode)

    if check and p.returncode != 0:
        raise RuntimeError(f"命令失败: {cmd}")

    return p.returncode


def ensure_dirs():
    for d in [
        ROOT,
        COMFY_DIR,
        MODELS_DIR,
        CHECKPOINT_DIR,
        IPADAPTER_DIR,
        LORAS_DIR,
        CLIP_VISION_DIR,
        UNET_DIR,
        CLIP_DIR,
        VAE_DIR,
        PULID_DIR,
        CUSTOM_NODES_DIR,
        INPUT_DIR,
        OUTPUT_DIR,
    ]:
        d.mkdir(parents=True, exist_ok=True)


def set_stable_env():
    os.environ["PYTHONHASHSEED"] = "0"
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "0"
    os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
    os.environ["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"


def port_open(host="127.0.0.1", port=8188):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(1)
    try:
        s.connect((host, port))
        s.close()
        return True
    except Exception:
        return False


def stop_existing_comfyui(port=PORT):
    if not port_open("127.0.0.1", port):
        return
    print("[INFO] 检测到已有 ComfyUI 进程，重启以加载最新 custom_nodes")
    run_cmd(f"pkill -f 'main.py.*--port {port}' || true", check=False)
    start = time.time()
    while time.time() - start < 30:
        if not port_open("127.0.0.1", port):
            print("[OK] 旧 ComfyUI 进程已停止")
            return
        time.sleep(1)
    raise RuntimeError("旧 ComfyUI 进程未能停止。请在 Kaggle 停止运行中的 Cell 1/Cell 3 后重试。")


def comfy_object_info(timeout=30, base_url=COMFY_URL):
    with urllib.request.urlopen(f"{base_url}/object_info", timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def verify_comfyui_nodes(base_url=COMFY_URL):
    start = time.time()
    timeout = 120
    last_error = None
    object_info = {}
    while time.time() - start < timeout:
        try:
            object_info = comfy_object_info(timeout=10, base_url=base_url)
            break
        except Exception as exc:
            last_error = exc
            time.sleep(2)
    else:
        raise RuntimeError(f"ComfyUI object_info 不可访问: {last_error!r}")

    required = []
    if "Flux" in usable_profiles:
        required.extend(["UNETLoader", "DualCLIPLoader", "VAELoader"])
        if PULID_NODE_DIR.exists():
            required.extend([
                "PulidFluxModelLoader",
                "PulidFluxFaceNetLoader",
                "PulidFluxEvaClipLoader",
                "ApplyPulidFlux",
            ])

    missing = [name for name in required if name not in object_info]
    if missing:
        print("[ERROR] ComfyUI object_info 缺少节点:", ", ".join(missing))
        print("[INFO] 请查看上方 [ComfyUI] 日志里的 custom_nodes import error。")
        raise RuntimeError("ComfyUI 节点加载失败: " + ", ".join(missing))

    if required:
        print("[OK] ComfyUI required nodes loaded:", base_url, "|", ", ".join(required))


def cuda_device_count():
    try:
        import torch
        return torch.cuda.device_count() if torch.cuda.is_available() else 0
    except Exception:
        return 0


def desired_comfy_instances():
    gpu_count = cuda_device_count()
    if MULTI_GPU_WORKERS in {"0", "false", "no", "off", "single"}:
        target = 1
    elif MULTI_GPU_WORKERS == "auto":
        target = gpu_count if gpu_count > 1 else 1
    else:
        try:
            target = int(MULTI_GPU_WORKERS)
        except ValueError:
            target = 1

    target = max(1, min(target, max(1, MAX_COMFY_INSTANCES)))
    if gpu_count:
        target = min(target, gpu_count)

    instances = []
    for index in range(target):
        port = PORT + index
        gpu_index = index if gpu_count else None
        name = f"gpu{index}" if gpu_index is not None else "main"
        instances.append({
            "name": name,
            "port": port,
            "url": f"http://127.0.0.1:{port}",
            "gpu_index": gpu_index,
            "worker_id": f"kaggle-t4-main-{name}",
        })
    return instances


def write_comfy_instances(instances):
    COMFY_INSTANCES_PATH.write_text(json.dumps(instances, ensure_ascii=False, indent=2), encoding="utf-8")
    print("[OK] comfy_instances.json:", COMFY_INSTANCES_PATH)
    print(json.dumps(instances, ensure_ascii=False, indent=2))


def file_ok(path, min_mb=1):
    path = Path(path)
    return path.exists() and path.stat().st_size > min_mb * 1024 * 1024


def git_commit(path):
    path = Path(path)
    if not (path / ".git").exists():
        return None

    try:
        return subprocess.check_output(
            ["git", "rev-parse", "HEAD"],
            cwd=str(path),
            text=True,
        ).strip()
    except Exception:
        return None


def checkout_git_commit(path, commit, name):
    if not commit:
        return

    path = Path(path)
    print(f"[INFO] 固定 {name} 到 commit:", commit)
    run_cmd(f"git fetch --depth 1 origin {commit}", cwd=path, check=False)
    run_cmd(f"git checkout {commit}", cwd=path, check=True)


def find_input_archive(file_name):
    if not INPUT_ROOT.exists():
        return None
    for path in INPUT_ROOT.rglob(file_name):
        if path.is_file() and path.stat().st_size > 100 * 1024:
            return path
    return None


def extract_archive_root(zip_path, target_dir, expected_file):
    zip_path = Path(zip_path)
    target_dir = Path(target_dir)
    tmp_dir = target_dir.with_name(target_dir.name + "_extract_tmp")
    if tmp_dir.exists():
        shutil.rmtree(tmp_dir)
    tmp_dir.mkdir(parents=True, exist_ok=True)

    print("[UNZIP]", zip_path, "->", tmp_dir)
    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(tmp_dir)

    candidates = [p for p in tmp_dir.iterdir() if p.is_dir()]
    source_dir = candidates[0] if len(candidates) == 1 else tmp_dir
    if not (source_dir / expected_file).exists():
        raise RuntimeError(f"压缩包结构不符合预期，找不到 {expected_file}: {zip_path}")

    if target_dir.exists():
        shutil.rmtree(target_dir)
    shutil.move(str(source_dir), str(target_dir))
    if tmp_dir.exists():
        shutil.rmtree(tmp_dir)
    print("[OK] extracted:", target_dir)


def save_runtime_versions(extra=None):
    info = {
        "python": sys.version,
        "desired_comfyui_commit": COMFYUI_COMMIT,
        "desired_ipadapter_plus_commit": IPADAPTER_PLUS_COMMIT,
        "comfyui_commit": git_commit(COMFY_DIR),
        "ipadapter_plus_commit": git_commit(IPADAPTER_NODE_DIR),
    }
    if extra:
        info.update(extra)

    try:
        import torch
        info["torch"] = torch.__version__
        info["cuda_available"] = torch.cuda.is_available()
        info["gpu_count"] = torch.cuda.device_count() if torch.cuda.is_available() else 0
        info["gpu"] = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
        info["gpus"] = [
            torch.cuda.get_device_name(index)
            for index in range(torch.cuda.device_count())
        ] if torch.cuda.is_available() else []
    except Exception as e:
        info["torch_error"] = repr(e)

    with open(RUNTIME_VERSIONS_PATH, "w", encoding="utf-8") as f:
        json.dump(info, f, ensure_ascii=False, indent=2)

    print("[OK] runtime_versions.json:", RUNTIME_VERSIONS_PATH)
    print(json.dumps(info, ensure_ascii=False, indent=2))


# ============================================================
# 1. 初始化和 GPU 检查
# ============================================================

show_step("1. 初始化和 GPU 检查")

set_stable_env()
ensure_dirs()

try:
    import torch
    print("Python:", sys.version)
    print("Torch:", torch.__version__)
    print("CUDA available:", torch.cuda.is_available())

    if torch.cuda.is_available():
        print("GPU:", torch.cuda.get_device_name(0))
    else:
        raise RuntimeError("CUDA 不可用。请先在 Kaggle Settings -> Accelerator 选择 GPU。")

except Exception as e:
    raise RuntimeError("GPU / Torch 检查失败") from e


# ============================================================
# 2. 克隆 ComfyUI 和 IPAdapter 节点
# ============================================================

show_step("2. 克隆 ComfyUI 和 IPAdapter 节点")

if not (COMFY_DIR / "main.py").exists():
    if COMFY_DIR.exists() and not (COMFY_DIR / "main.py").exists():
        print("[WARN] 发现不完整 ComfyUI，删除:", COMFY_DIR)
        shutil.rmtree(COMFY_DIR)

    comfy_zip = find_input_archive("ComfyUI-master.zip")
    if comfy_zip:
        extract_archive_root(comfy_zip, COMFY_DIR, "main.py")
    else:
        run_cmd(
            f"git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git {COMFY_DIR}",
            check=True
        )
else:
    print("[SKIP] ComfyUI 已存在:", COMFY_DIR)

checkout_git_commit(COMFY_DIR, COMFYUI_COMMIT, "ComfyUI")

ensure_dirs()

IPADAPTER_NODE_DIR = CUSTOM_NODES_DIR / "ComfyUI_IPAdapter_plus"

if not IPADAPTER_NODE_DIR.exists():
    ipadapter_zip = find_input_archive("ComfyUI_IPAdapter_plus-main.zip")
    if ipadapter_zip:
        extract_archive_root(ipadapter_zip, IPADAPTER_NODE_DIR, "README.md")
    else:
        run_cmd(
            f"git clone --depth 1 https://github.com/cubiq/ComfyUI_IPAdapter_plus.git {IPADAPTER_NODE_DIR}",
            check=True
        )
else:
    print("[SKIP] ComfyUI_IPAdapter_plus 已存在:", IPADAPTER_NODE_DIR)

checkout_git_commit(IPADAPTER_NODE_DIR, IPADAPTER_PLUS_COMMIT, "ComfyUI_IPAdapter_plus")

PULID_NODE_DIR = CUSTOM_NODES_DIR / "ComfyUI_PuLID_Flux_ll"
if not PULID_NODE_DIR.exists():
    pulid_zip = find_input_archive("ComfyUI_PuLID_Flux_ll-main.zip")
    if pulid_zip:
        extract_archive_root(pulid_zip, PULID_NODE_DIR, "requirements.txt")
    else:
        print("[INFO] 未在 Dataset 中找到 PuLID Flux 节点 zip，改为从 GitHub clone。")
        run_cmd(
            f"git clone --depth 1 https://github.com/lldacing/ComfyUI_PuLID_Flux_ll.git {PULID_NODE_DIR}",
            check=True,
        )
else:
    print("[SKIP] ComfyUI_PuLID_Flux_ll 已存在:", PULID_NODE_DIR)

save_runtime_versions()


# ============================================================
# 3. 准备模型：只使用 Kaggle Input Dataset
# 支持：
# 1. /kaggle/input/**/models/
# 2. /kaggle/input/**/models.zip
# 不自动从 HuggingFace 下载，避免误判后重下
# ============================================================

show_step("3. 准备模型")

SD15_MODEL_SPECS = [
    {
        "name": "SD1.5 checkpoint",
        "relative": Path("checkpoints") / "realisticVisionV60B1_v60B1VAE.safetensors",
        "target": CHECKPOINT_DIR / "realisticVisionV60B1_v60B1VAE.safetensors",
        "min_mb": 1800,
    },
    {
        "name": "SD1.5 FaceID IPAdapter",
        "relative": Path("ipadapter") / "ip-adapter-faceid-plusv2_sd15.bin",
        "target": IPADAPTER_DIR / "ip-adapter-faceid-plusv2_sd15.bin",
        "min_mb": 100,
    },
    {
        "name": "SD1.5 FaceID LoRA",
        "relative": Path("loras") / "ip-adapter-faceid-plusv2_sd15_lora.safetensors",
        "target": LORAS_DIR / "ip-adapter-faceid-plusv2_sd15_lora.safetensors",
        "min_mb": 30,
    },
    {
        "name": "clip_vision",
        "relative": Path("clip_vision") / "CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors",
        "target": CLIP_VISION_DIR / "CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors",
        "min_mb": 1000,
    },
]

SDXL_MODEL_SPECS = [
    {
        "name": "SDXL checkpoint",
        "relative": Path("checkpoints") / "realvisxlV50_v50Bakedvae.safetensors",
        "target": CHECKPOINT_DIR / "realvisxlV50_v50Bakedvae.safetensors",
        "min_mb": 1800,
    },
    {
        "name": "SDXL FaceID IPAdapter",
        "relative": Path("ipadapter") / "ipAdapterFaceid_faceidPlusv2Sdxl.bin",
        "target": IPADAPTER_DIR / "ipAdapterFaceid_faceidPlusv2Sdxl.bin",
        "min_mb": 1000,
    },
    {
        "name": "CLIP Vision",
        "relative": Path("clip_vision") / "CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors",
        "target": CLIP_VISION_DIR / "CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors",
        "min_mb": 1000,
    },
]

FLUX_CORE_MODEL_SPECS = [
    {
        "name": "FLUX UNET FP8",
        "relative": Path("unet") / "flux1-dev-fp8.safetensors",
        "target": UNET_DIR / "flux1-dev-fp8.safetensors",
        "min_mb": 10000,
        "aliases": [
            "flux1-dev-fp8.safetensors",
            "flux1-dev-fp8-e4m3fn.safetensors",
            "flux1-dev-fp8_e4m3fn.safetensors",
        ],
    },
    {
        "name": "FLUX CLIP-L",
        "relative": Path("clip") / "clip_l.safetensors",
        "target": CLIP_DIR / "clip_l.safetensors",
        "min_mb": 100,
    },
    {
        "name": "FLUX T5 XXL FP8",
        "relative": Path("clip") / "t5xxl_fp8_e4m3fn.safetensors",
        "target": CLIP_DIR / "t5xxl_fp8_e4m3fn.safetensors",
        "min_mb": 4000,
        "aliases": [
            "t5xxl_fp8_e4m3fn.safetensors",
            "t5xxl_fp8_e4m3fn_scaled.safetensors",
        ],
    },
    {
        "name": "FLUX AE / VAE",
        "relative": Path("vae") / "ae.safetensors",
        "target": VAE_DIR / "ae.safetensors",
        "min_mb": 250,
    },
]

MODEL_PROFILES = {
    "SD1.5": SD15_MODEL_SPECS,
    "SDXL": SDXL_MODEL_SPECS,
    "Flux": FLUX_CORE_MODEL_SPECS,
}

MODEL_SUBDIR_TARGETS = {
    "checkpoints": CHECKPOINT_DIR,
    "ipadapter": IPADAPTER_DIR,
    "loras": LORAS_DIR,
    "clip_vision": CLIP_VISION_DIR,
    "unet": UNET_DIR,
    "clip": CLIP_DIR,
    "vae": VAE_DIR,
    "pulid": PULID_DIR,
}

OPTIONAL_MODEL_HINTS = [
    {
        "label": "SDXL FaceID IPAdapter",
        "target": IPADAPTER_DIR / "ipAdapterFaceid_faceidPlusv2Sdxl.bin",
        "min_mb": 1000,
    },
    {
        "label": "SDXL FaceID LoRA",
        "target": LORAS_DIR / "ip-adapter-faceid-plusv2_sdxl_lora.safetensors",
        "min_mb": 30,
    },
    {
        "label": "FLUX PuLID v0.9.1",
        "target": PULID_DIR / "pulid_flux_v0.9.1.safetensors",
        "min_mb": 200,
    },
]


def link_or_copy(src, dst):
    src = Path(src)
    dst = Path(dst)

    if not src.exists():
        raise FileNotFoundError(f"源文件不存在: {src}")

    dst.parent.mkdir(parents=True, exist_ok=True)

    if dst.exists() or dst.is_symlink():
        try:
            if dst.exists() and dst.stat().st_size > 1024 * 1024:
                print("[SKIP] 已存在:", dst)
                return
            dst.unlink()
        except Exception:
            pass

    try:
        os.symlink(src, dst)
        print("[LINK]", dst)
        print("  ->  ", src)
    except Exception as e:
        print("[WARN] 软链接失败，改为复制:", e)
        shutil.copy2(src, dst)
        print("[COPY]", src)
        print("  ->  ", dst)


def merge_all_model_datasets(verbose=True):
    if not INPUT_ROOT.exists():
        return 0

    linked = 0
    seen = set()
    for models_root in INPUT_ROOT.rglob("models"):
        if not models_root.is_dir():
            continue
        for subdir, target_dir in MODEL_SUBDIR_TARGETS.items():
            source_dir = models_root / subdir
            if not source_dir.is_dir():
                continue
            for src in sorted(source_dir.iterdir()):
                if not src.is_file():
                    continue
                if src.suffix.lower() not in {".safetensors", ".bin", ".pt", ".pth", ".ckpt"}:
                    continue
                key = (subdir, src.name)
                if key in seen:
                    continue
                seen.add(key)
                dst = target_dir / src.name
                if dst.exists() or dst.is_symlink():
                    continue
                link_or_copy(src, dst)
                linked += 1

    # Kaggle Dataset 不一定保留 models/<subdir> 层级。按已知文件名再做一次
    # 全盘扫描，并把常见 FLUX FP8 文件名变体链接到 Worker 使用的规范名称。
    all_specs = [spec for specs in MODEL_PROFILES.values() for spec in specs]
    for spec in all_specs:
        target = spec["target"]
        if file_ok(target, spec["min_mb"]):
            continue
        names = spec.get("aliases") or [spec["relative"].name]
        candidates = []
        for name in names:
            candidates.extend(
                path for path in INPUT_ROOT.rglob(name)
                if path.is_file() and path.resolve() != target.resolve()
            )
        candidates.sort(key=lambda path: path.stat().st_size, reverse=True)
        source = next(
            (path for path in candidates if path.stat().st_size >= spec["min_mb"] * 1024 * 1024),
            None,
        )
        if source is not None:
            link_or_copy(source, target)
            linked += 1
            if source.name != target.name:
                print(f"[ALIAS] {source.name} -> {target.name}")

    if verbose:
        print(f"[OK] Model Dataset 合并完成: {linked} 个文件")
    return linked


def model_profile_ok(specs):
    return all(file_ok(spec["target"], spec["min_mb"]) for spec in specs)


def available_model_profiles():
    return [name for name, specs in MODEL_PROFILES.items() if model_profile_ok(specs)]


def print_model_profile_status():
    for profile_name, specs in MODEL_PROFILES.items():
        print(f"\n{profile_name} 文件检查：")
        for spec in specs:
            p = spec["target"]
            min_mb = spec["min_mb"]
            actual_mb = round(p.stat().st_size / 1024 / 1024, 2) if p.exists() else 0
            if file_ok(p, min_mb):
                print("[OK]", spec["name"], "|", p.name, "|", actual_mb, "MB")
            else:
                print("[MISS]", spec["name"], "|", p.name, "|", actual_mb, "MB | min", min_mb, "MB")


def model_asset_report():
    profiles = {}
    for profile_name, specs in MODEL_PROFILES.items():
        files = []
        missing = []
        for spec in specs:
            p = spec["target"]
            actual_mb = round(p.stat().st_size / 1024 / 1024, 2) if p.exists() else 0
            ok = file_ok(p, spec["min_mb"])
            item = {
                "name": spec["name"],
                "file": p.name,
                "relative": str(spec["relative"]).replace("\\", "/"),
                "size_mb": actual_mb,
                "min_mb": spec["min_mb"],
                "ok": ok,
            }
            files.append(item)
            if not ok:
                missing.append(item["relative"])
        profiles[profile_name] = {
            "available": not missing,
            "missing": missing,
            "files": files,
        }

    optional = []
    for item in OPTIONAL_MODEL_HINTS:
        p = item["target"]
        optional.append({
            "name": item["label"],
            "file": p.name,
            "size_mb": round(p.stat().st_size / 1024 / 1024, 2) if p.exists() else 0,
            "min_mb": item["min_mb"],
            "ok": file_ok(p, item["min_mb"]),
        })

    available_checkpoints = sorted({
        p.name
        for root in (CHECKPOINT_DIR, UNET_DIR)
        for p in root.glob("*")
        if p.is_file() and p.suffix.lower() in {".safetensors", ".ckpt"}
    })
    available_loras = sorted({
        p.name
        for p in LORAS_DIR.glob("*.safetensors")
        if p.is_file()
    })

    report = {
        "profiles": profiles,
        "optional": optional,
        "available_profiles": [name for name, data in profiles.items() if data["available"]],
        "available_checkpoints": available_checkpoints,
        "available_loras": available_loras,
    }
    report["summary"] = model_asset_summary(report)
    return report


def model_asset_summary(report):
    profiles = report.get("profiles", {})
    available_profiles = report.get("available_profiles", [])
    missing_profiles = {
        name: data.get("missing", [])
        for name, data in profiles.items()
        if not data.get("available")
    }
    return {
        "ready": bool(available_profiles),
        "profile_count": len(profiles),
        "available_profile_count": len(available_profiles),
        "missing_profile_count": len(missing_profiles),
        "checkpoint_count": len(report.get("available_checkpoints", [])),
        "lora_count": len(report.get("available_loras", [])),
        "available_profiles": available_profiles,
        "missing_profiles": missing_profiles,
    }


def print_model_asset_summary(report):
    summary = report.get("summary") or model_asset_summary(report)
    print("\n[ASSET SUMMARY]")
    print("ready:", summary["ready"])
    print("available_profiles:", ", ".join(summary["available_profiles"]) if summary["available_profiles"] else "-")
    print("checkpoint_count:", summary["checkpoint_count"])
    print("lora_count:", summary["lora_count"])
    for name, data in report.get("profiles", {}).items():
        state = "OK" if data.get("available") else "MISS"
        missing = data.get("missing") or []
        print(f"profile {name}: {state}")
        if missing:
            for item in missing:
                print("  missing:", item)


def find_models_zip():
    if not INPUT_ROOT.exists():
        return None

    for z in INPUT_ROOT.rglob("models.zip"):
        if z.is_file() and z.stat().st_size > 100 * 1024 * 1024:
            return z

    return None


def wait_for_model_dataset(timeout=900):
    start = time.time()

    while time.time() - start < timeout:
        merge_all_model_datasets(verbose=False)
        profiles = available_model_profiles()
        if profiles:
            return "merged", profiles

        z = find_models_zip()
        if z is not None:
            return "zip", z

        elapsed = int(time.time() - start)
        print(
            f"\r[WAIT] 等待模型 Dataset 出现在 /kaggle/input ... "
            f"{elapsed}s / {timeout}s",
            end=""
        )
        time.sleep(5)

    print()
    return None, None


merge_all_model_datasets()
usable_profiles = available_model_profiles()

if usable_profiles:
    print("[SKIP] 可用模型分支已就绪:", ", ".join(usable_profiles))
else:
    source_type, source_path = wait_for_model_dataset(timeout=900)

    if source_type == "zip":
        models_zip = source_path

        print("\n[INFO] 找到 models.zip:")
        print(models_zip)
        print("[SIZE]", round(models_zip.stat().st_size / 1024 / 1024, 2), "MB")

        with zipfile.ZipFile(models_zip, "r") as z:
            members = z.infolist()
            total = len(members)

            for i, member in enumerate(members, 1):
                print(f"\r[UNZIP] {i}/{total} {member.filename}", end="")
                z.extract(member, COMFY_DIR)

        print("\n[OK] models.zip 解压完成")
        merge_all_model_datasets()

    elif source_type != "merged":
        print("\n[ERROR] 没有找到模型 Dataset 或 models.zip。")
        print("\n当前 /kaggle/input 内容：")

        if INPUT_ROOT.exists():
            for p in INPUT_ROOT.rglob("*"):
                print(" -", p)

        raise RuntimeError(
            "模型 Dataset 尚未准备好。\n"
            "请等待 Kaggle 右侧 Input 下载完成后重新运行 Cell 1。"
        )

usable_profiles = available_model_profiles()
if not usable_profiles:
    print_model_profile_status()
    raise RuntimeError(
        "没有可用模型分支。至少需要满足 SD1.5、SDXL 或 Flux 其中一套模型文件。"
    )

print("\n可用模型分支:", ", ".join(usable_profiles))


def merge_extra_lora_datasets():
    if not INPUT_ROOT.exists():
        return

    linked = 0
    seen = set()
    for lora_dir in INPUT_ROOT.rglob("models/loras"):
        if not lora_dir.is_dir():
            continue
        for src in sorted(lora_dir.glob("*.safetensors")):
            if src.name in seen:
                continue
            seen.add(src.name)
            dst = LORAS_DIR / src.name
            if dst.exists() or dst.is_symlink():
                continue
            link_or_copy(src, dst)
            linked += 1

    print(f"[OK] 额外 LoRA Dataset 合并完成: {linked} 个")


merge_extra_lora_datasets()


print("\n模型文件检查：")
print_model_profile_status()


print("\nOptional identity file check:")
for item in OPTIONAL_MODEL_HINTS:
    p = item["target"]
    min_mb = item["min_mb"]
    if file_ok(p, min_mb):
        print("[OK]", item["label"], "|", p, "|", round(p.stat().st_size / 1024 / 1024, 2), "MB")
    else:
        print("[MISS]", item["label"], "|", p.name, "| mount the matching identity Dataset if you need it")

asset_report = model_asset_report()
print_model_asset_summary(asset_report)
save_runtime_versions({"model_assets": asset_report})


# ============================================================
# 4. 安装 ComfyUI / IPAdapter / InsightFace 依赖
# ============================================================

show_step("4. 安装运行依赖")

run_cmd(
    f"{sys.executable} -m pip install --no-input -r requirements.txt",
    cwd=COMFY_DIR,
    check=True
)

req = IPADAPTER_NODE_DIR / "requirements.txt"

if req.exists():
    run_cmd(
        f"{sys.executable} -m pip install --no-input -r {req}",
        check=True
    )
else:
    print("[INFO] IPAdapter 节点无 requirements.txt，跳过")


def install_cpu_onnxruntime_for_pulid():
    print("[INFO] 修正 PuLID/InsightFace 使用的 onnxruntime 为 CPU 兼容版")
    run_cmd(
        f"{sys.executable} -m pip uninstall -y onnxruntime onnxruntime-gpu",
        check=False,
    )
    run_cmd(
        f"{sys.executable} -m pip install --no-input --no-cache-dir --force-reinstall "
        f"onnxruntime==1.20.1 opencv-python-headless huggingface_hub tqdm",
        check=True,
    )
    run_cmd(
        f"""{sys.executable} - <<'PY'
import onnxruntime as ort
print("[OK] onnxruntime:", ort.__version__)
print("[OK] onnxruntime providers:", ort.get_available_providers())
PY""",
        check=True,
    )


def patch_pulid_flux_for_comfyui():
    import re

    candidate_files = [
        PULID_NODE_DIR / "PulidFluxHook.py",
        PULID_NODE_DIR / "pulidflux.py",
    ]
    candidate_files.extend(PULID_NODE_DIR.glob("**/*.py"))
    seen = set()
    pulid_file = None
    source = None
    pattern = re.compile(
        r"def\s+pulid_forward_orig\s*\((?P<args>.*?)\)\s*(?:->\s*[^:]+)?\s*:",
        re.S,
    )
    match = None

    for candidate in candidate_files:
        if candidate in seen or not candidate.exists():
            continue
        seen.add(candidate)
        candidate_source = candidate.read_text(encoding="utf-8")
        candidate_match = pattern.search(candidate_source)
        if candidate_match:
            pulid_file = candidate
            source = candidate_source
            match = candidate_match
            break

    if pulid_file is None or source is None or match is None:
        searched = ", ".join(str(p) for p in seen)
        print("[INFO] PuLID Flux patch searched files:")
        print(searched[:1800])
        raise RuntimeError("PuLID Flux patch failed: pulid_forward_orig signature not found")

    print(f"[INFO] PuLID Flux patch target: {pulid_file}")

    signature = match.group(0)
    if "**kwargs" in signature:
        print("[SKIP] PuLID Flux already accepts ComfyUI extra kwargs")
        print("[OK] PuLID Flux pulid_forward_orig signature:")
        print(signature[:900])
        return

    insert_at = source.rfind(")", match.start(), match.end())
    if insert_at < 0:
        raise RuntimeError("PuLID Flux patch failed: pulid_forward_orig close paren not found")
    before_signature_close = source[:insert_at]
    if before_signature_close.rstrip().endswith(","):
        patched = before_signature_close.rstrip() + "\n    **kwargs" + source[insert_at:]
    else:
        patched = before_signature_close + ", **kwargs" + source[insert_at:]
    pulid_file.write_text(patched, encoding="utf-8")

    verified = pulid_file.read_text(encoding="utf-8")
    verified_match = pattern.search(verified)
    if not verified_match or "**kwargs" not in verified_match.group(0):
        raise RuntimeError("PuLID Flux patch failed: pulid_forward_orig still does not accept **kwargs")

    print("[PATCH] PuLID Flux compatibility: pulid_forward_orig accepts **kwargs")
    print("[OK] PuLID Flux pulid_forward_orig signature:")
    print(verified_match.group(0)[:900])


pulid_req = PULID_NODE_DIR / "requirements.txt"
if PULID_NODE_DIR.exists() and pulid_req.exists():
    run_cmd(
        f"{sys.executable} -m pip install --no-input -r {pulid_req}",
        check=True
    )
    run_cmd(
        f"{sys.executable} -m pip install --no-input facenet-pytorch --no-deps",
        check=True
    )
elif PULID_NODE_DIR.exists():
    print("[INFO] PuLID Flux 节点无 requirements.txt，仅安装 facenet-pytorch")
    run_cmd(
        f"{sys.executable} -m pip install --no-input facenet-pytorch --no-deps",
        check=True
    )
else:
    print("[INFO] 未安装 PuLID Flux 节点依赖")

if PULID_NODE_DIR.exists():
    install_cpu_onnxruntime_for_pulid()
    patch_pulid_flux_for_comfyui()

needs_insightface = any(profile in usable_profiles for profile in ("SD1.5", "SDXL"))
if needs_insightface:
    run_cmd(
        "apt-get update -qq && "
        "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq "
        "build-essential python3-dev cmake ninja-build "
        "libgl1 libglib2.0-0 libopenblas-dev",
        check=False
    )

    run_cmd(
        f"{sys.executable} -m pip install --no-input --no-cache-dir "
        f"setuptools wheel cython packaging numpy "
        f"onnxruntime==1.20.1 opencv-python-headless huggingface_hub tqdm",
        check=True
    )

    try:
        import insightface
        print("[SKIP] insightface 已可导入")
    except Exception:
        run_cmd(
            f"{sys.executable} -m pip install --no-input --no-cache-dir "
            f"--no-build-isolation -v insightface==0.7.3",
            check=True
        )
else:
    print("[SKIP] 当前仅有 Flux/PuLID 分支可用，不需要初始化 InsightFace buffalo_l")


# ============================================================
# 5. 初始化 buffalo_l
# ============================================================

if needs_insightface:
    show_step("5. 初始化 InsightFace buffalo_l")

    init_script = ROOT / "init_buffalo_l.py"

    init_script.write_text(
        textwrap.dedent(
            r'''
            from insightface.app import FaceAnalysis
            import onnxruntime as ort

            print("[INFO] onnxruntime providers:", ort.get_available_providers())
            print("[INFO] 初始化 buffalo_l")

            app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
            app.prepare(ctx_id=-1, det_size=(640, 640))

            print("[OK] buffalo_l 初始化完成")
            '''
        ),
        encoding="utf-8"
    )

    run_cmd(f"{sys.executable} {init_script}", check=True)
else:
    show_step("5. 跳过 InsightFace buffalo_l")
    print("[SKIP] Flux PuLID FaceNet 不需要 buffalo_l；继续启动 ComfyUI")


# ============================================================
# 6. 启动 ComfyUI 本地 API
# ============================================================

show_step("6. 启动 ComfyUI 本地 API")

def start_comfyui():
    stop_existing_comfyui()

    env = os.environ.copy()
    env["PYTHONHASHSEED"] = "0"
    env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    env["NVIDIA_TF32_OVERRIDE"] = "0"
    env["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "0"
    env["PYTHONUNBUFFERED"] = "1"

    cmd = [
        sys.executable,
        "main.py",
        "--listen",
        "0.0.0.0",
        "--port",
        str(PORT),
        "--use-pytorch-cross-attention",
        "--disable-smart-memory",
    ]

    print("[CMD]", " ".join(cmd))

    p = subprocess.Popen(
        cmd,
        cwd=str(COMFY_DIR),
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    def log_reader():
        for line in p.stdout:
            print("[ComfyUI]", line.rstrip())

    threading.Thread(target=log_reader, daemon=True).start()

    start = time.time()
    timeout = 240

    while time.time() - start < timeout:
        if p.poll() is not None:
            raise RuntimeError(f"ComfyUI 进程提前退出: {p.returncode}")

        if port_open("127.0.0.1", PORT):
            print("\n[OK] ComfyUI API 已启动:", COMFY_URL)
            return p

        elapsed = int(time.time() - start)
        print(f"\r[WAIT] ComfyUI 启动中... {elapsed}s / {timeout}s", end="")
        time.sleep(2)

    raise TimeoutError("ComfyUI 启动超时")

def start_comfyui_instance(instance):
    port = int(instance["port"])
    base_url = instance["url"]
    gpu_index = instance.get("gpu_index")
    name = instance.get("name") or f"port{port}"

    stop_existing_comfyui(port)

    env = os.environ.copy()
    env["PYTHONHASHSEED"] = "0"
    env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    env["NVIDIA_TF32_OVERRIDE"] = "0"
    env["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "0"
    env["PYTHONUNBUFFERED"] = "1"
    if gpu_index is not None:
        env["CUDA_VISIBLE_DEVICES"] = str(gpu_index)

    user_dir = ROOT / f"ComfyUI_user_{name}"
    user_dir.mkdir(parents=True, exist_ok=True)

    cmd = [
        sys.executable,
        "main.py",
        "--listen",
        "0.0.0.0",
        "--port",
        str(port),
        "--user-directory",
        str(user_dir),
        "--use-pytorch-cross-attention",
        "--disable-smart-memory",
    ]

    if gpu_index is not None:
        print(f"[GPU MAP] {name}: physical GPU {gpu_index} -> {base_url}")
    print("[CMD]", " ".join(cmd))

    p = subprocess.Popen(
        cmd,
        cwd=str(COMFY_DIR),
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    def log_reader():
        for line in p.stdout:
            print(f"[ComfyUI:{name}]", line.rstrip())

    threading.Thread(target=log_reader, daemon=True).start()

    start = time.time()
    timeout = 240

    while time.time() - start < timeout:
        if p.poll() is not None:
            raise RuntimeError(f"ComfyUI process exited early on {base_url}: {p.returncode}")
        if port_open("127.0.0.1", port):
            print("\n[OK] ComfyUI API ready:", base_url)
            return p
        elapsed = int(time.time() - start)
        print(f"\r[WAIT] ComfyUI starting {base_url} ... {elapsed}s / {timeout}s", end="")
        time.sleep(2)

    raise TimeoutError(f"ComfyUI startup timeout: {base_url}")


comfy_instances = desired_comfy_instances()
print("[INFO] ComfyUI instances:", json.dumps(comfy_instances, ensure_ascii=False))
comfy_processes = [start_comfyui_instance(instance) for instance in comfy_instances]
for instance in comfy_instances:
    verify_comfyui_nodes(instance["url"])
write_comfy_instances(comfy_instances)
save_runtime_versions({
    "model_assets": asset_report,
    "comfy_instances": comfy_instances,
    "multi_gpu_workers": MULTI_GPU_WORKERS,
})

show_step("Cell 1 完成")
print("""
[OK] 环境和 ComfyUI 服务已经准备好。

下一步：
1. 确保 Kaggle Input 里有当前要用的 face.png
2. 运行 Cell 2 生图

如果生图失败或想换脸，只需要重新运行 Cell 2。
""")