cbpark84
1c919f1106
- run_batch(): sequential|concurrent|auto - chain(): 파이프라인 체이닝 - schema(): 스키마 자동 추출 - ResultCache: LRU+TTL 캐시 - preview_callback: 실시간 프리뷰 - upload_file/audio/video() - BridgeInputSEED: random/increment/fixed - MCPServer: MCP Tool 자동 생성
63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
"""WebSocket 통신 및 워크플로우 완료 대기."""
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from typing import Any, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def wait_for_completion(
|
|
ws,
|
|
target_prompt_id: str,
|
|
progress_callback: Optional[Callable[[float], None]] = None,
|
|
preview_callback: Optional[Callable[[bytes, int, int], None]] = None,
|
|
) -> None:
|
|
"""
|
|
WebSocket 메시지를 수신하며 워크플로우 완료를 기다린다.
|
|
|
|
완료 신호: { "type": "executing", "data": { "node": null } }
|
|
에러 신호: { "type": "execution_error" }
|
|
프리뷰 신호: 바이너리 메시지 (이미지)
|
|
"""
|
|
async for raw in ws:
|
|
try:
|
|
if isinstance(raw, bytes):
|
|
# 바이너리 프리뷰 메시지 (KSampler 중간 이미지)
|
|
if preview_callback and raw:
|
|
# 프리뷰 형식: [단계, 총단계] (4바이트) + PNG 바이너리
|
|
if len(raw) >= 8:
|
|
import struct
|
|
step, total = struct.unpack(">II", raw[:8])
|
|
preview_callback(raw[8:], step, total)
|
|
continue
|
|
data = json.loads(raw)
|
|
except (json.JSONDecodeError, Exception):
|
|
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}")
|