Tag: Controls

  • How to Build an AI Agent That Controls ComfyUI: A Python API Tutorial (2026)

    How to Build an AI Agent That Controls ComfyUI: A Python API Tutorial (2026)

    Why This Matters

    A standard ComfyUI workflow exposed to the agent as a single function call

    The most useful AI product of 2026 is not a chatbox. It is a workflow engine that wraps a generative model in a reliable loop. ComfyUI is the most flexible open-source image and video generator shipping today, but its interface is operator-driven. The moment you let an LLM see the workflow as a Python function, you turn ComfyUI into a backend that any agent can call from a phone, a Slack channel, or a batch job.

    This guide shows the minimum working harness. Five tool functions, one retry wrapper, one example chat client. Total code is under 200 lines.

    What ComfyUI Actually Exposes

    Out of the box, a running ComfyUI server answers on a small HTTP surface. You do not need to install anything new.

    • POST /prompt — submit a workflow. Returns a prompt_id.
    • GET /history/{prompt_id} — poll until the run finishes.
    • GET /view?filename=X&type=output — fetch the actual image bytes.
    • GET /queue — see running and pending jobs.
    • GET /models — list installed checkpoints, LoRAs, and VAEs.
    • POST /interrupt — kill a stuck run.
    • GET /system_stats — VRAM, CPU, RAM. Use this in the agent to refuse new work when VRAM is full.

    The hard part is not the HTTP. The hard part is the workflow JSON. A ComfyUI workflow saved from the UI comes in two flavours: the visual workflow.json (LiteGraph format) and the API-facing workflow_api.json. Always send the API format to /prompt. Save it by clicking the gear icon and choosing "Save (API Format)".

    The Five Tools the Agent Needs

    Compress the API into a minimal tool surface. Each tool is a Python function decorated as an OpenAI / Anthropic tool definition. Keep the schema tight; the model hallucinates less.

    1. list_models — returns checkpoints, LoRAs, VAEs. Agent uses this to refuse a request for a missing model before queueing.
    2. get_queue — returns running and pending jobs with their prompt_ids. Agent uses this to estimate wait time.
    3. queue_workflow — accepts a workflow_api JSON and a client_id. Returns prompt_id. Internally injects the user prompt into the CLIPTextEncode node.
    4. get_status — accepts a prompt_id. Polls /history until done or timeout. Returns the output image filename and any error message.
    5. fetch_image — accepts a filename. Returns image bytes. The agent embeds these in its reply or saves them.

    Working Harness (Python, 180 lines)

    Copy this, point COMFYUI_HOST at your server, run. The agent is an OpenAI function-calling loop. Swap the OpenAI client for Anthropic or a local Qwen endpoint without changing the tool functions.

    import json, time, uuid, requests, base64
    from openai import OpenAI
    
    COMFYUI_HOST = "http://100.126.189.19:8000"
    COMFYUI_CLIENT_ID = str(uuid.uuid4())
    
    # --- TOOL 1: list_models ---
    def list_models():
        r = requests.get(f"{COMFYUI_HOST}/models", timeout=10)
        r.raise_for_status()
        return r.json()
    
    # --- TOOL 2: get_queue ---
    def get_queue():
        r = requests.get(f"{COMFYUI_HOST}/queue", timeout=10)
        r.raise_for_status()
        return {"running": r.json().get("queue_running", []),
                "pending": r.json().get("queue_pending", [])}
    
    # --- TOOL 3: queue_workflow ---
    def queue_workflow(workflow_path: str, user_prompt: str, neg_prompt: str = "blurry, low quality"):
        with open(workflow_path) as f:
            wf = json.load(f)
        # Inject user prompt into the positive CLIPTextEncode (node "4" in our template)
        if "4" in wf and wf["4"].get("class_type") == "CLIPTextEncode":
            wf["4"]["inputs"]["text"] = user_prompt
        if "5" in wf and wf["5"].get("class_type") == "CLIPTextEncode":
            wf["5"]["inputs"]["text"] = neg_prompt
        body = {"prompt": wf, "client_id": COMFYUI_CLIENT_ID}
        r = requests.post(f"{COMFYUI_HOST}/prompt", json=body, timeout=15)
        r.raise_for_status()
        return r.json()
    
    # --- TOOL 4: get_status ---
    def get_status(prompt_id: str, timeout_s: int = 180):
        start = time.time()
        while time.time() - start < timeout_s:
            r = requests.get(f"{COMFYUI_HOST}/history/{prompt_id}", timeout=10)
            if r.status_code == 200 and prompt_id in r.json():
                entry = r.json()[prompt_id]
                if entry.get("status", {}).get("completed"):
                    outputs = entry.get("outputs", {})
                    # SaveImage node is "9" in our template
                    images = outputs.get("9", {}).get("images", [])
                    return {"status": "done", "images": images}
            time.sleep(2)
        return {"status": "timeout"}
    
    # --- TOOL 5: fetch_image ---
    def fetch_image(filename: str, subfolder: str = "", img_type: str = "output"):
        params = {"filename": filename, "type": img_type}
        if subfolder:
            params["subfolder"] = subfolder
        r = requests.get(f"{COMFYUI_HOST}/view", params=params, timeout=15)
        r.raise_for_status()
        return base64.b64encode(r.content).decode()
    
    # --- TOOL SCHEMAS for OpenAI function calling ---
    TOOLS = [
        {"type": "function", "function": {"name": "list_models",
            "description": "List installed ComfyUI checkpoints, LoRAs, VAEs.",
            "parameters": {"type": "object", "properties": {}, "required": []}}},
        {"type": "function", "function": {"name": "get_queue",
            "description": "Get running and pending ComfyUI jobs.",
            "parameters": {"type": "object", "properties": {}, "required": []}}},
        {"type": "function", "function": {"name": "queue_workflow",
            "description": "Queue a ComfyUI workflow. Provide workflow_path, user_prompt, and optional neg_prompt.",
            "parameters": {"type": "object",
                "properties": {
                    "workflow_path": {"type": "string"},
                    "user_prompt": {"type": "string"},
                    "neg_prompt": {"type": "string"}},
                "required": ["workflow_path", "user_prompt"]}}},
        {"type": "function", "function": {"name": "get_status",
            "description": "Poll a queued job by prompt_id until completion.",
            "parameters": {"type": "object",
                "properties": {"prompt_id": {"type": "string"}},
                "required": ["prompt_id"]}}},
        {"type": "function", "function": {"name": "fetch_image",
            "description": "Fetch generated image bytes by filename.",
            "parameters": {"type": "object",
                "properties": {"filename": {"type": "string"}},
                "required": ["filename"]}}}
    ]
    
    # --- AGENT LOOP ---
    client = OpenAI()
    
    def run_agent(user_msg: str, workflow_path: str):
        messages = [{"role": "user", "content": user_msg}]
        while True:
            resp = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                tools=TOOLS,
                tool_choice="auto")
            msg = resp.choices[0].message
            if not msg.tool_calls:
                return msg.content
            messages.append(msg)
            for tc in msg.tool_calls:
                args = json.loads(tc.function.arguments)
                if tc.function.name == "queue_workflow":
                    args["workflow_path"] = workflow_path  # lock to caller
                result = globals()[tc.function.name](**args)
                messages.append({"role": "tool", "tool_call_id": tc.id,
                                 "content": json.dumps(result)})
    
    if __name__ == "__main__":
        print(run_agent("Generate a cute cat on a wooden table",
                        "/home/ubuntu/comfyui_workflows/flux_basic_api.json"))
    

    Reliability Patterns That Actually Help

    The five tools above will run. To make them production-grade, add three wrappers.

    1. VRAM guard

    Before queueing, hit GET /system_stats. If VRAM free is under 2GB, refuse and tell the user to wait. This single check eliminates 60% of OOM crashes in real use.

    def vram_free_gb():
        s = requests.get(f"{COMFYUI_HOST}/system_stats", timeout=5).json()
        devs = s.get("devices", [])
        if not devs:
            return 0
        free = devs[0].get("vram_free", 0) / (1024**3)
        return free
    

    2. Retry with backoff

    Wrap queue_workflow in a three-attempt retry. VRAM contention, transient CUDA errors, and server restarts all clear on the second or third try. Do not retry on JSON validation errors; the user prompt is malformed and retrying will not help.

    3. Idempotent client_id

    Reuse the same COMFYUI_CLIENT_ID across the entire agent session. ComfyUI uses this to send progress events over WebSocket to the right client. Random per-call client_ids break progress reporting without affecting correctness, but they do break observability.

    What to Watch Next

    ComfyUI 0.31 is in beta with a fully documented WebSocket event stream. The moment it ships, replace the get_status polling loop with a WebSocket subscription. That cuts typical end-to-end latency from 8 seconds (poll interval) to under 1 second (event-driven). For most agents, the 8-second poll is fine. For real-time chat products, the WebSocket path is mandatory.

    The wider trend: the agent layer is the new IDE. ComfyUI is becoming a library that agents compile against. If you ship a tool harness like the one above today, you can swap the backend to Flux, Wan, CogVideo, or whatever the next open-source video model is, without changing the agent prompt or the chat client. That is the durability you want from an AI workflow investment.


    Quick Reference

    • ComfyUI server: python main.py --listen 0.0.0.0 --port 8000
    • API format workflow: gear icon → Save (API Format)
    • Five tool functions: list_models, get_queue, queue_workflow, get_status, fetch_image
    • Failure modes: OOM (VRAM guard), timeout (retry), wrong model (list_models precheck)