From 534bc19033031dca61b057af348fa02fb9b5dfa1 Mon Sep 17 00:00:00 2001 From: facok <128763816+facok@users.noreply.github.com> Date: Tue, 17 Mar 2026 01:43:26 +0800 Subject: [PATCH] Implement LCS plugin for training-free color control in FLUX Implements "The Latent Color Subspace" (arXiv:2603.12261v1, ICML 2026) as a ComfyUI plugin with V3 API + V2 backward compatibility. 6 nodes across 3 categories: - LCS/calibration: LCSCalibrate (PCA from VAE), LCSLoadData (cache load) - LCS/intervention: LCSColorIntervene (single color), LCSColorBatch (multi) - LCS/observe: LCSPreviewColors (latent color map), LCSStepObserver Core modules: patchify (16ch->64d patches), bicone HSL<->LCS mapping, alpha_t/beta_t tables (51 entries from Appendix F), PCA calibration, and timestep normalization with residual preservation. --- __init__.py | 37 ++++++ core/__init__.py | 9 ++ core/calibration.py | 169 +++++++++++++++++++++++++ core/color_space.py | 298 ++++++++++++++++++++++++++++++++++++++++++++ core/defaults.py | 63 ++++++++++ core/lcs_data.py | 22 ++++ core/patchify.py | 25 ++++ core/timestep.py | 75 +++++++++++ data/.gitkeep | 0 nodes/__init__.py | 23 ++++ nodes/calibrate.py | 126 +++++++++++++++++++ nodes/intervene.py | 265 +++++++++++++++++++++++++++++++++++++++ nodes/observe.py | 129 +++++++++++++++++++ requirements.txt | 2 + 14 files changed, 1243 insertions(+) create mode 100644 __init__.py create mode 100644 core/__init__.py create mode 100644 core/calibration.py create mode 100644 core/color_space.py create mode 100644 core/defaults.py create mode 100644 core/lcs_data.py create mode 100644 core/patchify.py create mode 100644 core/timestep.py create mode 100644 data/.gitkeep create mode 100644 nodes/__init__.py create mode 100644 nodes/calibrate.py create mode 100644 nodes/intervene.py create mode 100644 nodes/observe.py create mode 100644 requirements.txt diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..fb8829a --- /dev/null +++ b/__init__.py @@ -0,0 +1,37 @@ +"""ComfyUI-LCS: The Latent Color Subspace — training-free color control for FLUX. + +Paper: "The Latent Color Subspace" (arXiv:2603.12261v1, ICML 2026) +""" + +# V3 ComfyExtension entry point +from comfy_api.latest import ComfyExtension, io +from .nodes.calibrate import LCSCalibrate, LCSLoadData +from .nodes.intervene import LCSColorIntervene, LCSColorBatch +from .nodes.observe import LCSPreviewColors, LCSStepObserver + + +class LCSExtension(ComfyExtension): + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + LCSCalibrate, + LCSLoadData, + LCSColorIntervene, + LCSColorBatch, + LCSPreviewColors, + LCSStepObserver, + ] + + +async def comfy_entrypoint() -> LCSExtension: + return LCSExtension() + + +# V2 backward compatibility +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "LCSExtension", + "comfy_entrypoint", +] diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..7ab2fcd --- /dev/null +++ b/core/__init__.py @@ -0,0 +1,9 @@ +from .lcs_data import LCSData +from .patchify import patchify, unpatchify +from .timestep import sigma_to_paper_t, get_alpha_beta, normalize_to_t50, denormalize_from_t50 +from .color_space import decode_lcs_to_hsl, encode_hsl_to_lcs, hex_to_hsl, hsl_to_rgb + +# Lazy import for calibration (depends on comfy.utils) +def calibrate(*args, **kwargs): + from .calibration import calibrate as _calibrate + return _calibrate(*args, **kwargs) diff --git a/core/calibration.py b/core/calibration.py new file mode 100644 index 0000000..20ee704 --- /dev/null +++ b/core/calibration.py @@ -0,0 +1,169 @@ +"""PCA calibration from FLUX VAE: compute LCS basis, mean, and anchor positions.""" + +import math +import torch +import comfy.utils +from .patchify import patchify +from .lcs_data import LCSData + + +# 8 anchor colors: R, B, G, M, C, Y, Black, White +ANCHOR_COLORS_RGB = [ + (1.0, 0.0, 0.0), # Red + (0.0, 0.0, 1.0), # Blue + (0.0, 1.0, 0.0), # Green + (1.0, 0.0, 1.0), # Magenta + (0.0, 1.0, 1.0), # Cyan + (1.0, 1.0, 0.0), # Yellow + (0.0, 0.0, 0.0), # Black + (1.0, 1.0, 1.0), # White +] + + +def calibrate(vae, num_colors=512, image_size=512, batch_size=8): + """Compute LCS data (PCA basis, mean, anchors) from FLUX VAE. + + 1. Sample num_colors solid-color images uniformly from HSV + 2. VAE encode each → latent + 3. Patchify → average patches per image → vector in R^64 + 4. PCA on all vectors → basis B [64,3], mean μ [64] + 5. Encode 8 anchor colors → compute LCS coords + hue angles + + Returns: LCSData + """ + device = comfy.model_management.intermediate_device() + + # Step 1: Sample colors uniformly from HSV (full saturation, full value for diversity) + colors = [] + for i in range(num_colors): + # Uniform sampling in HSV + h = (i * 137.508) % 360.0 / 360.0 # Golden angle for uniform coverage + s = 0.3 + 0.7 * ((i * 73) % 100) / 100.0 # Vary saturation 0.3-1.0 + v = 0.3 + 0.7 * ((i * 47) % 100) / 100.0 # Vary value 0.3-1.0 + # HSV to RGB + r, g, b = _hsv_to_rgb(h, s, v) + colors.append((r, g, b)) + + # Step 2+3: Encode and average patches + vectors = [] + pbar = comfy.utils.ProgressBar(num_colors) + + for batch_start in range(0, num_colors, batch_size): + batch_end = min(batch_start + batch_size, num_colors) + batch_colors = colors[batch_start:batch_end] + actual_batch = len(batch_colors) + + # Create solid color images [B, H, W, 3] (BHWC format for ComfyUI VAE) + imgs = torch.zeros(actual_batch, image_size, image_size, 3, dtype=torch.float32, device="cpu") + for j, (r, g, b) in enumerate(batch_colors): + imgs[j, :, :, 0] = r + imgs[j, :, :, 1] = g + imgs[j, :, :, 2] = b + + # VAE encode → [B, 16, H/8, W/8] + latent = vae.encode(imgs[:, :, :, :3]) + + # Patchify → [B, L, 64] + patches, _, _ = patchify(latent) + + # Average across patches → [B, 64] + avg = patches.mean(dim=1).cpu() + for j in range(actual_batch): + vectors.append(avg[j]) + + pbar.update(actual_batch) + + # Stack all vectors: [N, 64] + X = torch.stack(vectors, dim=0).float() + + # Step 4: PCA + mean = X.mean(dim=0) # [64] + X_centered = X - mean + # SVD for PCA + U, S, Vh = torch.linalg.svd(X_centered, full_matrices=False) + # Top 3 components: B = V[:, :3] (columns are principal directions) + basis = Vh[:3].T # [64, 3] (Vh rows are right singular vectors) + + # Step 5: Encode 8 anchor colors → LCS coords + anchor_lcs_list = [] + for r, g, b in ANCHOR_COLORS_RGB: + img = torch.zeros(1, image_size, image_size, 3, dtype=torch.float32, device="cpu") + img[0, :, :, 0] = r + img[0, :, :, 1] = g + img[0, :, :, 2] = b + latent = vae.encode(img[:, :, :, :3]) + patches, _, _ = patchify(latent) + avg = patches.mean(dim=1).cpu().squeeze(0) # [64] + # Project to LCS + lcs_coord = (avg - mean) @ basis # [3] + anchor_lcs_list.append(lcs_coord) + + anchor_lcs = torch.stack(anchor_lcs_list, dim=0) # [8, 3] + + # Compute hue angles for 6 chromatic anchors + anchor_angles = _compute_anchor_angles(anchor_lcs, basis, mean) + + return LCSData( + basis=basis, + mean=mean, + anchor_lcs=anchor_lcs, + anchor_angles=anchor_angles, + ) + + +def _compute_anchor_angles(anchor_lcs, basis, mean): + """Compute hue angles of the 6 chromatic anchors in the chromatic plane. + + The chromatic plane is perpendicular to the achromatic axis (black→white). + Returns [6] tensor of angles in radians. + """ + black = anchor_lcs[6] # [3] + white = anchor_lcs[7] # [3] + chromatic = anchor_lcs[:6] # [6, 3] + + # Achromatic axis + a = white - black + a_unit = a / (a.norm() + 1e-10) + + # Build orthonormal basis in chromatic plane + arb = torch.zeros(3, dtype=a.dtype) + arb[0] = 1.0 + if a_unit[0].abs() > 0.9: + arb[0] = 0.0 + arb[1] = 1.0 + e1 = arb - (arb * a_unit).sum() * a_unit + e1 = e1 / (e1.norm() + 1e-10) + e2 = torch.linalg.cross(a_unit, e1) + + # Project each chromatic anchor onto the plane and compute angle + angles = [] + for i in range(6): + c = chromatic[i] + # Project onto achromatic axis + c_proj = black + ((c - black) * a).sum() / ((a * a).sum() + 1e-10) * a + # Chromatic residual + chroma = c - c_proj + x = (chroma * e1).sum() + y = (chroma * e2).sum() + angle = torch.atan2(y, x) % (2 * math.pi) + angles.append(angle) + + return torch.stack(angles) # [6] + + +def _hsv_to_rgb(h, s, v): + """Convert HSV to RGB (scalars in [0,1]).""" + if s < 1e-10: + return v, v, v + h6 = h * 6.0 + i = int(h6) % 6 + f = h6 - int(h6) + p = v * (1.0 - s) + q = v * (1.0 - s * f) + t = v * (1.0 - s * (1.0 - f)) + if i == 0: return v, t, p + if i == 1: return q, v, p + if i == 2: return p, v, t + if i == 3: return p, q, v + if i == 4: return t, p, v + return v, p, q diff --git a/core/color_space.py b/core/color_space.py new file mode 100644 index 0000000..7d5726f --- /dev/null +++ b/core/color_space.py @@ -0,0 +1,298 @@ +"""Bicone LCS ↔ HSL mapping using 8 anchor colors. + +Anchors are indexed as: [Red, Blue, Green, Magenta, Cyan, Yellow, Black, White] +Indices: 0=R, 1=B, 2=G, 3=M, 4=C, 5=Y, 6=Black, 7=White +""" + +import math +import torch + + +def hex_to_hsl(hex_str): + """Convert "#RRGGBB" to (h, s, l) where h∈[0,1], s∈[0,1], l∈[0,1].""" + hex_str = hex_str.lstrip("#") + r = int(hex_str[0:2], 16) / 255.0 + g = int(hex_str[2:4], 16) / 255.0 + b = int(hex_str[4:6], 16) / 255.0 + return rgb_to_hsl(r, g, b) + + +def rgb_to_hsl(r, g, b): + """Convert RGB [0,1] to HSL [0,1].""" + cmax = max(r, g, b) + cmin = min(r, g, b) + delta = cmax - cmin + l = (cmax + cmin) / 2.0 + + if delta < 1e-10: + return 0.0, 0.0, l + + s = delta / (1.0 - abs(2.0 * l - 1.0)) if abs(2.0 * l - 1.0) < 1.0 else 0.0 + + if cmax == r: + h = ((g - b) / delta) % 6.0 + elif cmax == g: + h = (b - r) / delta + 2.0 + else: + h = (r - g) / delta + 4.0 + h = h / 6.0 + if h < 0: + h += 1.0 + + return h, max(0.0, min(1.0, s)), max(0.0, min(1.0, l)) + + +def hsl_to_rgb(h, s, l): + """Convert HSL [0,1] to RGB [0,1]. Works with scalars or tensors.""" + if isinstance(h, torch.Tensor): + return _hsl_to_rgb_tensor(h, s, l) + + c = (1.0 - abs(2.0 * l - 1.0)) * s + x = c * (1.0 - abs((h * 6.0) % 2.0 - 1.0)) + m = l - c / 2.0 + + h6 = h * 6.0 + if h6 < 1: + r, g, b = c, x, 0 + elif h6 < 2: + r, g, b = x, c, 0 + elif h6 < 3: + r, g, b = 0, c, x + elif h6 < 4: + r, g, b = 0, x, c + elif h6 < 5: + r, g, b = x, 0, c + else: + r, g, b = c, 0, x + + return r + m, g + m, b + m + + +def _hsl_to_rgb_tensor(h, s, l): + """Vectorized HSL→RGB for tensors.""" + c = (1.0 - (2.0 * l - 1.0).abs()) * s + h6 = h * 6.0 + x = c * (1.0 - ((h6 % 2.0) - 1.0).abs()) + m = l - c / 2.0 + + r = torch.zeros_like(h) + g = torch.zeros_like(h) + b = torch.zeros_like(h) + + mask0 = h6 < 1 + mask1 = (h6 >= 1) & (h6 < 2) + mask2 = (h6 >= 2) & (h6 < 3) + mask3 = (h6 >= 3) & (h6 < 4) + mask4 = (h6 >= 4) & (h6 < 5) + mask5 = h6 >= 5 + + r[mask0] = c[mask0]; g[mask0] = x[mask0] + r[mask1] = x[mask1]; g[mask1] = c[mask1] + g[mask2] = c[mask2]; b[mask2] = x[mask2] + g[mask3] = x[mask3]; b[mask3] = c[mask3] + r[mask4] = x[mask4]; b[mask4] = c[mask4] + r[mask5] = c[mask5]; b[mask5] = x[mask5] + + return (r + m).clamp(0, 1), (g + m).clamp(0, 1), (b + m).clamp(0, 1) + + +def decode_lcs_to_hsl(c, anchor_lcs, anchor_angles): + """Decode LCS coordinates to HSL using bicone geometry. + + c: [..., 3] LCS coordinates (normalized to t=50) + anchor_lcs: [8, 3] anchor positions [R,B,G,M,C,Y,Black,White] + anchor_angles: [6] hue angles of chromatic anchors in radians + + Returns: (h, s, l) each [...] in [0,1] + """ + black = anchor_lcs[6] # [3] + white = anchor_lcs[7] # [3] + chromatic = anchor_lcs[:6] # [6, 3] + + # Achromatic axis + a = white - black # [3] + a_norm_sq = (a * a).sum() + 1e-10 + + # Lightness: project onto achromatic axis + diff = c - black # [..., 3] + l = (diff * a).sum(dim=-1) / a_norm_sq # [...] + l = l.clamp(0.0, 1.0) + + # Point on achromatic axis + c_L = black + l.unsqueeze(-1) * a # [..., 3] + + # Chromatic residual + chroma_vec = c - c_L # [..., 3] + chroma_dist = chroma_vec.norm(dim=-1) + 1e-10 # [...] + + # Compute hue angle in chromatic plane + # Build 2 orthonormal basis vectors in the plane perpendicular to a + a_unit = a / a.norm() + # Pick an arbitrary vector not parallel to a + arb = torch.zeros(3, device=a.device, dtype=a.dtype) + arb[0] = 1.0 + if (a_unit[0].abs() > 0.9): + arb[0] = 0.0 + arb[1] = 1.0 + e1 = arb - (arb * a_unit).sum() * a_unit + e1 = e1 / (e1.norm() + 1e-10) + e2 = torch.linalg.cross(a_unit, e1) + + # Project chromatic vector to 2D + x_coord = (chroma_vec * e1).sum(dim=-1) # [...] + y_coord = (chroma_vec * e2).sum(dim=-1) # [...] + angle = torch.atan2(y_coord, x_coord) # [...] radians + angle = angle % (2 * math.pi) + + # Map angle to hue [0,1] using sorted anchor angles + # anchor_angles are the angles of [R,B,G,M,C,Y] in the same coordinate system + # Standard HSL hue: R=0, Y=1/6, G=2/6, C=3/6, B=4/6, M=5/6 + # But anchors may not be in that order in angle-space, so we interpolate + sorted_angles, sort_idx = anchor_angles.sort() + # Standard hue for each anchor: R=0/6, B=4/6, G=2/6, M=5/6, C=3/6, Y=1/6 + anchor_hues = torch.tensor([0.0, 4.0/6.0, 2.0/6.0, 5.0/6.0, 3.0/6.0, 1.0/6.0], + device=c.device, dtype=c.dtype) + sorted_hues = anchor_hues[sort_idx] + + # Piecewise linear interpolation around the circle + h = _angle_to_hue(angle, sorted_angles, sorted_hues) + + # Saturation: distance to achromatic axis normalized by max distance + # Max distance at this hue and lightness + bicone_factor = 1.0 - (2.0 * l - 1.0).abs() # [...] + bicone_factor = bicone_factor.clamp(min=1e-10) + + # Find the hue point on the anchor polygon + c_H = _hue_to_polygon_point(h, chromatic, anchor_angles, a_unit, e1, e2) # [..., 3] + c_H_dist = (c_H - c_L).norm(dim=-1) + 1e-10 + s = chroma_dist / (c_H_dist * bicone_factor) + s = s.clamp(0.0, 1.0) + + return h, s, l + + +def encode_hsl_to_lcs(h, s, l, anchor_lcs, anchor_angles): + """Encode HSL to LCS coordinates using bicone geometry. + + h, s, l: [...] in [0,1] + anchor_lcs: [8, 3] + anchor_angles: [6] radians + + Returns: c [..., 3] LCS coordinates + """ + black = anchor_lcs[6] # [3] + white = anchor_lcs[7] # [3] + chromatic = anchor_lcs[:6] # [6, 3] + + a = white - black + a_unit = a / a.norm() + + # Build chromatic plane basis + arb = torch.zeros(3, device=a.device, dtype=a.dtype) + arb[0] = 1.0 + if (a_unit[0].abs() > 0.9): + arb[0] = 0.0 + arb[1] = 1.0 + e1 = arb - (arb * a_unit).sum() * a_unit + e1 = e1 / (e1.norm() + 1e-10) + e2 = torch.linalg.cross(a_unit, e1) + + # Lightness point on achromatic axis + c_L = black + l.unsqueeze(-1) * a # [..., 3] + + # Hue point on chromatic polygon + c_H = _hue_to_polygon_point(h, chromatic, anchor_angles, a_unit, e1, e2) # [..., 3] + + # Combine: c = c_L + s * (1 - |2l-1|) * (c_H - c_L) + bicone_factor = 1.0 - (2.0 * l - 1.0).abs() # [...] + c = c_L + (s * bicone_factor).unsqueeze(-1) * (c_H - c_L) + + return c + + +def _angle_to_hue(angle, sorted_angles, sorted_hues): + """Map an angle [...] to hue [0,1] via piecewise linear interpolation on anchor angles.""" + n = len(sorted_angles) + h = torch.zeros_like(angle) + + for i in range(n): + j = (i + 1) % n + a_start = sorted_angles[i] + a_end = sorted_angles[j] + h_start = sorted_hues[i] + h_end = sorted_hues[j] + + # Handle wraparound + if a_end < a_start: + a_end = a_end + 2 * math.pi + span = a_end - a_start + if span < 1e-10: + continue + + # Check which angles fall in this segment + angle_shifted = angle.clone() + if a_end > 2 * math.pi: + # Wraparound segment + mask = (angle >= a_start) | (angle < (a_end - 2 * math.pi)) + angle_shifted = torch.where(angle < a_start, angle + 2 * math.pi, angle) + else: + mask = (angle >= a_start) & (angle < a_end) + + frac = ((angle_shifted - a_start) / span).clamp(0, 1) + + # Interpolate hue (handling hue wraparound) + h_diff = h_end - h_start + if abs(h_diff) > 0.5: + if h_diff > 0: + h_diff -= 1.0 + else: + h_diff += 1.0 + interp = h_start + frac * h_diff + interp = interp % 1.0 + + h = torch.where(mask, interp, h) + + return h + + +def _hue_to_polygon_point(h, chromatic, anchor_angles, a_unit, e1, e2): + """Map hue values [...] to points on the chromatic anchor polygon in 3D LCS space. + + chromatic: [6, 3] anchor LCS positions + """ + # Standard hue for each anchor: R=0, B=4/6, G=2/6, M=5/6, C=3/6, Y=1/6 + anchor_hues = torch.tensor([0.0, 4.0/6.0, 2.0/6.0, 5.0/6.0, 3.0/6.0, 1.0/6.0], + device=chromatic.device, dtype=chromatic.dtype) + + # Sort anchors by hue for polygon interpolation + sorted_hues, sort_idx = anchor_hues.sort() + sorted_chromatic = chromatic[sort_idx] # [6, 3] + + # For each input hue, find which polygon segment it falls in and interpolate + result = torch.zeros(h.shape + (3,), device=chromatic.device, dtype=chromatic.dtype) + + for i in range(6): + j = (i + 1) % 6 + h_start = sorted_hues[i] + h_end = sorted_hues[j] + + if j == 0: + h_end = h_end + 1.0 # wraparound + + span = h_end - h_start + if span < 1e-10: + continue + + if j == 0: + mask = (h >= h_start) | (h < sorted_hues[0]) + h_shifted = torch.where(h < h_start, h + 1.0, h) + else: + mask = (h >= h_start) & (h < h_end) + h_shifted = h + + frac = ((h_shifted - h_start) / span).clamp(0, 1) + + interp = sorted_chromatic[i] + frac.unsqueeze(-1) * (sorted_chromatic[j] - sorted_chromatic[i]) + result = torch.where(mask.unsqueeze(-1), interp, result) + + return result diff --git a/core/defaults.py b/core/defaults.py new file mode 100644 index 0000000..9c23c58 --- /dev/null +++ b/core/defaults.py @@ -0,0 +1,63 @@ +"""Hardcoded alpha_t and beta_t tables from paper Appendix F (51 entries, t=0..50).""" + +import torch + +# Shift alpha_t: 3D vectors for each timestep t=0..50 +ALPHA_T = [ + [2.3413, -2.3586, 0.4266], [2.3574, -2.3833, 0.4644], [2.3638, -2.3904, 0.4883], + [2.3734, -2.3951, 0.5122], [2.3831, -2.3993, 0.5384], [2.3925, -2.4026, 0.5647], + [2.4023, -2.4047, 0.5919], [2.4124, -2.4060, 0.6198], [2.4226, -2.4064, 0.6484], + [2.4330, -2.4060, 0.6772], [2.4437, -2.4051, 0.7065], [2.4546, -2.4035, 0.7367], + [2.4659, -2.4011, 0.7668], [2.4775, -2.3981, 0.7974], [2.4897, -2.4009, 0.8312], + [2.5021, -2.4036, 0.8656], [2.5148, -2.4065, 0.9008], [2.5277, -2.4093, 0.9364], + [2.5408, -2.4123, 0.9727], [2.5542, -2.4154, 1.0099], [2.5680, -2.4186, 1.0481], + [2.5820, -2.4218, 1.0868], [2.5963, -2.4252, 1.1263], [2.6110, -2.4288, 1.1672], + [2.6261, -2.4324, 1.2090], [2.6416, -2.4363, 1.2520], [2.6575, -2.4403, 1.2957], + [2.6738, -2.4444, 1.3406], [2.6904, -2.4485, 1.3865], [2.7074, -2.4529, 1.4336], + [2.7250, -2.4574, 1.4818], [2.7432, -2.4621, 1.5314], [2.7618, -2.4669, 1.5823], + [2.7810, -2.4720, 1.6344], [2.8006, -2.4771, 1.6878], [2.8209, -2.4826, 1.7430], + [2.8418, -2.4883, 1.7995], [2.8631, -2.4944, 1.8578], [2.8853, -2.5005, 1.9179], + [2.9080, -2.5066, 1.9793], [2.9313, -2.5132, 2.0426], [2.9555, -2.5199, 2.1082], + [2.9804, -2.5268, 2.1756], [3.0060, -2.5338, 2.2450], [3.0328, -2.5411, 2.3172], + [3.0603, -2.5486, 2.3914], [3.0889, -2.5561, 2.4682], [3.1189, -2.5640, 2.5482], + [3.1497, -2.5725, 2.6302], [3.1824, -2.5796, 2.7175], [3.2152, -2.5889, 2.8050], +] + +# Scale beta_t: 3D vectors for each timestep t=0..50 +BETA_T = [ + [0.0163, 0.0172, 0.0295], [0.0905, 0.0716, 0.0999], [0.1345, 0.1123, 0.1544], + [0.1826, 0.1491, 0.2065], [0.2360, 0.1899, 0.2630], [0.2904, 0.2316, 0.3202], + [0.3471, 0.2749, 0.3793], [0.4050, 0.3191, 0.4394], [0.4640, 0.3641, 0.5003], + [0.5231, 0.4091, 0.5611], [0.5834, 0.4547, 0.6228], [0.6456, 0.5016, 0.6861], + [0.7077, 0.5481, 0.7488], [0.7713, 0.5958, 0.8127], [0.8410, 0.6496, 0.8866], + [0.9119, 0.7044, 0.9616], [0.9845, 0.7605, 1.0386], [1.0578, 0.8172, 1.1163], + [1.1325, 0.8750, 1.1957], [1.2094, 0.9344, 1.2771], [1.2880, 0.9953, 1.3606], + [1.3680, 1.0571, 1.4453], [1.4498, 1.1205, 1.5321], [1.5341, 1.1858, 1.6216], + [1.6206, 1.2526, 1.7131], [1.7094, 1.3214, 1.8072], [1.7998, 1.3913, 1.9030], + [1.8927, 1.4633, 2.0014], [1.9879, 1.5370, 2.1022], [2.0854, 1.6126, 2.2056], + [2.1853, 1.6900, 2.3114], [2.2881, 1.7696, 2.4202], [2.3939, 1.8515, 2.5321], + [2.5021, 1.9354, 2.6467], [2.6133, 2.0215, 2.7642], [2.7280, 2.1106, 2.8857], + [2.8455, 2.2017, 3.0101], [2.9668, 2.2957, 3.1386], [3.0921, 2.3929, 3.2712], + [3.2204, 2.4922, 3.4067], [3.3523, 2.5946, 3.5464], [3.4888, 2.7006, 3.6911], + [3.6292, 2.8097, 3.8398], [3.7741, 2.9222, 3.9931], [3.9247, 3.0394, 4.1527], + [4.0793, 3.1597, 4.3168], [4.2393, 3.2843, 4.4866], [4.4053, 3.4142, 4.6636], + [4.5760, 3.5480, 4.8461], [4.7541, 3.6886, 5.0383], [4.9407, 3.8364, 5.2390], +] + +# Pre-convert to tensors (lazily cached on first access) +_alpha_tensor = None +_beta_tensor = None + + +def get_alpha_table(): + global _alpha_tensor + if _alpha_tensor is None: + _alpha_tensor = torch.tensor(ALPHA_T, dtype=torch.float32) # [51, 3] + return _alpha_tensor + + +def get_beta_table(): + global _beta_tensor + if _beta_tensor is None: + _beta_tensor = torch.tensor(BETA_T, dtype=torch.float32) # [51, 3] + return _beta_tensor diff --git a/core/lcs_data.py b/core/lcs_data.py new file mode 100644 index 0000000..b1874d8 --- /dev/null +++ b/core/lcs_data.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +import torch + + +@dataclass +class LCSData: + basis: torch.Tensor # [64, 3] PCA basis B (orthonormal columns) + mean: torch.Tensor # [64] PCA mean mu + anchor_lcs: torch.Tensor # [8, 3] LCS coords of 8 anchor colors [R,B,G,M,C,Y,Black,White] + anchor_angles: torch.Tensor # [6] hue angles (radians) of the 6 chromatic anchors + + def to(self, device, dtype=None): + """Move all tensors to device/dtype.""" + kw = {"device": device} + if dtype is not None: + kw["dtype"] = dtype + return LCSData( + basis=self.basis.to(**kw), + mean=self.mean.to(**kw), + anchor_lcs=self.anchor_lcs.to(**kw), + anchor_angles=self.anchor_angles.to(**kw), + ) diff --git a/core/patchify.py b/core/patchify.py new file mode 100644 index 0000000..885b4d5 --- /dev/null +++ b/core/patchify.py @@ -0,0 +1,25 @@ +"""Patchify/unpatchify for FLUX latent tensors (patch_size=2, 16 channels → 64-dim patches).""" + +from einops import rearrange + + +def patchify(x): + """Convert latent [B, 16, H, W] → patch sequence [B, L, 64]. + + L = (H/2) * (W/2), d = 16 * 2 * 2 = 64. + Returns (patches, h_len, w_len) where h_len=H/2, w_len=W/2. + """ + B, C, H, W = x.shape + h_len = H // 2 + w_len = W // 2 + patches = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + return patches, h_len, w_len + + +def unpatchify(patches, h_len, w_len): + """Convert patch sequence [B, L, 64] → latent [B, 16, H, W]. + + h_len, w_len from patchify output. + """ + return rearrange(patches, "b (h w) (c ph pw) -> b c (h ph) (w pw)", + h=h_len, w=w_len, c=16, ph=2, pw=2) diff --git a/core/timestep.py b/core/timestep.py new file mode 100644 index 0000000..9694800 --- /dev/null +++ b/core/timestep.py @@ -0,0 +1,75 @@ +"""Sigma ↔ paper timestep conversion and α_t/β_t interpolation.""" + +import torch +from .defaults import get_alpha_table, get_beta_table + + +def sigma_to_paper_t(sigma): + """Convert FLUX sigma ∈ [0,1] to paper timestep t ∈ [0,50]. + + sigma=1 → noise → t=0, sigma=0 → clean → t=50. + """ + if isinstance(sigma, torch.Tensor): + return 50.0 * (1.0 - sigma.clamp(0.0, 1.0)) + return 50.0 * (1.0 - max(0.0, min(1.0, sigma))) + + +def get_alpha_beta(sigma, device=None): + """Get interpolated α_t and β_t [3] vectors for a given sigma. + + Returns (alpha_t, beta_t) as tensors on the specified device. + """ + t = sigma_to_paper_t(sigma) + if isinstance(t, torch.Tensor): + t = t.item() + + alpha_table = get_alpha_table() # [51, 3] + beta_table = get_beta_table() # [51, 3] + + t = max(0.0, min(50.0, t)) + t_low = int(t) + t_high = min(t_low + 1, 50) + frac = t - t_low + + alpha = (1.0 - frac) * alpha_table[t_low] + frac * alpha_table[t_high] + beta = (1.0 - frac) * beta_table[t_low] + frac * beta_table[t_high] + + if device is not None: + alpha = alpha.to(device) + beta = beta.to(device) + return alpha, beta + + +def get_alpha_beta_t50(device=None): + """Get α_50 and β_50 (reference timestep t=50, clean image).""" + alpha_table = get_alpha_table() + beta_table = get_beta_table() + alpha_50 = alpha_table[50] + beta_50 = beta_table[50] + if device is not None: + alpha_50 = alpha_50.to(device) + beta_50 = beta_50.to(device) + return alpha_50, beta_50 + + +def normalize_to_t50(c, alpha_t, beta_t, alpha_50, beta_50): + """Normalize LCS coords from timestep t to reference t=50. + + ĉ = (c - α_t) / β_t * β_50 + α_50 + c: [..., 3], alpha_t/beta_t/alpha_50/beta_50: [3] + """ + beta_t_safe = beta_t.clone() + beta_t_safe = torch.where(beta_t_safe.abs() < 1e-6, + torch.full_like(beta_t_safe, 1e-6), beta_t_safe) + return (c - alpha_t) / beta_t_safe * beta_50 + alpha_50 + + +def denormalize_from_t50(c_hat, alpha_t, beta_t, alpha_50, beta_50): + """Denormalize LCS coords from reference t=50 back to timestep t. + + c = (ĉ - α_50) / β_50 * β_t + α_t + """ + beta_50_safe = beta_50.clone() + beta_50_safe = torch.where(beta_50_safe.abs() < 1e-6, + torch.full_like(beta_50_safe, 1e-6), beta_50_safe) + return (c_hat - alpha_50) / beta_50_safe * beta_t + alpha_t diff --git a/data/.gitkeep b/data/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..1bd94a0 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,23 @@ +"""V2 backward compatibility exports for all LCS nodes.""" + +from .calibrate import LCSCalibrate, LCSLoadData +from .intervene import LCSColorIntervene, LCSColorBatch +from .observe import LCSPreviewColors, LCSStepObserver + +NODE_CLASS_MAPPINGS = { + "LCSCalibrate": LCSCalibrate, + "LCSLoadData": LCSLoadData, + "LCSColorIntervene": LCSColorIntervene, + "LCSColorBatch": LCSColorBatch, + "LCSPreviewColors": LCSPreviewColors, + "LCSStepObserver": LCSStepObserver, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LCSCalibrate": "LCS Calibrate", + "LCSLoadData": "LCS Load Data", + "LCSColorIntervene": "LCS Color Intervene", + "LCSColorBatch": "LCS Color Batch", + "LCSPreviewColors": "LCS Preview Colors", + "LCSStepObserver": "LCS Step Observer", +} diff --git a/nodes/calibrate.py b/nodes/calibrate.py new file mode 100644 index 0000000..c5894cc --- /dev/null +++ b/nodes/calibrate.py @@ -0,0 +1,126 @@ +"""Calibration nodes: LCSCalibrate and LCSLoadData.""" + +import os +import torch +from comfy_api.latest import io +from safetensors.torch import save_file, load_file + +from ..core.calibration import calibrate +from ..core.lcs_data import LCSData + +LCS_DATA = io.Custom("LCS_DATA") + +# Directory for cached calibration files +DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data") + + +class LCSCalibrate(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LCSCalibrate", + display_name="LCS Calibrate", + category="LCS/calibration", + description="Compute LCS basis and anchors from FLUX VAE via PCA on solid-color images", + inputs=[ + io.Vae.Input("vae", tooltip="FLUX VAE model"), + io.Int.Input("num_colors", default=512, min=64, max=2048, step=64, + tooltip="Number of solid-color samples for PCA"), + io.Int.Input("image_size", default=512, min=256, max=1024, step=128, + tooltip="Size of solid-color calibration images"), + ], + outputs=[ + LCS_DATA.Output(display_name="lcs_data"), + ], + ) + + @classmethod + def execute(cls, vae, num_colors, image_size) -> io.NodeOutput: + lcs_data = calibrate(vae, num_colors=num_colors, image_size=image_size) + + # Auto-save to data/ directory + os.makedirs(DATA_DIR, exist_ok=True) + save_path = os.path.join(DATA_DIR, "lcs_calibration.safetensors") + save_file({ + "basis": lcs_data.basis, + "mean": lcs_data.mean, + "anchor_lcs": lcs_data.anchor_lcs, + "anchor_angles": lcs_data.anchor_angles, + }, save_path) + + return io.NodeOutput(lcs_data) + + +class LCSLoadData(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + # Scan data/ dir for available calibration files + files = ["auto"] + if os.path.isdir(DATA_DIR): + for f in sorted(os.listdir(DATA_DIR)): + if f.endswith(".safetensors"): + files.append(f) + + return io.Schema( + node_id="LCSLoadData", + display_name="LCS Load Data", + category="LCS/calibration", + description="Load cached LCS calibration data, or auto-calibrate if VAE is provided", + inputs=[ + io.Combo.Input("calibration_file", options=files, default="auto", + tooltip="Select calibration file or 'auto' to use latest/auto-calibrate"), + io.Vae.Input("vae", optional=True, + tooltip="VAE for auto-calibration if no cached file exists"), + ], + outputs=[ + LCS_DATA.Output(display_name="lcs_data"), + ], + ) + + @classmethod + def execute(cls, calibration_file, vae=None) -> io.NodeOutput: + loaded = False + + if calibration_file != "auto": + path = os.path.join(DATA_DIR, calibration_file) + if os.path.exists(path): + data = load_file(path) + lcs_data = LCSData( + basis=data["basis"], + mean=data["mean"], + anchor_lcs=data["anchor_lcs"], + anchor_angles=data["anchor_angles"], + ) + loaded = True + + if not loaded: + # Try auto-loading the default calibration file + default_path = os.path.join(DATA_DIR, "lcs_calibration.safetensors") + if os.path.exists(default_path): + data = load_file(default_path) + lcs_data = LCSData( + basis=data["basis"], + mean=data["mean"], + anchor_lcs=data["anchor_lcs"], + anchor_angles=data["anchor_angles"], + ) + loaded = True + + if not loaded: + if vae is None: + raise ValueError( + "No calibration file found and no VAE provided. " + "Either run LCSCalibrate first or provide a VAE for auto-calibration." + ) + # Auto-calibrate + lcs_data = calibrate(vae, num_colors=512, image_size=512) + os.makedirs(DATA_DIR, exist_ok=True) + save_path = os.path.join(DATA_DIR, "lcs_calibration.safetensors") + save_file({ + "basis": lcs_data.basis, + "mean": lcs_data.mean, + "anchor_lcs": lcs_data.anchor_lcs, + "anchor_angles": lcs_data.anchor_angles, + }, save_path) + + return io.NodeOutput(lcs_data) diff --git a/nodes/intervene.py b/nodes/intervene.py new file mode 100644 index 0000000..8bdc24d --- /dev/null +++ b/nodes/intervene.py @@ -0,0 +1,265 @@ +"""Intervention nodes: LCSColorIntervene and LCSColorBatch.""" + +import torch +import torch.nn.functional as F +from comfy_api.latest import io + +from ..core.lcs_data import LCSData +from ..core.patchify import patchify, unpatchify +from ..core.timestep import get_alpha_beta, get_alpha_beta_t50, normalize_to_t50, denormalize_from_t50 +from ..core.color_space import hex_to_hsl, encode_hsl_to_lcs, decode_lcs_to_hsl + +LCS_DATA = io.Custom("LCS_DATA") + +# FLUX VAE constants +SCALE_FACTOR = 0.3611 +SHIFT_FACTOR = 0.1159 + + +def _build_post_cfg_fn(lcs_data, target_colors_hsl, strength, mode, start_step, end_step, mask): + """Build the post_cfg_function closure for color intervention. + + target_colors_hsl: list of (h, s, l) tuples, one per batch item (or one for all). + """ + def post_cfg_fn(args): + denoised = args["denoised"] # [B, 16, H, W] in process_in space + sigma = args["sigma"] + + # Determine current step + sigma_val = float(sigma.flatten()[0]) + + # Get step index from sample_sigmas + sigmas = args["model_options"]["transformer_options"]["sample_sigmas"] + step_index = None + matched = (sigmas == sigma_val).nonzero() + if len(matched) > 0: + step_index = matched[0].item() + else: + for i in range(len(sigmas) - 1): + if (sigmas[i] - sigma_val) * (sigmas[i + 1] - sigma_val) <= 0: + step_index = i + break + if step_index is None: + return denoised + + # Check if in intervention range + if step_index < start_step or step_index > end_step: + return denoised + + device = denoised.device + dtype = denoised.dtype + + # Move LCS data to device + ld = lcs_data.to(device, dtype) + B_mat = ld.basis # [64, 3] + mu = ld.mean # [64] + anchor_lcs = ld.anchor_lcs # [8, 3] + anchor_angles = ld.anchor_angles # [6] + + # Convert from process_in to raw VAE space + raw = denoised / SCALE_FACTOR + SHIFT_FACTOR # [B, 16, H, W] + + # Patchify + patches, h_len, w_len = patchify(raw) # [B, L, 64] + + # Project to LCS + projection = (patches - mu) @ B_mat # [B, L, 3] + + # Compute residual (61D orthogonal complement) + reconstruction = projection @ B_mat.T + mu # [B, L, 64] + residual = patches - reconstruction # [B, L, 64] + + # Get timestep statistics + alpha_t, beta_t = get_alpha_beta(sigma_val, device=device) + alpha_t, beta_t = alpha_t.to(dtype), beta_t.to(dtype) + alpha_50, beta_50 = get_alpha_beta_t50(device=device) + alpha_50, beta_50 = alpha_50.to(dtype), beta_50.to(dtype) + + # Normalize to t=50 + c_norm = normalize_to_t50(projection, alpha_t, beta_t, alpha_50, beta_50) # [B, L, 3] + + # Apply intervention per batch item + B_size = c_norm.shape[0] + new_c_norm = c_norm.clone() + + for b in range(B_size): + color_idx = b if b < len(target_colors_hsl) else 0 + t_h, t_s, t_l = target_colors_hsl[color_idx] + + # Encode target color to LCS at t=50 + t_h_t = torch.tensor(t_h, device=device, dtype=dtype) + t_s_t = torch.tensor(t_s, device=device, dtype=dtype) + t_l_t = torch.tensor(t_l, device=device, dtype=dtype) + target_lcs = encode_hsl_to_lcs(t_h_t, t_s_t, t_l_t, anchor_lcs, anchor_angles) # [3] + + c_b = c_norm[b] # [L, 3] + + if mode == "type_i": + # Type I: direct LCS translation + shift = target_lcs - c_b.mean(dim=0) + new_c_norm[b] = c_b + strength * shift + + elif mode == "type_ii": + # Type II: decode → shift in HSL → re-encode + h_cur, s_cur, l_cur = decode_lcs_to_hsl(c_b, anchor_lcs, anchor_angles) + # Shift towards target HSL + h_new = t_h_t.expand_as(h_cur) + s_new = t_s_t.expand_as(s_cur) + l_new = t_l_t.expand_as(l_cur) + # Interpolate in HSL + h_interp = _hue_lerp(h_cur, h_new, strength) + s_interp = s_cur + strength * (s_new - s_cur) + l_interp = l_cur + strength * (l_new - l_cur) + new_c_norm[b] = encode_hsl_to_lcs(h_interp, s_interp.clamp(0, 1), + l_interp.clamp(0, 1), + anchor_lcs, anchor_angles) + + else: # interpolated (default) + # gamma_t = sigma (high sigma → Type I, low sigma → Type II) + gamma = sigma_val + + # Type I + shift = target_lcs - c_b.mean(dim=0) + c_type_i = c_b + strength * shift + + # Type II + h_cur, s_cur, l_cur = decode_lcs_to_hsl(c_b, anchor_lcs, anchor_angles) + h_new = t_h_t.expand_as(h_cur) + s_new = t_s_t.expand_as(s_cur) + l_new = t_l_t.expand_as(l_cur) + h_interp = _hue_lerp(h_cur, h_new, strength) + s_interp = s_cur + strength * (s_new - s_cur) + l_interp = l_cur + strength * (l_new - l_cur) + c_type_ii = encode_hsl_to_lcs(h_interp, s_interp.clamp(0, 1), + l_interp.clamp(0, 1), + anchor_lcs, anchor_angles) + + # Interpolate: gamma * Type_I + (1-gamma) * Type_II + new_c_norm[b] = gamma * c_type_i + (1.0 - gamma) * c_type_ii + + # Apply mask if provided + if mask is not None: + mask_dev = mask.to(device=device, dtype=dtype) + if mask_dev.ndim == 3: + mask_dev = mask_dev[:1] # Use first mask for all batch items + # Downsample mask to patch grid: → [B, 1, h_len, w_len] + if mask_dev.ndim == 2: + mask_4d = mask_dev.unsqueeze(0).unsqueeze(0) # [1, 1, H, W] + elif mask_dev.ndim == 3: + mask_4d = mask_dev.unsqueeze(1) # [B, 1, H, W] + else: + mask_4d = mask_dev # assume already 4D + mask_resized = F.interpolate( + mask_4d, size=(h_len, w_len), mode="bilinear", align_corners=False + ) + # Flatten to [1, L, 1] + mask_flat = mask_resized.reshape(1, -1, 1) + if mask_flat.shape[1] != new_c_norm.shape[1]: + mask_flat = mask_flat[:, :new_c_norm.shape[1], :] + # Blend: masked areas get intervention, unmasked keep original + new_c_norm = c_norm + mask_flat * (new_c_norm - c_norm) + + # Denormalize back to timestep t + new_projection = denormalize_from_t50(new_c_norm, alpha_t, beta_t, alpha_50, beta_50) + + # Reconstruct patches + patches_new = new_projection @ B_mat.T + mu + residual # [B, L, 64] + + # Unpatchify + raw_new = unpatchify(patches_new, h_len, w_len) # [B, 16, H, W] + + # Convert back to process_in space + modified = (raw_new - SHIFT_FACTOR) * SCALE_FACTOR + + return modified.to(dtype) + + return post_cfg_fn + + +def _hue_lerp(h1, h2, t): + """Lerp hues on the circle [0,1], taking the shortest path.""" + diff = h2 - h1 + # Wrap to [-0.5, 0.5] + diff = diff - (diff > 0.5).float() + (diff < -0.5).float() + result = h1 + t * diff + return result % 1.0 + + +class LCSColorIntervene(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LCSColorIntervene", + display_name="LCS Color Intervene", + category="LCS/intervention", + description="Steer colors during FLUX generation via Latent Color Subspace", + inputs=[ + io.Model.Input("model"), + LCS_DATA.Input("lcs_data", tooltip="Calibration data from LCSCalibrate or LCSLoadData"), + io.Color.Input("color", default="#FF0000", tooltip="Target color"), + io.Float.Input("strength", default=1.0, min=0.0, max=2.0, step=0.05, + tooltip="Intervention strength (1.0 = full, 0.0 = none)"), + io.Combo.Input("mode", options=["interpolated", "type_i", "type_ii"], + default="interpolated", + tooltip="Interpolated blends Type I (LCS shift) and Type II (HSL shift)"), + io.Int.Input("start_step", default=8, min=0, max=50, + tooltip="First step to apply intervention (paper optimal: 8)"), + io.Int.Input("end_step", default=10, min=0, max=50, + tooltip="Last step to apply intervention (paper optimal: 10)"), + io.Mask.Input("mask", optional=True, + tooltip="Optional mask for localized color control"), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, lcs_data, color, strength, mode, start_step, end_step, + mask=None) -> io.NodeOutput: + m = model.clone() + h, s, l = hex_to_hsl(color) + hook = _build_post_cfg_fn(lcs_data, [(h, s, l)], strength, mode, start_step, end_step, mask) + m.set_model_sampler_post_cfg_function(hook) + return io.NodeOutput(m) + + +class LCSColorBatch(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LCSColorBatch", + display_name="LCS Color Batch", + category="LCS/intervention", + description="Apply different target colors to each batch item for multi-color generation", + inputs=[ + io.Model.Input("model"), + LCS_DATA.Input("lcs_data", tooltip="Calibration data from LCSCalibrate or LCSLoadData"), + io.String.Input("colors", default="#FF0000,#00FF00,#0000FF", + tooltip="Comma-separated hex colors, one per batch item"), + io.Float.Input("strength", default=1.0, min=0.0, max=2.0, step=0.05), + io.Combo.Input("mode", options=["interpolated", "type_i", "type_ii"], + default="interpolated"), + io.Int.Input("start_step", default=8, min=0, max=50), + io.Int.Input("end_step", default=10, min=0, max=50), + io.Mask.Input("mask", optional=True), + ], + outputs=[ + io.Model.Output(display_name="model"), + io.Int.Output(display_name="batch_size"), + ], + ) + + @classmethod + def execute(cls, model, lcs_data, colors, strength, mode, start_step, end_step, + mask=None) -> io.NodeOutput: + m = model.clone() + + # Parse comma-separated hex colors + color_list = [c.strip() for c in colors.split(",") if c.strip()] + target_hsl = [hex_to_hsl(c) for c in color_list] + batch_size = len(target_hsl) + + hook = _build_post_cfg_fn(lcs_data, target_hsl, strength, mode, start_step, end_step, mask) + m.set_model_sampler_post_cfg_function(hook) + return io.NodeOutput(m, batch_size) diff --git a/nodes/observe.py b/nodes/observe.py new file mode 100644 index 0000000..c4eedcc --- /dev/null +++ b/nodes/observe.py @@ -0,0 +1,129 @@ +"""Observation nodes: LCSPreviewColors and LCSStepObserver.""" + +import os +import torch +import numpy as np +import torch.nn.functional as F +from PIL import Image as PILImage +from comfy_api.latest import io +import folder_paths + +from ..core.lcs_data import LCSData +from ..core.patchify import patchify +from ..core.timestep import get_alpha_beta, get_alpha_beta_t50, normalize_to_t50 +from ..core.color_space import decode_lcs_to_hsl, hsl_to_rgb + +LCS_DATA = io.Custom("LCS_DATA") + +# FLUX VAE constants +SCALE_FACTOR = 0.3611 +SHIFT_FACTOR = 0.1159 + + +def _latent_to_color_preview(samples, lcs_data, sigma, upscale=8): + """Convert latent tensor to LCS color preview image. + + samples: [B, 16, H, W] in process_in space + Returns: [B, H_up, W_up, 3] float32 in [0,1] + """ + device = samples.device + dtype = samples.dtype + ld = lcs_data.to(device, dtype) + + raw = samples / SCALE_FACTOR + SHIFT_FACTOR + patches, h_len, w_len = patchify(raw) + projection = (patches - ld.mean) @ ld.basis + + alpha_t, beta_t = get_alpha_beta(sigma, device=device) + alpha_t, beta_t = alpha_t.to(dtype), beta_t.to(dtype) + alpha_50, beta_50 = get_alpha_beta_t50(device=device) + alpha_50, beta_50 = alpha_50.to(dtype), beta_50.to(dtype) + c_norm = normalize_to_t50(projection, alpha_t, beta_t, alpha_50, beta_50) + + B = c_norm.shape[0] + images = [] + for b in range(B): + c_b = c_norm[b] + h_vals, s_vals, l_vals = decode_lcs_to_hsl(c_b, ld.anchor_lcs, ld.anchor_angles) + r, g, b_ch = hsl_to_rgb(h_vals, s_vals, l_vals) + rgb = torch.stack([r, g, b_ch], dim=-1).reshape(h_len, w_len, 3) + if upscale > 1: + rgb = F.interpolate( + rgb.permute(2, 0, 1).unsqueeze(0), + scale_factor=upscale, mode="nearest" + ).squeeze(0).permute(1, 2, 0) + images.append(rgb) + + return torch.stack(images, dim=0).clamp(0, 1).cpu().float() + + +class LCSPreviewColors(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LCSPreviewColors", + display_name="LCS Preview Colors", + category="LCS/observe", + description="Visualize latent colors without VAE decoding — pure math color preview from LCS", + inputs=[ + io.Latent.Input("latent", tooltip="Latent from KSampler or similar"), + LCS_DATA.Input("lcs_data"), + io.Float.Input("sigma", default=0.0, min=0.0, max=1.0, step=0.01, + tooltip="Sigma for normalization (0.0 = final/clean, use sigma from sampler)"), + ], + outputs=[ + io.Image.Output(display_name="preview"), + ], + ) + + @classmethod + def execute(cls, latent, lcs_data, sigma) -> io.NodeOutput: + samples = latent["samples"] + result = _latent_to_color_preview(samples, lcs_data, sigma, upscale=8) + return io.NodeOutput(result) + + +class LCSStepObserver(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LCSStepObserver", + display_name="LCS Step Observer", + category="LCS/observe", + description="Patches model to save per-step LCS color previews to temp directory", + inputs=[ + io.Model.Input("model"), + LCS_DATA.Input("lcs_data"), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, lcs_data) -> io.NodeOutput: + m = model.clone() + step_counter = [0] + + def observer_fn(args): + denoised = args["denoised"] + sigma = args["sigma"] + sigma_val = float(sigma.flatten()[0]) + + # Generate color preview for first batch item + preview = _latent_to_color_preview( + denoised[:1], lcs_data, sigma_val, upscale=4 + ) + + # Save to temp directory + temp_dir = folder_paths.get_temp_directory() + img_np = (preview[0].numpy() * 255).clip(0, 255).astype(np.uint8) + pil_img = PILImage.fromarray(img_np) + filename = f"lcs_step_{step_counter[0]:03d}_s{sigma_val:.3f}.png" + pil_img.save(os.path.join(temp_dir, filename)) + step_counter[0] += 1 + + return denoised # Don't modify + + m.set_model_sampler_post_cfg_function(observer_fn) + return io.NodeOutput(m) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..146237c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +einops +safetensors