feat: binary streaming for image/audio/video
- /comfybridge/file/{prompt_id}/{name}: binary streaming endpoint
- Images < 2MB: base64 inline + stream URL
- Images >= 2MB (4K+): stream URL only
- Audio: WAV encoding (no external deps)
- Video: MP4 via imageio+ffmpeg, ZIP of PNGs as fallback
- client: download() / save() / as_pil_async() / as_bytes_async()
- client: BridgeResult.info() / summary() / type_of()
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
7608cf5aa2
commit
26bd9affdf
+3
-1
@@ -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,
|
||||
|
||||
+174
-18
@@ -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)"
|
||||
|
||||
+399
-147
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user