Files
AiVideo/scripts/check_comfy.py
2026-03-25 13:33:48 +08:00

68 lines
2.0 KiB
Python

from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Any
import httpx
import yaml
def fetch_object_info(base_url: str, timeout_s: float = 5.0) -> dict[str, Any]:
url = base_url.rstrip("/") + "/object_info"
with httpx.Client(timeout=timeout_s) as client:
r = client.get(url)
r.raise_for_status()
data = r.json()
if not isinstance(data, dict):
raise RuntimeError(f"Unexpected object_info type: {type(data)}")
return data
def read_base_url_from_config(config_path: str) -> str | None:
p = Path(config_path)
if not p.exists():
return None
try:
raw = yaml.safe_load(p.read_text(encoding="utf-8"))
except Exception:
return None
if not isinstance(raw, dict):
return None
app = raw.get("app")
if not isinstance(app, dict):
return None
v = app.get("comfy_base_url")
if isinstance(v, str) and v.strip():
return v.strip()
return None
def main() -> int:
parser = argparse.ArgumentParser(description="Check ComfyUI API connectivity")
parser.add_argument(
"--base-url",
default="",
help="ComfyUI base URL (if empty, read from config app.comfy_base_url)",
)
parser.add_argument("--config", default="./configs/config.yaml", help="Config yaml path")
parser.add_argument("--timeout", type=float, default=5.0, help="Request timeout seconds")
parser.add_argument("--pretty", action="store_true", help="Pretty print JSON")
args = parser.parse_args()
try:
base_url = args.base_url.strip() or read_base_url_from_config(args.config) or "http://127.0.0.1:8188"
data = fetch_object_info(base_url, timeout_s=args.timeout)
out = json.dumps(data, ensure_ascii=False, indent=2 if args.pretty else None)
sys.stdout.write(out + "\n")
return 0
except Exception as e:
sys.stderr.write(f"[check_comfy] ERROR: {e}\n")
return 2
if __name__ == "__main__":
raise SystemExit(main())