# ComfyBridgeIO 아키텍처 설계

## 1. 프로젝트 개요

### 목적
ComfyBridgeIO는 ComfyUI 커스텀 노드 패키지와 Python 클라이언트 래퍼로 구성된다. AI Agent나 외부 스크립트가 ComfyUI 워크플로우를 **런타임에 동적으로 제어**할 수 있도록 표준화된 입출력 브리지를 제공한다.

### 핵심 동작 원리

```
[AI Agent / 외부 스크립트]
        │
        │  bridge.run("workflow.json", inputs={...})
        ▼
[ComfyBridge 클라이언트]
  1. workflow.json 파싱 → BridgeInput 노드 탐색
  2. 노드 inputs.value 필드 패치 (런타임 값 주입)
  3. POST /prompt  →  ComfyUI 서버에 실행 요청
  4. WebSocket 수신  →  실행 완료 대기
  5. GET /history/{prompt_id}  →  BridgeOutput 노드 결과 수집
  6. GET /view  or  GET /comfybridge/results/{prompt_id}  →  바이너리/스칼라 회수
        │
        │  BridgeResult 객체 반환
        ▼
[AI Agent]
  result.get("output_image")  →  PIL.Image / bytes
  result.get("caption")       →  str
```

### 설계 원칙
- **YAGNI**: 현재 요구사항 범위 외 기능 추가 금지
- **ComfyUI 버전 호환성**: AnyType 트릭을 사용해 와일드카드 타입 안정성 확보
- **비동기 우선**: 클라이언트는 `asyncio` 기반, 동기 래퍼 별도 제공
- **단방향 패치**: 클라이언트가 워크플로우 JSON을 수정하되 원본 파일은 변경하지 않음

---

## 2. 디렉토리 구조

```
ComfyBridgeIO/
├── __init__.py                    # ComfyUI 노드 등록 진입점
├── nodes/
│   ├── __init__.py
│   ├── bridge_input.py            # BridgeInput 노드 (타입별 팩토리 생성)
│   ├── bridge_output.py           # BridgeOutput 노드 (타입별 팩토리 생성)
│   ├── any_type.py                # AnyType 와일드카드 헬퍼
│   └── server_routes.py           # 커스텀 API 라우트 (/comfybridge/*)
├── client/
│   ├── __init__.py                # ComfyBridge, BridgeResult export
│   ├── bridge.py                  # ComfyBridge 메인 클래스 (async)
│   ├── bridge_sync.py             # 동기 래퍼 (asyncio.run 기반)
│   ├── workflow.py                # 워크플로우 JSON 파서 / 패처
│   ├── websocket_client.py        # WebSocket 통신 및 완료 대기
│   ├── result.py                  # BridgeResult 클래스
│   └── types.py                   # 타입 상수, 변환 유틸리티
├── examples/
│   ├── basic_text_to_image.py     # 기본 텍스트-이미지 예제
│   ├── agent_loop.py              # AI Agent 루프 예제
│   └── workflows/
│       └── sdxl_bridge_template.json
├── tests/
│   ├── __init__.py
│   ├── test_workflow_parser.py
│   ├── test_bridge_client.py
│   └── fixtures/
│       └── sample_workflow.json
├── pyproject.toml
├── requirements.txt               # 클라이언트 의존성
├── requirements-dev.txt
├── README.md
└── ARCHITECTURE.md
```

---

## 3. 커스텀 노드 설계

### 3-1. AnyType 와일드카드 (`nodes/any_type.py`)

ComfyUI의 타입 검증 시스템을 우회하기 위해 AnyType 클래스를 사용한다. `__ne__` 를 항상 `False` 로 반환하여 어떤 타입과도 동등 비교되도록 한다.

```python
class AnyType(str):
    """
    ComfyUI 타입 검증을 통과하는 와일드카드 타입.
    str을 상속하여 직렬화 호환성 유지.
    """
    def __ne__(self, other: object) -> bool:
        return False

ANY = AnyType("*")
```

이 타입은 BridgeOutput 노드의 입력부에 사용하며, 어떤 타입의 ComfyUI 출력도 연결 가능하게 한다.

---

### 3-2. BridgeInput 노드 (`nodes/bridge_input.py`)

#### 책임
워크플로우 안에서 **외부에서 주입 가능한 값의 입력 포인트** 역할을 한다. 클라이언트가 이 노드를 이름으로 탐색하여 `value` 위젯 값을 패치한다.

#### 타입별 팩토리 패턴

