- ComfyUI custom nodes: BridgeInput (11 types) + BridgeOutput
- AnyType wildcard for universal output capture
- Custom server route: GET /comfybridge/results/{prompt_id}
- Python async client (ComfyBridge) + sync wrapper (ComfyBridgeSync)
- Workflow JSON parser and patcher (deepcopy, original unchanged)
- WebSocket completion detection
- BridgeResult container with PIL/bytes helpers
- Examples: basic_text_to_image.py, agent_loop.py
- Unit tests for workflow parser and client
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
"""WebSocket 통신 및 워크플로우 완료 대기."""
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def wait_for_completion(
|
|
ws,
|
|
target_prompt_id: str,
|
|
progress_callback: Callable[[float], None] | None = None,
|
|
) -> None:
|
|
"""
|
|
WebSocket 메시지를 수신하며 워크플로우 완료를 기다린다.
|
|
|
|
완료 신호: { "type": "executing", "data": { "node": null } }
|
|
에러 신호: { "type": "execution_error" }
|
|
"""
|
|
async for raw in ws:
|
|
try:
|
|
if isinstance(raw, bytes):
|
|
continue # 바이너리 프리뷰 메시지 무시
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
msg_type = data.get("type")
|
|
msg_data = data.get("data", {})
|
|
|
|
# 다른 prompt_id 메시지 무시
|
|
pid = msg_data.get("prompt_id")
|
|
if pid and pid != target_prompt_id:
|
|
continue
|
|
|
|
if msg_type == "executing":
|
|
node = msg_data.get("node")
|
|
if node is None:
|
|
logger.debug(f"[ComfyBridgeIO] Execution complete: {target_prompt_id}")
|
|
return # 완료
|
|
|
|
elif msg_type == "progress" and progress_callback:
|
|
value = msg_data.get("value", 0)
|
|
maximum = msg_data.get("max", 1)
|
|
if maximum > 0:
|
|
progress_callback(value / maximum)
|
|
|
|
elif msg_type == "execution_error":
|
|
err = msg_data.get("exception_message", "Unknown error")
|
|
raise RuntimeError(f"[ComfyBridgeIO] Execution error: {err}")
|