bootstrap.py

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent
MANIFEST = json.loads((REPO_ROOT / "manifest.json").read_text())


def run(command: list[str], *, dry_run: bool) -> None:
    print("+", " ".join(command), flush=True)
    if not dry_run:
        subprocess.run(command, check=True)


def find_comfyui(value: str | None) -> Path:
    candidates = []
    if value:
        candidates.append(Path(value).expanduser())
    candidates.extend((Path.cwd(), REPO_ROOT.parent / "ComfyUI"))
    for candidate in candidates:
        resolved = candidate.resolve()
        if (resolved / "main.py").is_file() and (resolved / "models").is_dir():
            return resolved
    raise SystemExit("ComfyUI was not found. Pass --comfyui /path/to/ComfyUI.")


def ensure_huggingface_hub(*, dry_run: bool):
    if importlib.util.find_spec("huggingface_hub") is None:
        run(
            [sys.executable, "-m", "pip", "install", "huggingface_hub>=0.28"],
            dry_run=dry_run,
        )
        if dry_run:
            return None
    from huggingface_hub import hf_hub_download

    return hf_hub_download


def install_requirements(node_dir: Path, *, dry_run: bool) -> None:
    requirements = node_dir / "requirements.txt"
    if requirements.is_file():
        run(
            [sys.executable, "-m", "pip", "install", "-r", str(requirements)],
            dry_run=dry_run,
        )


def update_node(node_dir: Path, *, dry_run: bool) -> None:
    if not (node_dir / ".git").exists():
        print(f"[node] update skipped (not a Git checkout): {node_dir}")
        return
    status = subprocess.run(
        ["git", "-C", str(node_dir), "status", "--porcelain"],
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()
    if status:
        print(f"[node] update skipped (local changes): {node_dir}")
        return
    upstream = subprocess.run(
        [
            "git",
            "-C",
            str(node_dir),
            "rev-parse",
            "--abbrev-ref",
            "--symbolic-full-name",
            "@{upstream}",
        ],
        capture_output=True,
        text=True,
    )
    if upstream.returncode != 0:
        print(f"[node] update skipped (no upstream branch): {node_dir}")
        return
    run(["git", "-C", str(node_dir), "pull", "--ff-only"], dry_run=dry_run)


def install_nodes(comfyui: Path, *, dry_run: bool, update_existing: bool) -> None:
    custom_nodes = comfyui / "custom_nodes"
    custom_nodes.mkdir(parents=True, exist_ok=True)
    for entry in MANIFEST["nodes"]:
        destination = custom_nodes / entry["name"]
        candidates = [destination]
        candidates.extend(custom_nodes / alias for alias in entry.get("aliases", []))
        present = next((candidate for candidate in candidates if candidate.exists()), None)
        if present is not None:
            print(f"[node] present: {present}")
            if update_existing:
                update_node(present, dry_run=dry_run)
        else:
            run(
                ["git", "clone", "--depth", "1", entry["url"], str(destination)],
                dry_run=dry_run,
            )
        installed = present or destination
        if installed.exists() or not dry_run:
            install_requirements(installed, dry_run=dry_run)


def model_is_valid(path: Path, minimum: int) -> bool:
    try:
        return path.is_file() and path.stat().st_size >= minimum
    except OSError:
        return False


def resolve_memory_tier(value: str) -> str:
    if value != "auto":
        return value
    memory_gb = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 1e9
    for limit, tier in ((28, "24"), (40, "32"), (56, "48"), (80, "64"), (176, "96")):
        if memory_gb < limit:
            return tier
    return "192"


def install_mlx_profile(
    comfyui: Path,
    profile: str,
    stream_chunk: int | None,
    *,
    dry_run: bool,
) -> None:
    node_dir = comfyui / "custom_nodes" / "ComfyUI-MiniMax-H3-MLX-SolAttn"
    setup = node_dir / "scripts" / "setup_comfyui.py"
    run(
        [
            sys.executable, str(setup), "--models-dir", str(comfyui / "models"),
            "--profile", profile, "--accept-license",
        ],
        dry_run=dry_run,
    )
    if stream_chunk is None:
        return
    profile_name = profile if profile in ("bf16", "bf16-pruned") else f"{profile}-bit"
    source = comfyui / "models" / "minimax_h3" / "transformers" / profile_name
    output = comfyui / "models" / "minimax_h3" / "transformers" / f"{profile_name}-stream{stream_chunk}"
    if (output / "streaming_manifest.json").is_file():
        print(f"[model] present: {output}")
        return
    run(
        [
            sys.executable, str(node_dir / "scripts" / "build_streaming_checkpoint.py"),
            "--source", str(source), "--out", str(output),
            "--chunk-size", str(stream_chunk),
        ],
        dry_run=dry_run,
    )


def install_pytorch_models(comfyui: Path, selected: set[str], *, dry_run: bool) -> None:
    downloader = ensure_huggingface_hub(dry_run=dry_run)
    for entry in (item for item in MANIFEST["models"] if item["id"] in selected):
        destination = comfyui / entry["destination"]
        if model_is_valid(destination, entry["min_bytes"]):
            print(f"[model] present: {destination}")
            continue
        if destination.exists() or destination.is_symlink():
            raise RuntimeError(
                f"Incomplete model already exists at {destination}. Remove it and rerun."
            )
        destination.parent.mkdir(parents=True, exist_ok=True)
        print(f"[model] downloading {entry['repo_id']}:{entry['filename']}")
        if dry_run:
            continue
        if entry["repo_id"] == "Comfy-Org/MiniMax-H3":
            local_dir = comfyui / "models"
        else:
            local_dir = destination.parent
        downloaded = Path(
            downloader(
                repo_id=entry["repo_id"],
                filename=entry["filename"],
                local_dir=local_dir,
            )
        )
        if downloaded.resolve() != destination.resolve():
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(downloaded), destination)
        if not model_is_valid(destination, entry["min_bytes"]):
            raise RuntimeError(f"Downloaded model failed size validation: {destination}")