동일한 로직을 반복하지 않기 위해 팩토리 함수로 각 타입별 노드 클래스를 생성한다.

```python
SUPPORTED_TYPES = {
    "STRING":       ("STRING",       {"default": "", "multiline": True}),
    "INT":          ("INT",          {"default": 0, "min": -2**31, "max": 2**31}),
    "FLOAT":        ("FLOAT",        {"default": 0.0, "min": -1e9, "max": 1e9}),
    "BOOLEAN":      ("BOOLEAN",      {"default": False}),
    "IMAGE":        ("IMAGE",        {}),   # 위젯 없음, 입력 포트만
    "MASK":         ("MASK",         {}),
    "LATENT":       ("LATENT",       {}),
    "MODEL":        ("MODEL",        {}),
    "CLIP":         ("CLIP",         {}),
    "VAE":          ("VAE",          {}),
    "CONDITIONING": ("CONDITIONING", {}),
}

def make_bridge_input(type_name: str, comfy_type, widget_config: dict):
    """
    타입별 BridgeInput 노드 클래스를 동적으로 생성한다.
    """
    is_primitive = type_name in ("STRING", "INT", "FLOAT", "BOOLEAN")

    class BridgeInputNode:
        CATEGORY = "ComfyBridgeIO"
        FUNCTION = "passthrough"
        RETURN_TYPES = (comfy_type,)
        RETURN_NAMES = ("value",)

        @classmethod
        def INPUT_TYPES(cls):
            inputs = {
                "required": {
                    "name": ("STRING", {"default": "input_name"}),
                }
            }
            if is_primitive:
                # 스칼라 타입: 위젯으로 기본값 설정
                inputs["required"]["value"] = (comfy_type, widget_config)
            else:
                # 복합 타입(IMAGE 등): 선택적 연결 포트
                inputs["optional"] = {
                    "value": (comfy_type, {})
                }
            return inputs

        def passthrough(self, name: str, value=None):
            return (value,)

    BridgeInputNode.__name__ = f"BridgeInput{type_name}"
    BridgeInputNode.__qualname__ = f"BridgeInput{type_name}"
    return BridgeInputNode
```

#### 클라이언트와의 계약

| 노드 필드 | 역할 |
|-----------|------|
| `class_type` | `"BridgeInputSTRING"`, `"BridgeInputINT"` 등 (패턴: `BridgeInput*`) |
| `inputs.name` | 클라이언트 딕셔너리의 키와 매칭되는 식별자 |
| `inputs.value` | 클라이언트가 런타임에 패치하는 필드 |

---

### 3-3. BridgeOutput 노드 (`nodes/bridge_output.py`)

#### 책임
워크플로우 안에서 **결과를 이름으로 태깅하는 출력 포인트** 역할을 한다. `OUTPUT_NODE = True` 로 선언하여 ComfyUI가 이 노드의 실행을 보장하게 한다. 실행 후 결과를 서버 메모리에 저장하고 커스텀 API 라우트로 노출한다.

#### 설계

```python
from .any_type import ANY
from .server_routes import store_bridge_result

class BridgeOutput:
    CATEGORY = "ComfyBridgeIO"
    FUNCTION = "capture"
    OUTPUT_NODE = True
    RETURN_TYPES = ()      # 터미널 노드: 다운스트림 연결 없음

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "name":  ("STRING", {"default": "output_name"}),
                "value": (ANY, {}),
            },
            "hidden": {
                "prompt_id": "PROMPT_ID",   # ComfyUI가 런타임에 주입
                "unique_id": "UNIQUE_ID",
            }
        }

    def capture(self, name: str, value, prompt_id: str, unique_id: str):
        store_bridge_result(prompt_id=prompt_id, name=name, value=value)
        return {}
```

**IMAGE 타입 처리**: BridgeOutput이 IMAGE 값을 받으면 `store_bridge_result` 내부에서 PIL Image로 변환 후 메모리에 보관. 클라이언트는 `/comfybridge/results/{prompt_id}` 로 base64 인코딩된 PNG를 수신.

---

### 3-4. 서버 커스텀 라우트 (`nodes/server_routes.py`)

ComfyUI의 `PromptServer`에 라우트를 추가하여 BridgeOutput 결과를 외부로 노출한다.

```
GET /comfybridge/results/{prompt_id}
    → { "outputs": { "output_name": <value>, ... } }
    → IMAGE 타입은 base64 PNG 문자열로 직렬화
    → 스칼라 타입은 JSON 네이티브 값으로 직렬화
```

