Tool_Agent/tests/run_comfy/run_workflow.py

#!/usr/bin/env python3
"""RunComfy Server API + ComfyUI Backend API

流程：启动机器 → 上传 input/ 目录文件 → 提交 workflow → WebSocket 监听 → 下载结果 → 关机

input/ 目录结构：
    input/
    ├── images/          → 上传到 ComfyUI input/（LoadImage 节点用）
    ├── loras/           → 上传到 ComfyUI models/loras/
    ├── checkpoints/     → 上传到 ComfyUI models/checkpoints/
    ├── vae/             → 上传到 ComfyUI models/vae/
    └── (其他文件)        → 上传到 ComfyUI input/

用法:
    python run_workflow.py workflow_api.json
    python run_workflow.py workflow_api.json --input-dir ./input
    python run_workflow.py workflow_api.json --input-dir ./input --server-type large
"""

import argparse
import json
import os
import sys
import time
import urllib.parse
import uuid
from pathlib import Path

import requests
import websocket
from dotenv import load_dotenv
from check_workflow import analyze, check_files_exist

load_dotenv(Path(__file__).parent.parent.parent / ".env")

BASE_URL = "https://beta-api.runcomfy.net/prod/api"
USER_ID = os.getenv("RUNCOMFY_USER_ID")
API_TOKEN = os.getenv("API_TOKEN")

HEADERS = {
    "Authorization": f"Bearer {API_TOKEN}",
    "Content-Type": "application/json",
}

DEFAULT_VERSION_ID = "90f77137-ba75-400d-870f-204c614ae8a3"  # RunComfy/ComfyUI-NodesLoaded

# input/ 子目录 → ComfyUI 上传类型和 subfolder 映射
# type="input" 对应 ComfyUI 的 input 目录
# type="model" 暂无官方支持，lora 等模型走 subfolder 区分
SUBDIR_UPLOAD_MAP = {
    "images":      {"type": "input",  "subfolder": ""},
    "loras":       {"type": "input",  "subfolder": "loras"},
    "checkpoints": {"type": "input",  "subfolder": "checkpoints"},
    "vae":         {"type": "input",  "subfolder": "vae"},
    "controlnet":  {"type": "input",  "subfolder": "controlnet"},
    "upscale":     {"type": "input",  "subfolder": "upscale_models"},
}


# ── 机器管理 ──────────────────────────────────────────────

def launch_machine(version_id: str, server_type: str = "medium", duration: int = 3600) -> str:
    payload = {
        "workflow_version_id": version_id,
        "server_type": server_type,
        "estimated_duration": duration,
    }
    resp = requests.post(f"{BASE_URL}/users/{USER_ID}/servers", headers=HEADERS, json=payload)
    if not resp.ok:
        print(f"  HTTP {resp.status_code}: {resp.text}")
    resp.raise_for_status()
    print(f"  响应: {resp.json()}")
    server_id = resp.json()["server_id"]
    print(f"机器已创建: {server_id}  (type={server_type})")
    return server_id


def wait_for_ready(server_id: str, timeout: int = 300) -> str:
    print("等待机器就绪...")
    start = time.time()
    while time.time() - start < timeout:
        resp = requests.get(f"{BASE_URL}/users/{USER_ID}/servers/{server_id}", headers=HEADERS)
        resp.raise_for_status()
        data = resp.json()
        status = data.get("current_status", "")
        print(f"  状态: {status}")
        if status == "Ready":
            url = data["main_service_url"].rstrip("/")
            print(f"  就绪: {url}")
            return url
        if status in ("Error", "Failed"):
            raise Exception(f"机器启动失败: {status}")
        time.sleep(5)
    raise TimeoutError(f"等待超时 ({timeout}s)")


def stop_machine(server_id: str):
    resp = requests.delete(f"{BASE_URL}/users/{USER_ID}/servers/{server_id}", headers=HEADERS)
    resp.raise_for_status()
    print(f"机器已关闭: {server_id}")


# ── 文件上传 ──────────────────────────────────────────────

def upload_file(comfy_url: str, file_path: Path, file_type: str = "input", subfolder: str = "") -> str:
    """上传文件到 ComfyUI，返回服务器上的实际文件名"""
    with open(file_path, "rb") as f:
        files = [("image", (file_path.name, f, "application/octet-stream"))]
        data = {"overwrite": "true", "type": file_type, "subfolder": subfolder}
        resp = requests.post(f"{comfy_url}/upload/image", data=data, files=files)
    resp.raise_for_status()
    server_name = resp.json()["name"]
    subfolder_str = f" → {subfolder}/{server_name}" if subfolder else f" → {server_name}"
    print(f"  上传: {file_path.name}{subfolder_str}")
    return server_name


