55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
"""Pack source + merged LUT for in-node live preview (Comfy UI payload)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
|
|
import numpy as np
|
|
|
|
LIVE_PREVIEW_MAX_EDGE = 512
|
|
|
|
|
|
def _downscale_hwc(hwc: np.ndarray, max_edge: int) -> np.ndarray:
|
|
hwc = np.asarray(hwc, dtype=np.float32)
|
|
h, w, _ = hwc.shape
|
|
scale = min(1.0, max_edge / max(h, w))
|
|
if scale >= 1.0:
|
|
return hwc.copy()
|
|
nw = max(1, int(w * scale))
|
|
nh = max(1, int(h * scale))
|
|
try:
|
|
from PIL import Image
|
|
|
|
arr = (np.clip(hwc, 0, 1) * 255.0).astype(np.uint8)
|
|
pil = Image.fromarray(arr, "RGB").resize((nw, nh), Image.Resampling.BILINEAR)
|
|
return np.asarray(pil, dtype=np.float32) / 255.0
|
|
except ImportError:
|
|
# Nearest stride fallback
|
|
ys = (np.linspace(0, h - 1, nh)).astype(np.int32)
|
|
xs = (np.linspace(0, w - 1, nw)).astype(np.int32)
|
|
return hwc[np.ix_(ys, xs)]
|
|
|
|
|
|
def pack_live_cache(source_hwc: np.ndarray, merged_lut: np.ndarray) -> dict:
|
|
"""JSON-serializable bundle for the Comfy front-end (PNG + float32 LUT)."""
|
|
from PIL import Image
|
|
|
|
small = _downscale_hwc(source_hwc, LIVE_PREVIEW_MAX_EDGE)
|
|
sh, sw, _ = small.shape
|
|
merged = np.asarray(merged_lut, dtype=np.float32).reshape(-1)
|
|
|
|
buf = io.BytesIO()
|
|
Image.fromarray((np.clip(small, 0, 1) * 255.0).astype(np.uint8), "RGB").save(
|
|
buf, format="PNG", optimize=True
|
|
)
|
|
|
|
return {
|
|
"src_png": base64.b64encode(buf.getvalue()).decode("ascii"),
|
|
"w": int(sw),
|
|
"h": int(sh),
|
|
"lut": base64.b64encode(merged.tobytes()).decode("ascii"),
|
|
"lut_len": int(merged.size),
|
|
"lut_size": 25,
|
|
}
|