결과 저장소: 모듈 레벨 `dict[prompt_id → dict[name → value]]`. 최대 100개 prompt_id 보관 (LRU 방식으로 오래된 항목 제거).

---

### 3-5. 노드 등록 (`__init__.py`)

```python
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}

# BridgeInput 타입별 등록
for type_name, (comfy_type, widget_cfg) in SUPPORTED_TYPES.items():
    cls = make_bridge_input(type_name, comfy_type, widget_cfg)
    node_id = f"BridgeInput{type_name}"
    NODE_CLASS_MAPPINGS[node_id] = cls
    NODE_DISPLAY_NAME_MAPPINGS[node_id] = f"Bridge Input ({type_name})"

# BridgeOutput 등록
NODE_CLASS_MAPPINGS["BridgeOutput"] = BridgeOutput
NODE_DISPLAY_NAME_MAPPINGS["BridgeOutput"] = "Bridge Output"
```

---

## 4. Python 클라이언트 설계

### 4-1. ComfyBridge 클래스 API (`client/bridge.py`)

```python
class ComfyBridge:
    def __init__(
        self,
        host: str = "localhost",
        port: int = 8188,
        timeout: float = 300.0,
    ): ...

    async def run(
        self,
        workflow: str | dict,        # 파일 경로 또는 파싱된 dict
        inputs: dict[str, Any] = {},
        progress_callback: Callable[[float], None] | None = None,
    ) -> BridgeResult: ...

    async def upload_image(
        self,
        image: bytes | str,          # bytes 또는 파일 경로
        filename: str = "input.png",
    ) -> str: ...                    # ComfyUI 서버 내 파일명 반환

    async def get_queue_status(self) -> dict: ...

    async def cancel(self, prompt_id: str) -> None: ...
```

#### 동기 래퍼 (`client/bridge_sync.py`)

```python
class ComfyBridgeSync:
    """asyncio가 없는 환경을 위한 동기 래퍼."""

    def run(self, workflow, inputs={}, **kwargs) -> BridgeResult:
        return asyncio.run(self._bridge.run(workflow, inputs, **kwargs))
```

---

### 4-2. 워크플로우 파서 / 패처 (`client/workflow.py`)

#### 파싱 전략

ComfyUI 워크플로우 JSON (API 포맷)은 노드 ID를 키로 하는 flat dict다.

```json
{
  "1": { "class_type": "BridgeInputSTRING", "inputs": { "name": "prompt", "value": "default text" } },
  "2": { "class_type": "KSampler", "inputs": { "text": ["1", 0], ... } },
  "9": { "class_type": "BridgeOutput", "inputs": { "name": "output_image", "value": ["8", 0] } }
}
```

#### BridgeInput 노드 탐색

```python
BRIDGE_INPUT_PREFIX = "BridgeInput"

def find_bridge_inputs(workflow: dict) -> dict[str, tuple[str, str]]:
    """
    Returns: { "input_name": (node_id, type_name) }
    예: { "prompt": ("1", "STRING"), "steps": ("3", "INT") }
    """
    result = {}
    for node_id, node in workflow.items():
        class_type = node.get("class_type", "")
        if class_type.startswith(BRIDGE_INPUT_PREFIX):
            name = node["inputs"].get("name")
            type_name = class_type[len(BRIDGE_INPUT_PREFIX):]
            if name:
                result[name] = (node_id, type_name)
    return result
```

#### 값 주입 (패치)

```python
def patch_inputs(workflow: dict, inputs: dict[str, Any]) -> dict:
    """
    원본 workflow를 deepcopy 후 BridgeInput 노드의 value를 패치한다.
    IMAGE 타입: 이미 upload_image()로 업로드된 파일명(str)을 주입.
    """
    patched = copy.deepcopy(workflow)
    bridge_map = find_bridge_inputs(patched)

    for name, value in inputs.items():
        if name not in bridge_map:
            raise ValueError(f"BridgeInput '{name}' not found in workflow")
        node_id, type_name = bridge_map[name]
        converted = convert_value(value, type_name)
        patched[node_id]["inputs"]["value"] = converted

    return patched
```

#### 타입 변환 테이블 (`client/types.py`)

| type_name | Python 입력 타입 | 변환 |
|-----------|-----------------|------|
| STRING | str | 그대로 |
| INT | int / str | int() 캐스팅 |
| FLOAT | float / str | float() 캐스팅 |
| BOOLEAN | bool / int | bool() 캐스팅 |
| IMAGE | str (파일명) | 업로드 후 반환된 서버 파일명 |
| MASK, LATENT, MODEL, CLIP, VAE, CONDITIONING | 노드 링크 `[node_id, idx]` | 그대로 전달 |