def upload_input_dir(comfy_url: str, input_dir: Path) -> dict[str, str]:
    """
    扫描 input_dir，按子目录上传文件，返回 {原文件名: 服务器文件名} 映射
    - input/images/  → type=input, subfolder=""
    - input/loras/   → type=input, subfolder="loras"
    - input/*.png    → type=input, subfolder=""（根目录文件）
    """
    if not input_dir.exists():
        print(f"  input 目录不存在: {input_dir}")
        return {}

    uploaded = {}
    IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
    VIDEO_EXTS = {".mp4", ".avi", ".mov", ".webm"}
    MODEL_EXTS = {".safetensors", ".ckpt", ".pt", ".pth", ".gguf"}
    ALL_EXTS = IMAGE_EXTS | VIDEO_EXTS | MODEL_EXTS

    # 根目录文件 → input
    for f in input_dir.iterdir():
        if f.is_file() and f.suffix.lower() in ALL_EXTS:
            server_name = upload_file(comfy_url, f, "input", "")
            uploaded[f.name] = server_name

    # 子目录文件
    for subdir in input_dir.iterdir():
        if not subdir.is_dir():
            continue
        mapping = SUBDIR_UPLOAD_MAP.get(subdir.name, {"type": "input", "subfolder": subdir.name})
        for f in subdir.iterdir():
            if f.is_file() and f.suffix.lower() in ALL_EXTS:
                server_name = upload_file(comfy_url, f, mapping["type"], mapping["subfolder"])
                uploaded[f.name] = server_name

    return uploaded


# ── 提交 workflow ─────────────────────────────────────────

def submit_prompt(comfy_url: str, workflow_api: dict, client_id: str) -> str:
    payload = {"prompt": workflow_api, "client_id": client_id}
    resp = requests.post(f"{comfy_url}/prompt", json=payload)
    resp.raise_for_status()
    data = resp.json()
    if data.get("node_errors"):
        print(f"  节点错误: {data['node_errors']}")
    prompt_id = data["prompt_id"]
    print(f"任务已提交: {prompt_id}")
    return prompt_id


# ── WebSocket 监听 ────────────────────────────────────────

def wait_for_completion(comfy_url: str, client_id: str, prompt_id: str, timeout: int = 600):
    scheme = "wss" if comfy_url.startswith("https") else "ws"
    ws_url = f"{scheme}://{comfy_url.split('://', 1)[-1]}/ws?clientId={client_id}"
    print("WebSocket 监听中...")

    ws = websocket.WebSocket()
    ws.settimeout(timeout)
    ws.connect(ws_url)

    try:
        while True:
            out = ws.recv()
            if not out or isinstance(out, bytes):
                continue
            msg = json.loads(out)
            msg_type = msg.get("type", "")
            data = msg.get("data", {})

            if msg_type == "executing":
                node = data.get("node")
                if data.get("prompt_id") == prompt_id and node is None:
                    print("  执行完成")
                    break
                if node:
                    print(f"  执行节点: {node}")

            elif msg_type == "progress":
                value = data.get("value", 0)
                max_val = data.get("max", 1)
                print(f"  进度: {value}/{max_val}")

            elif msg_type == "execution_error":
                if data.get("prompt_id") == prompt_id:
                    raise Exception(f"执行错误: {data.get('exception_message', 'unknown')}")
    finally:
        ws.close()


# ── 下载结果 ──────────────────────────────────────────────

def download_outputs(comfy_url: str, prompt_id: str, output_dir: Path) -> list[str]:
    resp = requests.get(f"{comfy_url}/history/{prompt_id}")
    resp.raise_for_status()
    data = resp.json().get(prompt_id, {})
    outputs = data.get("outputs", {})

    output_dir.mkdir(parents=True, exist_ok=True)
    saved = []

    for node_id, node_output in outputs.items():
        if "images" in node_output:
            for image in node_output["images"]:
                params = {"filename": image["filename"], "subfolder": image.get("subfolder", ""), "type": image.get("temp") or image.get("type", "output")}
                resp = requests.get(f"{comfy_url}/view?{urllib.parse.urlencode(params)}")
                resp.raise_for_status()
                out_path = output_dir / image["filename"]
                out_path.write_bytes(resp.content)
                print(f"  图片: {out_path}")
                saved.append(str(out_path))

        if "gifs" in node_output:
            for video in node_output["gifs"]:
                params = {"filename": video["filename"], "subfolder": video.get("subfolder", ""), "format": video.get("format", "mp4")}
                resp = requests.get(f"{comfy_url}/view?{urllib.parse.urlencode(params)}")
                resp.raise_for_status()
                out_path = output_dir / video["filename"]
                out_path.write_bytes(resp.content)
                print(f"  视频: {out_path}")
                saved.append(str(out_path))

    return saved


