cbpark84-ComfyBridgeIO/client/bridge.py

"""ComfyBridge — ComfyUI 워크플로우 실행 클라이언트 (async)."""

from __future__ import annotations

import json
import logging
import uuid
from collections.abc import Callable
from pathlib import Path
from typing import Any

import aiohttp

from .result import BridgeResult
from .websocket_client import wait_for_completion
from .workflow import load_workflow, patch_inputs

logger = logging.getLogger(__name__)


class ComfyBridgeError(Exception):
    """ComfyBridge 관련 에러."""


class ComfyBridge:
    """
    ComfyUI 워크플로우 실행 클라이언트.

    사용 예시::

        bridge = ComfyBridge(host="localhost", port=8188)
        result = await bridge.run(
            "workflow_api.json",
            inputs={"prompt": "a cat", "steps": 20},
        )
        image = result.as_pil("output_image")
    """

    def __init__(
        self,
        host: str = "localhost",
        port: int = 8188,
        timeout: float = 300.0,
    ):
        self.base_url = f"http://{host}:{port}"
        self.ws_url = f"ws://{host}:{port}/ws"
        self.timeout = aiohttp.ClientTimeout(total=timeout)

    async def run(
        self,
        workflow: str | dict | Path,
        inputs: dict[str, Any] | None = None,
        progress_callback: Callable[[float], None] | None = None,
    ) -> BridgeResult:
        """
        워크플로우를 실행하고 결과를 반환한다.

        Args:
            workflow: 워크플로우 파일 경로 또는 파싱된 dict
            inputs:   BridgeInput 노드에 주입할 값 딕셔너리
            progress_callback: 진행률(0.0~1.0)을 받는 콜백 함수

        Returns:
            BridgeResult 객체
        """
        wf = load_workflow(workflow)
        patched = patch_inputs(wf, inputs or {})
        client_id = str(uuid.uuid4())

        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            # WebSocket 연결
            ws_url = f"{self.ws_url}?clientId={client_id}"
            async with session.ws_connect(ws_url) as ws:
                # /prompt API 호출
                prompt_id = await self._post_prompt(session, patched, client_id)
                logger.info(f"[ComfyBridgeIO] Queued prompt_id={prompt_id}")

                # 완료 대기
                await wait_for_completion(ws, prompt_id, progress_callback)

            # 결과 수집
            outputs = await self._collect_outputs(session, prompt_id)

        return BridgeResult(outputs, base_url=self.base_url, timeout=self.timeout.total or 300.0)

    async def upload_image(
        self,
        image: bytes | str | Path,
        filename: str = "input.png",
    ) -> str:
        """
        이미지를 ComfyUI 서버에 업로드하고 서버 파일명을 반환한다.

        Returns:
            서버 내 파일명 (BridgeInputIMAGE value에 주입할 값)
        """
        if isinstance(image, (str, Path)):
            with open(image, "rb") as f:
                image = f.read()

        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            form = aiohttp.FormData()
            form.add_field("image", image, filename=filename, content_type="image/png")
            async with session.post(f"{self.base_url}/upload/image", data=form) as resp:
                resp.raise_for_status()
                data = await resp.json()
                return data["name"]

    async def get_queue_status(self) -> dict:
        """현재 큐 상태를 반환한다."""
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.get(f"{self.base_url}/queue") as resp:
                resp.raise_for_status()
                return await resp.json()

    async def _post_prompt(
        self,
        session: aiohttp.ClientSession,
        workflow: dict,
        client_id: str,
    ) -> str:
        payload = {"prompt": workflow, "client_id": client_id}
        async with session.post(
            f"{self.base_url}/prompt",
            data=json.dumps(payload),
            headers={"Content-Type": "application/json"},
        ) as resp:
            resp.raise_for_status()
            data = await resp.json()
            return data["prompt_id"]

    async def _collect_outputs(
        self,
        session: aiohttp.ClientSession,
        prompt_id: str,
    ) -> dict[str, Any]:
        outputs: dict[str, Any] = {}

        # 1. /history에서 표준 이미지 결과 수집
        async with session.get(f"{self.base_url}/history/{prompt_id}") as resp:
            if resp.status == 200:
                history = await resp.json()
                node_outputs = history.get(prompt_id, {}).get("outputs", {})
                for node_id, node_out in node_outputs.items():
                    if "images" in node_out:
                        for i, img_info in enumerate(node_out["images"]):
                            params = {
                                "filename": img_info["filename"],
                                "subfolder": img_info.get("subfolder", ""),
                                "type": img_info.get("type", "output"),
                            }
                            async with session.get(
                                f"{self.base_url}/view", params=params
                            ) as img_resp:
                                if img_resp.status == 200:
                                    key = f"__node_{node_id}_image_{i}"
                                    outputs[key] = await img_resp.read()

        # 2. /comfybridge/results에서 BridgeOutput 결과 수집
        async with session.get(
            f"{self.base_url}/comfybridge/results/{prompt_id}"
        ) as resp:
            if resp.status == 200:
                data = await resp.json()
                outputs.update(data.get("outputs", {}))

        return outputs