---

### 4-3. WebSocket 통신 흐름 (`client/websocket_client.py`)

```
클라이언트                              ComfyUI 서버
    │                                       │
    ├─── GET /ws?clientId={uuid} ──────────►│  WebSocket 연결 수립
    │                                       │
    ├─── POST /prompt ──────────────────────►│  { prompt, client_id, extra_data }
    │◄── { prompt_id, number, ... } ─────────┤  즉시 응답 (큐에 삽입)
    │                                       │
    │◄── { type: "execution_start" } ────────┤  실행 시작
    │◄── { type: "executing",               │
    │       data: { node: "1" } } ───────────┤  노드별 진행 상황
    │◄── { type: "progress",               │
    │       data: { value: 5, max: 20 } } ───┤  KSampler 스텝 진행률
    │◄── { type: "executed",               │
    │       data: { node: "9", ... } } ──────┤  OUTPUT_NODE 결과 (SaveImage 등)
    │◄── { type: "executing",               │
    │       data: { node: null } } ──────────┤  완료 신호
    │                                       │
    ├─── GET /history/{prompt_id} ──────────►│
    │◄── { prompt_id: { outputs: {...} } } ──┤
    │                                       │
    ├─── GET /comfybridge/results/{id} ─────►│  비이미지 BridgeOutput 결과
    │◄── { outputs: { "name": value } } ─────┤
```

#### 완료 감지 조건

```python
# WebSocket 메시지 처리 루프
async def wait_for_completion(ws, target_prompt_id: str) -> None:
    async for message in ws:
        data = json.loads(message)
        msg_type = data.get("type")
        msg_data = data.get("data", {})

        # 다른 prompt의 메시지 무시
        if msg_data.get("prompt_id") and msg_data["prompt_id"] != target_prompt_id:
            continue

        if msg_type == "executing" and msg_data.get("node") is None:
            break   # node가 null이면 이 prompt의 실행 완료

        if msg_type == "execution_error":
            raise ComfyBridgeError(msg_data.get("exception_message"))
```

---

### 4-4. 결과 회수 전략 (`client/result.py`)

#### BridgeResult 클래스

```python
class BridgeResult:
    def __init__(self, raw_outputs: dict[str, Any]):
        self._outputs = raw_outputs

    def get(self, name: str, default=None) -> Any:
        return self._outputs.get(name, default)

    def keys(self) -> list[str]:
        return list(self._outputs.keys())

    def as_pil(self, name: str) -> "PIL.Image.Image":
        """IMAGE 타입 결과를 PIL Image로 반환."""
        ...

    def as_bytes(self, name: str) -> bytes:
        """IMAGE 타입 결과를 PNG bytes로 반환."""
        ...
```

#### 결과 수집 흐름

```python
async def collect_results(
    session: aiohttp.ClientSession,
    base_url: str,
    prompt_id: str,
    workflow: dict,
) -> BridgeResult:
    outputs = {}

    # 1. /history/{prompt_id} 에서 SaveImage / PreviewImage 등 표준 노드 결과 수집
    history = await fetch_history(session, base_url, prompt_id)
    node_outputs = history[prompt_id]["outputs"]

    for node_id, node_output in node_outputs.items():
        if "images" in node_output:
            for img_info in node_output["images"]:
                # /view?filename=...&subfolder=...&type=... 로 다운로드
                img_bytes = await fetch_image(session, base_url, img_info)
                outputs[f"__node_{node_id}_image"] = img_bytes

    # 2. /comfybridge/results/{prompt_id} 에서 BridgeOutput 노드 결과 수집
    bridge_outputs = await fetch_bridge_results(session, base_url, prompt_id)
    outputs.update(bridge_outputs)

    return BridgeResult(outputs)
```

---

## 5. 데이터 플로우 다이어그램

