cbpark84-ComfyBridgeIO/client/result.py

"""BridgeResult — 워크플로우 실행 결과 컨테이너."""

from __future__ import annotations

import base64
import io
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

logger = logging.getLogger(__name__)


class BridgeResult:
    """
    워크플로우 실행 결과 컨테이너.

    BridgeOutput 노드에 지정한 이름으로 값을 회수한다.

    동기 메서드 (inline base64 있을 때):
        result.get("caption")           → str
        result.as_pil("small_image")    → PIL.Image (2MB 미만 이미지)
        result.as_bytes("small_image")  → bytes

    비동기 메서드 (모든 크기, URL 스트리밍 폴백):
        await result.as_pil_async("4k_image")   → PIL.Image
        await result.as_bytes_async("audio")    → bytes
        await result.download("video")          → bytes
        await result.save("video", "out.mp4")   → Path
    """

    def __init__(
        self,
        outputs: Dict[str, Any],
        base_url: str = "",
        timeout: float = 300.0,
    ):
        self._outputs = outputs
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout

    # ── 기본 접근 ────────────────────────────────────────────

    def get(self, name: str, default: Any = None) -> Any:
        """이름으로 결과 메타데이터를 반환한다."""
        return self._outputs.get(name, default)

    def keys(self) -> List[str]:
        """사용 가능한 출력 이름 목록."""
        return list(self._outputs.keys())

    def type_of(self, name: str) -> Optional[str]:
        """출력 타입 반환 ('image', 'audio', 'video', 'latent', 'str', ...)"""
        value = self._outputs.get(name)
        if value is None:
            return None
        if isinstance(value, dict):
            return value.get("__type__", "unknown")
        return type(value).__name__

    def has_inline(self, name: str) -> bool:
        """inline base64 데이터가 있는지 확인 (sync as_bytes 사용 가능 여부)."""
        value = self._outputs.get(name)
        return isinstance(value, dict) and "data" in value

    def has_url(self, name: str) -> bool:
        """스트리밍 URL이 있는지 확인."""
        value = self._outputs.get(name)
        return isinstance(value, dict) and "url" in value

    # ── 동기 메서드 (inline base64 only) ─────────────────────

    def as_pil(self, name: str):
        """
        이미지를 PIL Image로 반환한다 (동기, inline base64만 지원).
        4K 등 대형 이미지는 as_pil_async() 사용.
        """
        from PIL import Image
        data = self._get_inline_bytes(name)
        return Image.open(io.BytesIO(data))

    def as_bytes(self, name: str) -> bytes:
        """
        결과를 bytes로 반환한다 (동기, inline base64만 지원).
        대형 파일은 await result.download() 또는 as_bytes_async() 사용.
        """
        return self._get_inline_bytes(name)

    def _get_inline_bytes(self, name: str) -> bytes:
        value = self._outputs.get(name)
        if value is None:
            raise KeyError(f"Output '{name}' not found. Available: {self.keys()}")
        if isinstance(value, bytes):
            return value
        if isinstance(value, dict):
            if "data" in value:
                return base64.b64decode(value["data"])
            if "url" in value:
                t = self.type_of(name)
                raise RuntimeError(
                    f"'{name}' ({t}) is too large for inline transfer. "
                    f"Use: `await result.download('{name}')` or `await result.as_bytes_async('{name}')`"
                )
        raise TypeError(f"'{name}' is not a binary type: {type(value)}")

    # ── 비동기 메서드 (URL 스트리밍 폴백) ────────────────────

    async def as_pil_async(self, name: str):
        """이미지를 PIL Image로 반환한다 (비동기, 모든 크기 지원)."""
        from PIL import Image
        data = await self.as_bytes_async(name)
        return Image.open(io.BytesIO(data))

    async def as_bytes_async(self, name: str) -> bytes:
        """결과를 bytes로 반환한다 (비동기, 모든 크기 지원)."""
        value = self._outputs.get(name)
        if value is None:
            raise KeyError(f"Output '{name}' not found. Available: {self.keys()}")
        if isinstance(value, bytes):
            return value
        if isinstance(value, dict):
            # inline base64 우선
            if "data" in value:
                return base64.b64decode(value["data"])
            # URL 스트리밍 폴백
            if "url" in value:
                return await self.download(name)
        raise TypeError(f"'{name}' is not a binary type: {type(value)}")

    async def download(self, name: str) -> bytes:
        """
        바이너리 결과를 서버 URL에서 스트리밍 다운로드한다.

        이미지(4K 포함) / 오디오 / 영상 모두 지원.
        """
        value = self._outputs.get(name)
        if value is None:
            raise KeyError(f"Output '{name}' not found. Available: {self.keys()}")

        # 이미 bytes라면 바로 반환
        if isinstance(value, bytes):
            return value

        # inline base64 우선 (이미 캐시됨)
        if isinstance(value, dict) and "data" in value:
            return base64.b64decode(value["data"])

        # URL 스트리밍
        if isinstance(value, dict) and "url" in value:
            url = f"{self._base_url}{value['url']}"
            return await self._fetch_url(url)

        raise ValueError(f"'{name}' has no downloadable data: {value}")

    async def save(self, name: str, path: Union[str, Path]) -> Path:
        """
        결과를 파일로 저장한다.

        사용 예::
            await result.save("result_video", "output.mp4")
            await result.save("result_audio", "output.wav")
            await result.save("result_image", "4k_output.png")
        """
        data = await self.download(name)
        out = Path(path)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_bytes(data)
        logger.info(f"[ComfyBridgeIO] Saved '{name}' → {out} ({len(data):,} bytes)")
        return out

    async def _fetch_url(self, url: str) -> bytes:
        """URL에서 바이너리를 다운로드한다."""
        import aiohttp
        timeout = aiohttp.ClientTimeout(total=self._timeout)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(url) as resp:
                if resp.status == 404:
                    raise FileNotFoundError(f"Server returned 404 for: {url}")
                resp.raise_for_status()
                return await resp.read()

    # ── 편의 메서드 ──────────────────────────────────────────

    def info(self, name: str) -> Dict:
        """출력의 메타데이터를 반환한다 (크기, 포맷, URL 등)."""
        value = self._outputs.get(name)
        if isinstance(value, dict):
            # data 필드는 크므로 제외
            return {k: v for k, v in value.items() if k != "data"}
        return {"value": value, "__type__": type(value).__name__}

    def summary(self) -> Dict:
        """모든 출력의 요약 정보를 반환한다."""
        return {name: self.info(name) for name in self.keys()}

    def __repr__(self) -> str:
        parts = []
        for name in self.keys():
            t = self.type_of(name)
            inline = " [inline]" if self.has_inline(name) else ""
            url = " [url]" if self.has_url(name) else ""
            parts.append(f"  {name}: {t}{inline}{url}")
        return "BridgeResult(\n" + "\n".join(parts) + "\n)"