diff --git a/client/bridge.py b/client/bridge.py index cf11bc5..72c2860 100644 --- a/client/bridge.py +++ b/client/bridge.py @@ -1,5 +1,7 @@ """ComfyBridge — ComfyUI 워크플로우 실행 클라이언트 (async).""" +from __future__ import annotations + import json import logging import uuid @@ -79,7 +81,7 @@ class ComfyBridge: # 결과 수집 outputs = await self._collect_outputs(session, prompt_id) - return BridgeResult(outputs) + return BridgeResult(outputs, base_url=self.base_url, timeout=self.timeout.total or 300.0) async def upload_image( self, diff --git a/client/result.py b/client/result.py index 25fd44d..01dacba 100644 --- a/client/result.py +++ b/client/result.py @@ -1,47 +1,203 @@ """BridgeResult — 워크플로우 실행 결과 컨테이너.""" +from __future__ import annotations + import base64 import io +import logging +from pathlib import Path from typing import Any +logger = logging.getLogger(__name__) + class BridgeResult: """ - 워크플로우 실행 결과를 담는 컨테이너. + 워크플로우 실행 결과 컨테이너. - BridgeOutput 노드에 지정한 이름으로 값을 회수할 수 있다. + 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]): + 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) -> str | None: + """출력 타입 반환 ('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): - """IMAGE 타입 결과를 PIL Image로 반환한다.""" + """ + 이미지를 PIL Image로 반환한다 (동기, inline base64만 지원). + 4K 등 대형 이미지는 as_pil_async() 사용. + """ from PIL import Image - value = self._outputs[name] - if isinstance(value, bytes): - return Image.open(io.BytesIO(value)) - if isinstance(value, dict) and value.get("__type__") == "image": - data = base64.b64decode(value["data"]) - return Image.open(io.BytesIO(data)) - raise TypeError(f"'{name}' is not an image type: {type(value)}") + data = self._get_inline_bytes(name) + return Image.open(io.BytesIO(data)) def as_bytes(self, name: str) -> bytes: - """IMAGE 타입 결과를 PNG bytes로 반환한다.""" - value = self._outputs[name] + """ + 결과를 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) and value.get("__type__") == "image": + 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"]) - raise TypeError(f"'{name}' is not an image type: {type(value)}") + + # 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: 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: - return f"BridgeResult(keys={self.keys()})" + 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)" diff --git a/nodes/server_routes.py b/nodes/server_routes.py index 73fda22..e2bfc66 100644 --- a/nodes/server_routes.py +++ b/nodes/server_routes.py @@ -2,198 +2,421 @@ ComfyUI PromptServer에 마운트되는 커스텀 라우트. GET /comfybridge/results/{prompt_id} - → BridgeOutput 노드가 저장한 결과를 JSON으로 반환. - → IMAGE 타입은 base64 PNG 문자열로 직렬화. - → 스칼라 타입은 JSON 네이티브 값. + → BridgeOutput 결과 메타데이터 (JSON) + → 소형 이미지: base64 인라인 + url + → 대형 이미지/오디오/영상: 메타데이터 + url만 + +GET /comfybridge/file/{prompt_id}/{name} + → 바이너리 스트리밍 (image/png, audio/wav, video/mp4, application/zip) """ +from __future__ import annotations + import base64 import io import json import logging +import struct +import zipfile from collections import OrderedDict +from pathlib import Path logger = logging.getLogger(__name__) -# 결과 저장소: { prompt_id: { name: value } } +# ── 저장소 ──────────────────────────────────────────────────── _MAX_ENTRIES = 100 + +# 메타데이터 저장소: { prompt_id: { name: serialized_dict_or_scalar } } _results_store: OrderedDict[str, dict] = OrderedDict() +# 바이너리 저장소: { prompt_id: { name: (bytes, content_type) } } +_binary_store: OrderedDict[str, dict[str, tuple[bytes, str]]] = OrderedDict() + +# 소형 이미지 임계값 (원본 raw 바이트 기준) +_INLINE_THRESHOLD = 2 * 1024 * 1024 # 2 MB + + +# ── 저장 진입점 ─────────────────────────────────────────────── def store_bridge_result(prompt_id: str, name: str, value) -> None: - """BridgeOutput 노드 실행 시 결과를 저장한다.""" + """BridgeOutput 노드 실행 시 값을 인코딩하여 저장한다.""" if prompt_id not in _results_store: _results_store[prompt_id] = {} - # LRU: 최대 개수 초과 시 가장 오래된 항목 제거 + _binary_store[prompt_id] = {} + # LRU 제거 if len(_results_store) > _MAX_ENTRIES: + oldest = next(iter(_results_store)) _results_store.popitem(last=False) - _results_store[prompt_id][name] = value - logger.debug(f"[ComfyBridgeIO] stored result: prompt_id={prompt_id}, name={name}") + _binary_store.pop(oldest, None) + + meta, binary = _encode_value(prompt_id, name, value) + _results_store[prompt_id][name] = meta + if binary is not None: + raw_bytes, content_type = binary + _binary_store[prompt_id][name] = (raw_bytes, content_type) + + logger.debug(f"[ComfyBridgeIO] stored: prompt_id={prompt_id}, name={name}, type={type(value).__name__}") -def _serialize_value(value) -> object: +# ── 인코딩 핵심 ─────────────────────────────────────────────── + +def _encode_value(prompt_id: str, name: str, value) -> tuple[object, tuple[bytes, str] | None]: """ - 결과 값을 JSON 직렬화 가능한 형태로 변환한다. - - 타입별 직렬화 전략: - - IMAGE (torch.Tensor [B,H,W,C]): base64 PNG - - LATENT (dict{"samples": Tensor}): base64 PNG preview + shape 메타 - - CONDITIONING (list of [Tensor, dict]): 각 조건부 텐서 shape 메타 반환 - - MODEL/CLIP/VAE: 클래스명 + 메타정보 반환 (가중치는 직렬화 불가) - - 스칼라: 그대로 반환 - - PIL Image: base64 PNG + 값을 (JSON 직렬화 가능한 메타, (bytes, content_type) | None) 로 변환한다. + 바이너리가 없으면 None 반환. + """ + file_url = f"/comfybridge/file/{prompt_id}/{name}" + + # ── JSON 네이티브 스칼라 ───────────────────────────────── + if isinstance(value, (str, int, float, bool, type(None))): + return value, None + + # ── torch.Tensor ───────────────────────────────────────── + try: + import torch + if isinstance(value, torch.Tensor): + t = value + # IMAGE: [B, H, W, C] 또는 [H, W, C] + if t.dim() in (3, 4): + return _encode_image_tensor(t, file_url) + # VIDEO: [B, F, H, W, C] (드물지만 일부 노드) + if t.dim() == 5: + return _encode_video_frames(t[0], fps=24.0, file_url=file_url) + # 기타 텐서 + return {"__type__": "tensor", "shape": list(t.shape), "dtype": str(t.dtype)}, None + except ImportError: + pass + + # ── PIL Image ──────────────────────────────────────────── + try: + from PIL import Image as PILImage + if isinstance(value, PILImage.Image): + return _encode_pil_image(value, file_url) + except ImportError: + pass + + # ── AUDIO dict: {"waveform": Tensor [B, C, T], "sample_rate": int} ── + if isinstance(value, dict) and "waveform" in value and "sample_rate" in value: + return _encode_audio(value, file_url) + + # ── VIDEO dict: {"images": Tensor [F,H,W,C], "frame_rate": float, ...} ── + if isinstance(value, dict) and ("images" in value or "frames" in value): + frames_key = "images" if "images" in value else "frames" + fps = float(value.get("frame_rate", value.get("fps", 24.0))) + return _encode_video_frames(value[frames_key], fps=fps, file_url=file_url) + + # ── LATENT dict: {"samples": Tensor [B,C,H,W]} ────────── + if isinstance(value, dict) and "samples" in value: + return _encode_latent(value), None + + # ── CONDITIONING: list of [Tensor, dict] ───────────────── + if isinstance(value, list) and len(value) > 0 and isinstance(value[0], (list, tuple)): + return _encode_conditioning(value), None + + # ── 파일 경로 (str / Path) ──────────────────────────────── + if isinstance(value, (str, Path)): + p = Path(value) + if p.exists() and p.is_file(): + return _encode_file(p, file_url) + + # ── ComfyUI 내부 객체 (MODEL / CLIP / VAE 등) ───────────── + return _encode_comfy_object(value), None + + +# ── 이미지 인코딩 ───────────────────────────────────────────── + +def _encode_image_tensor(t, file_url: str): + """torch.Tensor [B,H,W,C] or [H,W,C] → PNG bytes → meta + binary""" + try: + import numpy as np + from PIL import Image as PILImage + if t.dim() == 4: + arr = t[0] + else: + arr = t + img_np = (arr.cpu().float().numpy() * 255).clip(0, 255).astype("uint8") + img = PILImage.fromarray(img_np) + return _encode_pil_image(img, file_url) + except Exception as e: + return {"__type__": "image", "error": str(e)}, None + + +def _encode_pil_image(img, file_url: str): + """PIL Image → PNG bytes → meta (inline base64 if small) + binary""" + buf = io.BytesIO() + img.save(buf, format="PNG") + png_bytes = buf.getvalue() + size_bytes = len(png_bytes) + + meta = { + "__type__": "image", + "url": file_url, + "width": img.width, + "height": img.height, + "format": "PNG", + "size_bytes": size_bytes, + } + + # 소형 이미지만 base64 인라인 포함 + if size_bytes < _INLINE_THRESHOLD: + meta["data"] = base64.b64encode(png_bytes).decode() + + return meta, (png_bytes, "image/png") + + +# ── 오디오 인코딩 ───────────────────────────────────────────── + +def _encode_audio(value: dict, file_url: str): + """{"waveform": Tensor [B,C,T], "sample_rate": int} → WAV bytes → meta + binary""" + try: + import torch + waveform = value["waveform"] + sample_rate = int(value["sample_rate"]) + + if isinstance(waveform, torch.Tensor): + # [B, C, T] → take first batch + if waveform.dim() == 3: + audio_np = waveform[0].cpu().float().numpy() # [C, T] + elif waveform.dim() == 2: + audio_np = waveform.cpu().float().numpy() # [C, T] + else: + audio_np = waveform.cpu().float().numpy().reshape(1, -1) + + wav_bytes = _numpy_to_wav(audio_np, sample_rate) + channels, num_samples = audio_np.shape if audio_np.ndim == 2 else (1, audio_np.shape[0]) + duration = num_samples / sample_rate + + meta = { + "__type__": "audio", + "url": file_url, + "format": "WAV", + "sample_rate": sample_rate, + "channels": channels, + "duration": round(duration, 3), + "num_samples": num_samples, + "size_bytes": len(wav_bytes), + } + return meta, (wav_bytes, "audio/wav") + except Exception as e: + logger.warning(f"[ComfyBridgeIO] Audio encoding failed: {e}") + + return {"__type__": "audio", "error": "encoding failed"}, None + + +def _numpy_to_wav(audio_np, sample_rate: int) -> bytes: + """numpy [C, T] float32 → WAV bytes (의존성 없음)""" + if audio_np.ndim == 1: + audio_np = audio_np.reshape(1, -1) + + channels = audio_np.shape[0] + # float32 → int16 + import numpy as np + audio_int16 = (audio_np * 32767).clip(-32768, 32767).astype(np.int16) + # [C, T] → interleave [T, C] → bytes + interleaved = audio_int16.T.flatten() + data = interleaved.tobytes() + + num_samples = audio_int16.shape[1] + byte_rate = sample_rate * channels * 2 + block_align = channels * 2 + data_size = len(data) + + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", 36 + data_size, b"WAVE", + b"fmt ", 16, 1, channels, sample_rate, + byte_rate, block_align, 16, + b"data", data_size, + ) + return header + data + + +# ── 영상 인코딩 ─────────────────────────────────────────────── + +def _encode_video_frames(frames, fps: float, file_url: str): + """ + Tensor [F,H,W,C] → MP4 bytes (imageio+ffmpeg) 또는 ZIP of PNGs (폴백) """ - # ── torch.Tensor 처리 ─────────────────────────────────────── try: import torch import numpy as np - from PIL import Image + from PIL import Image as PILImage - if isinstance(value, torch.Tensor): - t = value - if t.dim() == 4: - # IMAGE 포맷: [B, H, W, C] float32 - img_np = (t[0].cpu().float().numpy() * 255).clip(0, 255).astype("uint8") - img = Image.fromarray(img_np) - buf = io.BytesIO() - img.save(buf, format="PNG") - return { - "__type__": "image", - "data": base64.b64encode(buf.getvalue()).decode(), - "shape": list(t.shape), - } - elif t.dim() == 3: - # [H, W, C] 형태 - img_np = (t.cpu().float().numpy() * 255).clip(0, 255).astype("uint8") - img = Image.fromarray(img_np) - buf = io.BytesIO() - img.save(buf, format="PNG") - return { - "__type__": "image", - "data": base64.b64encode(buf.getvalue()).decode(), - "shape": list(t.shape), - } - else: - return {"__type__": "tensor", "shape": list(t.shape), "dtype": str(t.dtype)} + if isinstance(frames, torch.Tensor): + frames_np = (frames.cpu().float().numpy() * 255).clip(0, 255).astype(np.uint8) + else: + frames_np = frames - except ImportError: - pass + num_frames, h, w = frames_np.shape[:3] - # ── LATENT dict 처리 ──────────────────────────────────────── - # ComfyUI LATENT 포맷: {"samples": Tensor [B,C,H,W]} - if isinstance(value, dict) and "samples" in value: + # ── imageio + ffmpeg 시도 ────────────────────────── try: - import torch - samples = value["samples"] - if isinstance(samples, torch.Tensor): - shape = list(samples.shape) - # 잠재 공간 디코딩 없이 shape + 간략 통계만 반환 - return { - "__type__": "latent", - "shape": shape, - "batch_size": shape[0] if shape else 0, - "channels": shape[1] if len(shape) > 1 else 0, - "height_latent": shape[2] if len(shape) > 2 else 0, - "width_latent": shape[3] if len(shape) > 3 else 0, - # 픽셀 크기 추정 (SDXL 기본 VAE scale factor = 8) - "estimated_pixel_height": (shape[2] * 8) if len(shape) > 2 else 0, - "estimated_pixel_width": (shape[3] * 8) if len(shape) > 3 else 0, - } - except Exception: - pass - # samples 외 키도 함께 반환 - return {"__type__": "latent", "keys": list(value.keys())} + import imageio + import tempfile, os + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f: + tmp_path = f.name + writer = imageio.get_writer(tmp_path, fps=fps, codec="libx264", quality=8, macro_block_size=1) + for frame in frames_np: + writer.append_data(frame) + writer.close() + with open(tmp_path, "rb") as f: + mp4_bytes = f.read() + os.unlink(tmp_path) - # ── CONDITIONING 처리 ─────────────────────────────────────── - # ComfyUI CONDITIONING 포맷: list of [Tensor, dict] - if isinstance(value, list) and len(value) > 0 and isinstance(value[0], (list, tuple)): - try: - import torch - cond_info = [] - for item in value: - if isinstance(item, (list, tuple)) and len(item) >= 1: - tensor = item[0] - meta = item[1] if len(item) > 1 else {} - entry: dict = {} - if isinstance(tensor, torch.Tensor): - entry["shape"] = list(tensor.shape) - if isinstance(meta, dict): - # pooled_output shape 등 직렬화 가능한 메타만 추출 - safe_meta = {} - for k, v in meta.items(): - if isinstance(v, torch.Tensor): - safe_meta[k] = {"shape": list(v.shape)} - elif isinstance(v, (str, int, float, bool, type(None))): - safe_meta[k] = v - entry["meta"] = safe_meta - cond_info.append(entry) - return {"__type__": "conditioning", "conditions": cond_info} - except Exception: - pass - - # ── PIL Image 처리 ────────────────────────────────────────── - try: - from PIL import Image - if isinstance(value, Image.Image): - buf = io.BytesIO() - value.save(buf, format="PNG") - return { - "__type__": "image", - "data": base64.b64encode(buf.getvalue()).decode(), - "size": list(value.size), - "mode": value.mode, + meta = { + "__type__": "video", + "url": file_url, + "format": "MP4", + "fps": fps, + "num_frames": num_frames, + "width": w, + "height": h, + "duration": round(num_frames / fps, 3), + "size_bytes": len(mp4_bytes), } - except ImportError: - pass + return meta, (mp4_bytes, "video/mp4") + except Exception: + pass - # ── JSON 네이티브 타입 ───────────────────────────────────── - if isinstance(value, (str, int, float, bool, type(None))): - return value + # ── 폴백: ZIP of PNGs ───────────────────────────── + zip_buf = io.BytesIO() + with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf: + for i, frame in enumerate(frames_np): + img = PILImage.fromarray(frame) + frame_buf = io.BytesIO() + img.save(frame_buf, format="PNG") + zf.writestr(f"frame_{i:06d}.png", frame_buf.getvalue()) + # fps 메타 파일 + zf.writestr("meta.json", json.dumps({"fps": fps, "num_frames": num_frames})) + zip_bytes = zip_buf.getvalue() - # ── MODEL / CLIP / VAE 및 기타 ComfyUI 내부 객체 ──────────── - # 가중치 직렬화는 불가 → 클래스 정보 + 접근 가능한 메타만 반환 - type_name = type(value).__name__ - module_name = type(value).__module__ + meta = { + "__type__": "video", + "url": file_url, + "format": "ZIP_PNG_FRAMES", + "fps": fps, + "num_frames": num_frames, + "width": w, + "height": h, + "duration": round(num_frames / fps, 3), + "size_bytes": len(zip_bytes), + "note": "ffmpeg not available; frames exported as ZIP of PNGs", + } + return meta, (zip_bytes, "application/zip") - meta: dict = { - "__type__": "comfy_object", - "class": type_name, - "module": module_name, + except Exception as e: + logger.warning(f"[ComfyBridgeIO] Video encoding failed: {e}") + return {"__type__": "video", "error": str(e)}, None + + +# ── 파일 경로 처리 ──────────────────────────────────────────── + +def _encode_file(path: Path, file_url: str): + """파일 경로 → 바이트 읽기 → Content-Type 추정""" + _MIME = { + ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", + ".gif": "image/gif", ".webp": "image/webp", + ".wav": "audio/wav", ".mp3": "audio/mpeg", ".ogg": "audio/ogg", ".flac": "audio/flac", + ".mp4": "video/mp4", ".webm": "video/webm", ".avi": "video/avi", ".mov": "video/quicktime", + } + ext = path.suffix.lower() + content_type = _MIME.get(ext, "application/octet-stream") + raw = path.read_bytes() + + meta = { + "__type__": content_type.split("/")[0], # "image" / "audio" / "video" + "url": file_url, + "format": ext.lstrip(".").upper(), + "filename": path.name, + "size_bytes": len(raw), } - # ComfyUI ModelPatcher (MODEL 타입) - if hasattr(value, "model") and hasattr(value, "load_device"): - inner = value.model - meta["comfy_type"] = "MODEL" - if hasattr(inner, "__class__"): - meta["model_class"] = inner.__class__.__name__ - if hasattr(value, "load_device"): - try: - meta["device"] = str(value.load_device) - except Exception: - pass + # 소형 이미지이면 inline base64도 포함 + if content_type.startswith("image/") and len(raw) < _INLINE_THRESHOLD: + meta["data"] = base64.b64encode(raw).decode() - # CLIP + return meta, (raw, content_type) + + +# ── 기타 ComfyUI 타입 ───────────────────────────────────────── + +def _encode_latent(value: dict) -> dict: + try: + import torch + samples = value["samples"] + if isinstance(samples, torch.Tensor): + shape = list(samples.shape) + return { + "__type__": "latent", + "shape": shape, + "batch_size": shape[0] if shape else 0, + "channels": shape[1] if len(shape) > 1 else 0, + "height_latent": shape[2] if len(shape) > 2 else 0, + "width_latent": shape[3] if len(shape) > 3 else 0, + "estimated_pixel_height": shape[2] * 8 if len(shape) > 2 else 0, + "estimated_pixel_width": shape[3] * 8 if len(shape) > 3 else 0, + } + except Exception: + pass + return {"__type__": "latent", "keys": list(value.keys())} + + +def _encode_conditioning(value: list) -> dict: + try: + import torch + cond_info = [] + for item in value: + if isinstance(item, (list, tuple)) and len(item) >= 1: + tensor = item[0] + meta = item[1] if len(item) > 1 else {} + entry: dict = {} + if isinstance(tensor, torch.Tensor): + entry["shape"] = list(tensor.shape) + if isinstance(meta, dict): + safe = {} + for k, v in meta.items(): + if isinstance(v, torch.Tensor): + safe[k] = {"shape": list(v.shape)} + elif isinstance(v, (str, int, float, bool, type(None))): + safe[k] = v + entry["meta"] = safe + cond_info.append(entry) + return {"__type__": "conditioning", "conditions": cond_info} + except Exception: + pass + return {"__type__": "conditioning", "error": "parse failed"} + + +def _encode_comfy_object(value) -> dict: + meta: dict = { + "__type__": "comfy_object", + "class": type(value).__name__, + "module": type(value).__module__, + } + if hasattr(value, "model") and hasattr(value, "load_device"): + meta["comfy_type"] = "MODEL" + if hasattr(value.model, "__class__"): + meta["model_class"] = value.model.__class__.__name__ + try: + meta["device"] = str(value.load_device) + except Exception: + pass elif hasattr(value, "tokenize") and hasattr(value, "encode_from_tokens"): meta["comfy_type"] = "CLIP" - if hasattr(value, "patcher") and hasattr(value.patcher, "model"): - meta["model_class"] = value.patcher.model.__class__.__name__ - - # VAE elif hasattr(value, "decode") and hasattr(value, "encode"): meta["comfy_type"] = "VAE" - if hasattr(value, "first_stage_model"): - meta["model_class"] = value.first_stage_model.__class__.__name__ - return meta +# ── 라우트 등록 ─────────────────────────────────────────────── + def register_routes(): """PromptServer에 커스텀 라우트를 등록한다.""" try: from aiohttp import web from server import PromptServer + # ── GET /comfybridge/results/{prompt_id} ────────────── @PromptServer.instance.routes.get("/comfybridge/results/{prompt_id}") async def get_bridge_results(request: web.Request) -> web.Response: prompt_id = request.match_info["prompt_id"] @@ -203,15 +426,44 @@ def register_routes(): content_type="application/json", body=json.dumps({"error": f"No results for prompt_id '{prompt_id}'"}), ) - outputs = { - name: _serialize_value(val) - for name, val in _results_store[prompt_id].items() - } return web.Response( content_type="application/json", - body=json.dumps({"outputs": outputs}), + body=json.dumps({"outputs": _results_store[prompt_id]}), ) - logger.info("[ComfyBridgeIO] Registered route: GET /comfybridge/results/{prompt_id}") + # ── GET /comfybridge/file/{prompt_id}/{name} ────────── + @PromptServer.instance.routes.get("/comfybridge/file/{prompt_id}/{name}") + async def get_bridge_file(request: web.Request) -> web.Response: + prompt_id = request.match_info["prompt_id"] + name = request.match_info["name"] + + store = _binary_store.get(prompt_id, {}) + if name not in store: + return web.Response( + status=404, + content_type="application/json", + body=json.dumps({"error": f"No binary for '{name}' in prompt '{prompt_id}'"}), + ) + + raw_bytes, content_type = store[name] + + # Content-Disposition 파일명 결정 + _EXT = { + "image/png": ".png", "image/jpeg": ".jpg", "image/gif": ".gif", + "audio/wav": ".wav", "audio/mpeg": ".mp3", + "video/mp4": ".mp4", "video/webm": ".webm", + "application/zip": ".zip", + } + ext = _EXT.get(content_type, "") + disposition = f'attachment; filename="{name}{ext}"' + + return web.Response( + status=200, + headers={"Content-Disposition": disposition}, + content_type=content_type, + body=raw_bytes, + ) + + logger.info("[ComfyBridgeIO] Routes registered: /comfybridge/results, /comfybridge/file") except Exception as e: logger.warning(f"[ComfyBridgeIO] Failed to register routes: {e}")