```
[AI Agent]
    │
    │  1. bridge.run("workflow.json", inputs={"prompt": "a cat", "steps": 20})
    ▼
[ComfyBridge.run()]
    │
    ├─ 2. workflow.py: JSON 로드 → BridgeInput 노드 탐색
    ├─ 3. workflow.py: inputs dict → 노드 value 필드 패치 (deepcopy, 원본 불변)
    ├─ 4. IMAGE 타입이 있으면: POST /upload/image → 서버 파일명 획득 후 패치
    │
    ├─ 5. uuid4() → client_id 생성
    ├─ 6. WebSocket 연결 수립: ws://host:port/ws?clientId={client_id}
    │
    ├─ 7. POST /prompt → { prompt: patched_workflow, client_id }
    │       ← { prompt_id: "abc-123" }
    │
    ├─ 8. WebSocket 수신 루프 (완료 대기)
    │       progress_callback 호출 (선택)
    │       execution_error 시 ComfyBridgeError raise
    │
    ├─ 9. GET /history/{prompt_id}  →  표준 노드 이미지 결과 수집
    │       GET /view?filename=...  →  이미지 바이너리 다운로드
    │
    ├─ 10. GET /comfybridge/results/{prompt_id}  →  BridgeOutput 스칼라 결과 수집
    │
    └─ 11. BridgeResult 객체 생성 후 반환
    │
    ▼
[AI Agent]
    result.get("output_image")  # bytes (PNG)
    result.get("caption")       # str
    result.get("latent_code")   # dict (LATENT)
```

---

## 6. 설치 및 사용법 개요

### 커스텀 노드 설치

```bash
# ComfyUI 커스텀 노드 디렉토리에 클론
cd /path/to/ComfyUI/custom_nodes
git clone https://github.com/your-org/ComfyBridgeIO
# ComfyUI 재시작
```

### Python 클라이언트 설치

```bash
pip install comfybridgeio
# 또는 로컬 설치
pip install -e /path/to/ComfyBridgeIO
```

### 워크플로우 준비

1. ComfyUI에서 워크플로우 편집
2. 외부에서 주입할 값의 위치에 **Bridge Input (STRING)** 등 노드 삽입, `name` 위젯에 식별자 입력
3. 외부로 회수할 결과에 **Bridge Output** 노드 연결, `name` 위젯에 식별자 입력
4. `Save (API Format)` 으로 JSON 내보내기

### 클라이언트 사용 예시

```python
import asyncio
from comfybridgeio import ComfyBridge

async def main():
    bridge = ComfyBridge(host="localhost", port=8188)

    result = await bridge.run(
        workflow="workflow_api.json",
        inputs={
            "prompt":          "a cat in space, photorealistic",
            "negative_prompt": "blurry, low quality",
            "steps":           20,
            "cfg":             7.5,
            "seed":            42,
        },
        progress_callback=lambda p: print(f"Progress: {p:.0%}"),
    )

    # 이미지 저장
    with open("output.png", "wb") as f:
        f.write(result.as_bytes("output_image"))

    print("Done")

asyncio.run(main())
```

```python
# 동기 사용 예시
from comfybridgeio import ComfyBridgeSync

bridge = ComfyBridgeSync(host="localhost", port=8188)
result = bridge.run("workflow_api.json", inputs={"prompt": "a dog"})
image = result.as_pil("output_image")
image.show()
```

---

## 7. 의존성

### requirements.txt (클라이언트)

```
aiohttp>=3.9.0
websockets>=12.0
Pillow>=10.0.0
```

### requirements-dev.txt

```
pytest>=8.0.0
pytest-asyncio>=0.23.0
pytest-mock>=3.12.0
aioresponses>=0.7.6
```

### pyproject.toml

```toml
[project]
name = "comfybridgeio"
version = "0.1.0"
description = "ComfyUI dynamic I/O bridge for AI agents"
requires-python = ">=3.10"
dependencies = [
    "aiohttp>=3.9.0",
    "websockets>=12.0",
    "Pillow>=10.0.0",
]

[project.optional-dependencies]
dev = [
    "pytest>=8.0.0",
    "pytest-asyncio>=0.23.0",
    "pytest-mock>=3.12.0",
    "aioresponses>=0.7.6",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
```

---

## 8. 미결 사항 (PM 승인 필요)

| # | 항목 | 선택지 |
|---|------|--------|
| 1 | MODEL / CLIP / VAE / CONDITIONING 타입의 BridgeOutput 직렬화 | (A) 미지원 (에러 반환) / (B) JSON-safe repr만 반환 |
| 2 | BridgeOutput 결과 서버 메모리 보관 최대 개수 | 기본 100개 LRU / 조정 가능 |
| 3 | 클라이언트 패키지 PyPI 배포 여부 | 배포 / GitHub only |
| 4 | ComfyUI Manager 연동 (`node_list.json` 포함 여부) | 포함 / 미포함 |