def install_models(comfyui: Path, tier: dict, workflow_set: str, *, dry_run: bool) -> None:
    if tier["backend"] == "mlx":
        if workflow_set in ("fast", "all"):
            install_mlx_profile(
                comfyui,
                tier["mlx_profile"],
                5 if tier.get("streaming") else None,
                dry_run=dry_run,
            )
        if workflow_set in ("quality", "all"):
            install_pytorch_models(
                comfyui,
                set(tier["quality_models"]),
                dry_run=dry_run,
            )
        return
    install_pytorch_models(comfyui, set(tier["models"]), dry_run=dry_run)


def install_workflows(comfyui: Path, names: list[str], *, dry_run: bool) -> None:
    destination_dir = comfyui / "user" / "default" / "workflows"
    if not dry_run:
        destination_dir.mkdir(parents=True, exist_ok=True)
    for name in names:
        source = REPO_ROOT / "workflows" / name
        destination = destination_dir / name
        if not source.is_file():
            raise RuntimeError(f"Workflow is missing from config repository: {source}")
        print(f"[workflow] install: {destination}")
        if not dry_run:
            shutil.copy2(source, destination)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Install MiniMax-H3 staged workflows, nodes, and checkpoints."
    )
    parser.add_argument("--comfyui", help="Path to the ComfyUI checkout")
    parser.add_argument(
        "--memory-tier",
        choices=["auto", "24", "32", "48", "64", "96", "192"],
        default="auto",
        help="Install the models and creator workflows for this unified-memory tier.",
    )
    parser.add_argument(
        "--workflow-set",
        choices=["fast", "quality", "all"],
        default="fast",
        help="Install the tier's fast quantized workflows, BF16 quality workflows, or both.",
    )
    parser.add_argument("--skip-nodes", action="store_true")
    parser.add_argument(
        "--skip-node-updates",
        action="store_true",
        help="Keep existing custom-node Git checkouts at their current revisions.",
    )
    parser.add_argument("--skip-models", action="store_true")
    parser.add_argument("--skip-workflows", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    comfyui = find_comfyui(args.comfyui)
    memory_tier = resolve_memory_tier(args.memory_tier)
    tier = MANIFEST["memory_tiers"][memory_tier]
    if args.workflow_set in ("quality", "all") and "quality_workflows" not in tier:
        raise SystemExit(f"Memory tier {memory_tier} has no separate quality workflow set.")
    workflows = list(tier["workflows"]) if args.workflow_set in ("fast", "all") else []
    if args.workflow_set in ("quality", "all"):
        workflows.extend(tier["quality_workflows"])
    print(f"ComfyUI: {comfyui}")
    print(f"Memory tier: {memory_tier} GB{'+' if memory_tier == '192' else ''}")
    if not args.skip_nodes:
        install_nodes(
            comfyui,
            dry_run=args.dry_run,
            update_existing=not args.skip_node_updates,
        )
    if not args.skip_models:
        install_models(comfyui, tier, args.workflow_set, dry_run=args.dry_run)
    if not args.skip_workflows:
        install_workflows(comfyui, workflows, dry_run=args.dry_run)
    if tier["backend"] == "pytorch" or (
        args.workflow_set in ("quality", "all") and tier.get("quality_models")
    ):
        print(
            "Done. Restart ComfyUI with --use-pytorch-cross-attention. "
            "Do not use --force-fp16, --fp16-unet, or --fp16-text-enc."
        )
    else:
        print("Done. Restart ComfyUI and load the installed memory-tier workflow.")


if __name__ == "__main__":
    main()