242 lines
7.6 KiB
Python
242 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_MANIFEST = REPO_ROOT / "model_manifest.json"
|
|
|
|
|
|
def log(message: str) -> None:
|
|
print(f"[illusion-architect] {message}")
|
|
|
|
|
|
def load_manifest(path: Path) -> dict:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def available_presets(manifest: dict) -> set[str]:
|
|
return set(manifest.get("presets", {}).keys())
|
|
|
|
|
|
def detect_comfyui_root(provided: str | None) -> Path | None:
|
|
if provided:
|
|
candidate = Path(provided).expanduser().resolve()
|
|
return candidate
|
|
|
|
env_value = os.environ.get("COMFYUI_ROOT")
|
|
if env_value:
|
|
return Path(env_value).expanduser().resolve()
|
|
|
|
if REPO_ROOT.parent.name == "custom_nodes":
|
|
candidate = REPO_ROOT.parent.parent.resolve()
|
|
if (candidate / "main.py").exists() or (candidate / "ComfyUI.exe").exists():
|
|
return candidate
|
|
return candidate
|
|
|
|
return None
|
|
|
|
|
|
def verify_environment() -> None:
|
|
log(f"Python {sys.version.split()[0]}")
|
|
if sys.version_info < (3, 10):
|
|
raise SystemExit("Python 3.10+ is required.")
|
|
|
|
git_path = shutil.which("git")
|
|
if not git_path:
|
|
raise SystemExit("Git is required but was not found on PATH.")
|
|
log(f"Git detected at {git_path}")
|
|
|
|
try:
|
|
import torch # type: ignore
|
|
|
|
log(f"Torch {torch.__version__}")
|
|
except Exception as exc:
|
|
log(f"Torch check failed: {exc}")
|
|
|
|
|
|
def clone_repo(repo_url: str, target_dir: Path) -> None:
|
|
if target_dir.exists():
|
|
log(f"Custom node already present: {target_dir}")
|
|
return
|
|
|
|
target_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
log(f"Cloning {repo_url} -> {target_dir}")
|
|
subprocess.run(
|
|
["git", "clone", "--depth", "1", repo_url, str(target_dir)],
|
|
check=True,
|
|
)
|
|
|
|
|
|
def install_custom_nodes(comfyui_root: Path, manifest: dict) -> None:
|
|
for entry in manifest.get("custom_nodes", []):
|
|
target_dir = comfyui_root / entry["path"]
|
|
clone_repo(entry["repo"], target_dir)
|
|
|
|
|
|
def selected_families(raw_values: list[str] | None) -> set[str]:
|
|
if not raw_values:
|
|
return {"sd15", "sdxl", "flux"}
|
|
if "all" in raw_values:
|
|
return {"sd15", "sdxl", "flux"}
|
|
return set(raw_values)
|
|
|
|
|
|
def selected_presets(raw_values: list[str] | None, manifest: dict) -> set[str]:
|
|
if not raw_values:
|
|
return {"all"}
|
|
presets = set(raw_values)
|
|
if "all" in presets:
|
|
return {"all"}
|
|
|
|
unknown = presets - available_presets(manifest)
|
|
if unknown:
|
|
raise SystemExit(f"Unknown preset(s): {', '.join(sorted(unknown))}")
|
|
return presets
|
|
|
|
|
|
def model_matches(entry: dict, families: set[str], presets: set[str], manifest: dict) -> bool:
|
|
if entry.get("family") not in families:
|
|
return False
|
|
|
|
if "all" in presets:
|
|
return True
|
|
|
|
entry_styles = set(entry.get("styles") or [entry.get("style")])
|
|
for preset_name in presets:
|
|
preset = manifest["presets"][preset_name]
|
|
preset_families = set(preset.get("families", []))
|
|
preset_styles = set(preset.get("styles", []))
|
|
if entry.get("family") in preset_families and entry_styles & preset_styles:
|
|
return True
|
|
return False
|
|
|
|
|
|
def resolve_hf_token(cli_token: str | None, entry: dict) -> str | None:
|
|
if cli_token:
|
|
return cli_token
|
|
|
|
token_env = entry.get("token_env")
|
|
if token_env:
|
|
token = os.environ.get(token_env)
|
|
if token:
|
|
return token
|
|
|
|
return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
|
|
|
|
|
|
def download_model(entry: dict, comfyui_root: Path, token: str | None) -> None:
|
|
destination_dir = comfyui_root / entry["destination"]
|
|
destination_dir.mkdir(parents=True, exist_ok=True)
|
|
output_path = destination_dir / entry.get("output_name", Path(entry["filename"]).name)
|
|
|
|
if output_path.exists():
|
|
log(f"Model already present: {output_path.name}")
|
|
return
|
|
|
|
if entry.get("requires_token") and not token:
|
|
log(
|
|
f"Skipping {entry['id']} because it requires a Hugging Face token. Use --hf-token or set {entry.get('token_env', 'HF_TOKEN')}."
|
|
)
|
|
return
|
|
|
|
log(f"Downloading {entry['id']} from {entry['repo_id']}:{entry['filename']}")
|
|
cached_path = hf_hub_download(
|
|
repo_id=entry["repo_id"],
|
|
filename=entry["filename"],
|
|
token=token,
|
|
resume_download=True,
|
|
)
|
|
shutil.copy2(cached_path, output_path)
|
|
log(f"Saved {output_path}")
|
|
|
|
|
|
def install_models(
|
|
comfyui_root: Path,
|
|
manifest: dict,
|
|
families: set[str],
|
|
presets: set[str],
|
|
cli_token: str | None,
|
|
) -> None:
|
|
for entry in manifest.get("models", []):
|
|
if not model_matches(entry, families, presets, manifest):
|
|
continue
|
|
try:
|
|
download_model(entry, comfyui_root, resolve_hf_token(cli_token, entry))
|
|
except Exception as exc:
|
|
log(f"Failed to download {entry['id']}: {exc}")
|
|
|
|
|
|
def print_manual_followups(manifest: dict) -> None:
|
|
manual_models = manifest.get("manual_models", [])
|
|
if not manual_models:
|
|
return
|
|
|
|
log("Manual follow-up models:")
|
|
for entry in manual_models:
|
|
print(f" - {entry['id']} -> {entry['destination']}: {entry['notes']}")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Check and fill ComfyUI hidden-image dependencies.")
|
|
parser.add_argument("--comfyui-root", help="Path to the ComfyUI root directory.")
|
|
parser.add_argument("--manifest", default=str(DEFAULT_MANIFEST), help="Path to the JSON manifest.")
|
|
parser.add_argument(
|
|
"--family",
|
|
action="append",
|
|
choices=["all", "sd15", "sdxl", "flux"],
|
|
help="Limit downloads to a family. Repeat for multiple values.",
|
|
)
|
|
parser.add_argument(
|
|
"--preset",
|
|
action="append",
|
|
help="Install a named preset from the manifest, such as qr, text, ghost, or all. Repeat for multiple values.",
|
|
)
|
|
parser.add_argument("--list-presets", action="store_true", help="Print available install presets and exit.")
|
|
parser.add_argument("--hf-token", help="Explicit Hugging Face token used for gated or private downloads.")
|
|
parser.add_argument("--skip-custom-nodes", action="store_true", help="Skip cloning third-party custom nodes.")
|
|
parser.add_argument("--skip-models", action="store_true", help="Skip model downloads.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
verify_environment()
|
|
manifest = load_manifest(Path(args.manifest).expanduser().resolve())
|
|
|
|
if args.list_presets:
|
|
log("Available presets:")
|
|
for preset_name, preset in manifest.get("presets", {}).items():
|
|
print(f" - {preset_name}: {preset.get('notes', '')}")
|
|
return
|
|
|
|
comfyui_root = detect_comfyui_root(args.comfyui_root)
|
|
|
|
if comfyui_root is None:
|
|
raise SystemExit(
|
|
"Could not detect ComfyUI root. Re-run with --comfyui-root /path/to/ComfyUI or clone this repo into ComfyUI/custom_nodes."
|
|
)
|
|
|
|
log(f"Using ComfyUI root: {comfyui_root}")
|
|
families = selected_families(args.family)
|
|
presets = selected_presets(args.preset, manifest)
|
|
log(f"Selected families: {', '.join(sorted(families))}")
|
|
log(f"Selected presets: {', '.join(sorted(presets))}")
|
|
|
|
if not args.skip_custom_nodes:
|
|
install_custom_nodes(comfyui_root, manifest)
|
|
if not args.skip_models:
|
|
install_models(comfyui_root, manifest, families, presets, args.hf_token)
|
|
print_manual_followups(manifest)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |