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, Dict, Optional, Union

import aiohttp

from .cache import ResultCache
from .result import BridgeResult
from .schema import extract_schema
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: Union[str, dict, Path],
        inputs: Optional[Dict[str, Any]] = None,
        progress_callback: Optional[Callable[[float], None]] = None,
        preview_callback: Optional[Callable[[bytes, int, int], None]] = None,
    ) -> BridgeResult:
        """
        워크플로우를 실행하고 결과를 반환한다.

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

        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, preview_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: Union[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 run_batch(
        self,
        workflow: Union[str, dict, Path],
        inputs: list[Dict[str, Any]],
        mode: str = "auto",
        max_concurrent: int = 4,
    ) -> list[BridgeResult]:
        """
        여러 입력 셋을 배치로 실행한다.

        Args:
            workflow: 워크플로우 파일 경로 또는 dict
            inputs: 입력 딕셔너리 리스트
            mode: "sequential" (순차), "concurrent" (병렬), "auto" (영상이면 순차)
            max_concurrent: 동시 실행 수

        Returns:
            BridgeResult 리스트
        """
        if mode == "auto":
            wf = load_workflow(workflow)
            # 영상 출력이 있으면 순차 모드로
            mode = "sequential" if self._has_video_output(wf) else "concurrent"

        if mode == "sequential":
            return await self._run_batch_sequential(workflow, inputs)
        else:
            return await self._run_batch_concurrent(workflow, inputs, max_concurrent)

    async def chain(
        self,
        workflows: list[tuple[Union[str, dict, Path], Dict[str, Any]]],
        progress_callback: Optional[Callable[[float], None]] = None,
    ) -> BridgeResult:
        """
        여러 워크플로우를 순차적으로 체이닝한다.

        Args:
            workflows: (workflow_path, inputs) 튜플 리스트
            progress_callback: 진행률 콜백

        Returns:
            마지막 워크플로우 결과
        """
        result = None
        for i, (wf, inputs) in enumerate(workflows):
            logger.info(f"[ComfyBridgeIO] Chain step {i+1}/{len(workflows)}")
            if progress_callback:
                progress_callback(i / len(workflows))
            result = await self.run(wf, inputs, progress_callback)
        if progress_callback:
            progress_callback(1.0)
        return result

    async def upload_file(
        self,
        file_path: Union[str, Path],
        upload_type: str = "input",
    ) -> str:
        """
        범용 파일 업로드 (이미지/오디오/비디오).

        Args:
            file_path: 파일 경로
            upload_type: "input" | "temp" 등

        Returns:
            서버 내 파일명
        """
        path = Path(file_path)
        suffix = path.suffix.lower()

        if suffix in (".png", ".jpg", ".jpeg", ".webp"):
            return await self.upload_image(str(path))
        elif suffix in (".wav", ".mp3", ".aac", ".flac"):
            return await self.upload_audio(str(path))
        elif suffix in (".mp4", ".webm", ".mov", ".avi"):
            return await self.upload_video(str(path))
        else:
            raise ValueError(f"Unsupported file type: {suffix}")

    async def upload_audio(
        self,
        audio: Union[bytes, str, Path],
        filename: str = "input.wav",
    ) -> str:
        """오디오 파일을 업로드한다."""
        if isinstance(audio, (str, Path)):
            with open(audio, "rb") as f:
                audio = f.read()

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

    async def upload_video(
        self,
        video: Union[bytes, str, Path],
        filename: str = "input.mp4",
    ) -> str:
        """비디오 파일을 업로드한다."""
        if isinstance(video, (str, Path)):
            with open(video, "rb") as f:
                video = f.read()

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

    def get_schema(self, workflow: Union[str, dict, Path]) -> Dict[str, Any]:
        """워크플로우 스키마를 추출한다."""
        wf = load_workflow(workflow)
        return extract_schema(wf)

    def new_cache(self, max_size: int = 100, ttl: float = 3600.0) -> ResultCache:
        """새로운 ResultCache를 생성한다."""
        return ResultCache(max_size=max_size, ttl=ttl)

    # ── Private 메서드 ──────────────────────────────────────────

    async def _run_batch_sequential(
        self,
        workflow: Union[str, dict, Path],
        inputs: list[Dict[str, Any]],
    ) -> list[BridgeResult]:
        """배치를 순차적으로 실행한다."""
        results = []
        for i, input_dict in enumerate(inputs):
            logger.info(f"[ComfyBridgeIO] Batch {i+1}/{len(inputs)}")
            result = await self.run(workflow, input_dict)
            results.append(result)
        return results

    async def _run_batch_concurrent(
        self,
        workflow: Union[str, dict, Path],
        inputs: list[Dict[str, Any]],
        max_concurrent: int,
    ) -> list[BridgeResult]:
        """배치를 병렬로 실행한다."""
        import asyncio
        semaphore = asyncio.Semaphore(max_concurrent)

        async def _run_with_semaphore(input_dict):
            async with semaphore:
                return await self.run(workflow, input_dict)

        tasks = [_run_with_semaphore(inp) for inp in inputs]
        return await asyncio.gather(*tasks)

    def _has_video_output(self, workflow: dict) -> bool:
        """워크플로우에 비디오 출력이 있는지 확인한다."""
        for node in workflow.values():
            if not isinstance(node, dict):
                continue
            class_type = node.get("class_type", "").lower()
            if "video" in class_type or "vhs" in class_type:
                return True
        return False

    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