# ── 主流程 ────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="RunComfy workflow runner")
    parser.add_argument("workflow", help="workflow_api.json 路径")
    parser.add_argument("--input-dir", default="input", metavar="DIR",
                        help="输入文件目录，默认 input/。子目录 images/loras/checkpoints/vae/ 自动上传到对应位置")
    parser.add_argument("--version-id", default=DEFAULT_VERSION_ID, help="RunComfy workflow version_id")
    parser.add_argument("--server-type", default="medium",
                        choices=["medium", "large", "extra-large", "2x-large", "2xl-turbo"])
    parser.add_argument("--duration", type=int, default=3600, help="预估运行时长(秒)，默认3600")
    parser.add_argument("--keep-alive", action="store_true", help="完成后不自动关机")
    parser.add_argument("--server-id", metavar="ID", help="复用已有机器，跳过启动步骤")
    parser.add_argument("--skip-upload", action="store_true", help="跳过文件上传，直接提交 workflow")
    parser.add_argument("--output-dir", default="output", metavar="DIR", help="结果下载目录，默认 output/")
    args = parser.parse_args()

    if not USER_ID or not API_TOKEN:
        print("ERROR: 请设置 RUNCOMFY_USER_ID 和 API_TOKEN 环境变量")
        sys.exit(1)

    print(f"USER_ID : {USER_ID}")
    print(f"API_TOKEN: {API_TOKEN[:8]}...")

    workflow_path = Path(args.workflow)
    if not workflow_path.exists():
        print(f"ERROR: 文件不存在: {workflow_path}")
        sys.exit(1)

    with open(workflow_path, "r", encoding="utf-8") as f:
        workflow_api = json.load(f)

    # 提交前 check
    input_dir = Path(args.input_dir)
    all_input_files = list(input_dir.rglob("*")) if input_dir.exists() else []
    result = analyze(workflow_api)
    missing = check_files_exist(result["file_inputs"], all_input_files)

    if result["issues"] or missing:
        print("❌ workflow 检查未通过：")
        for issue in result["issues"]:
            print(f"  {issue}")
        for fi in missing:
            print(f"  缺少文件: {fi['filename']}  (节点 [{fi['node_id']}] {fi['class_type']})")
            print(f"  请将该文件放入 {input_dir}/ 目录")
        sys.exit(1)

    if result["widget_params"]:
        print("⚠️  存在 widget_* 占位参数，参数名可能不准确，继续运行...")
    else:
        print("✓ workflow 检查通过")

    client_id = str(uuid.uuid4())
    server_id = args.server_id  # None if not provided

    try:
        # 1. 启动机器（或复用已有机器）
        if server_id:
            print(f"复用已有机器: {server_id}")
            comfy_url = wait_for_ready(server_id)
        else:
            server_id = launch_machine(args.version_id, args.server_type, args.duration)
            comfy_url = wait_for_ready(server_id)

        # 2. 上传 input 目录
        if args.skip_upload:
            print("跳过文件上传 (--skip-upload)")
        else:
            print(f"\n上传 input 目录: {input_dir}")
            upload_input_dir(comfy_url, input_dir)

        # 3. 提交 workflow
        print(f"\n提交 workflow...")
        prompt_id = submit_prompt(comfy_url, workflow_api, client_id)

        # 4. 监听执行进度
        wait_for_completion(comfy_url, client_id, prompt_id)

        # 5. 下载结果
        print(f"\n下载结果...")
        saved = download_outputs(comfy_url, prompt_id, Path(args.output_dir))
        print(f"\n完成，共 {len(saved)} 个文件")

        if args.keep_alive:
            print(f"\n--keep-alive 模式，机器保持运行: {server_id}")
            print(f"ComfyUI URL: {comfy_url}")
        else:
            print("\n关闭机器...")
            stop_machine(server_id)

    except Exception as e:
        print(f"\n错误: {e}")
        print(f"机器 {server_id} 未自动关闭，请手动处理")
        sys.exit(1)


if __name__ == "__main__":
    main()