From c2babe3fc86596157496fb89676c65c2ce66f0d7 Mon Sep 17 00:00:00 2001 From: Satrio Dewantono Date: Fri, 30 Jan 2026 18:43:46 -0800 Subject: [PATCH] Initial release: Reference Chain Conditioning nodes - ReferenceChainConditioning: Takes multiple images via UI, scales them, VAE encodes, and chains as reference latents into single conditioning - ReferenceChainConditioningBase64: Same functionality but accepts base64-encoded images for API usage Features: - Custom image list widget with drag-and-drop support - ComfyUI-native styling - Configurable scaling (megapixels) and upscale method - Outputs first scaled image for preview/further processing Co-Authored-By: Claude Opus 4.5 --- .gitignore | 36 +++ __init__.py | 15 + js/reference_chain.js | 653 ++++++++++++++++++++++++++++++++++++++++++ nodes.py | 380 ++++++++++++++++++++++++ pyproject.toml | 15 + 5 files changed, 1099 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 js/reference_chain.js create mode 100644 nodes.py create mode 100644 pyproject.toml diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..152a99f --- /dev/null +++ b/.gitignore @@ -0,0 +1,36 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so + +# Distribution / packaging +dist/ +build/ +*.egg-info/ +*.egg + +# Virtual environments +venv/ +env/ +.venv/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Testing +.pytest_cache/ +.coverage +htmlcov/ + +# Local development files +*.log +*.bak diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..822de69 --- /dev/null +++ b/__init__.py @@ -0,0 +1,15 @@ +from .nodes import ReferenceChainConditioning, ReferenceChainConditioningBase64 + +NODE_CLASS_MAPPINGS = { + "ReferenceChainConditioning": ReferenceChainConditioning, + "ReferenceChainConditioningBase64": ReferenceChainConditioningBase64, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ReferenceChainConditioning": "Reference Chain Conditioning", + "ReferenceChainConditioningBase64": "Reference Chain Conditioning (Base64)", +} + +WEB_DIRECTORY = "./js" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/js/reference_chain.js b/js/reference_chain.js new file mode 100644 index 0000000..24f1933 --- /dev/null +++ b/js/reference_chain.js @@ -0,0 +1,653 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +// Constants for the image list widget - matching ComfyUI style +const IMAGE_HEIGHT = 200; // Height for image display area +const BUTTON_HEIGHT = 26; // Standard ComfyUI button height +const DROP_ZONE_HEIGHT = 78; // Height for drop zone (3x button height) +const FILENAME_HEIGHT = 20; // Height for filename display +const DIMENSIONS_HEIGHT = 16; // Height for dimensions text +const SLOT_PADDING = 8; +const REMOVE_BUTTON_SIZE = 16; + +/** + * Custom widget for managing a list of images with previews + * Styled to match ComfyUI's LoadImage node + */ +class ImageListWidget { + constructor(node, inputName, inputData) { + this.node = node; + this.name = inputName; + this.type = "custom"; + this.options = { serialize: true }; + + // Array of image filenames + this._images = []; + + // Cache for loaded image elements with dimensions + this.imageCache = new Map(); + + // Track if we're currently uploading + this.isUploading = false; + + // Hover state for remove buttons + this.hoverRemoveIndex = -1; + + // Hover state for add button + this.hoverAddButton = false; + + // Hover/drag state for drop zone + this.dropZoneActive = false; + + // Initialize from default value + try { + const defaultVal = inputData?.[1]?.default || "[]"; + this._images = JSON.parse(defaultVal); + if (!Array.isArray(this._images)) { + this._images = []; + } + } catch (e) { + this._images = []; + } + } + + get value() { + return JSON.stringify(this._images); + } + + set value(v) { + try { + if (typeof v === "string") { + this._images = JSON.parse(v); + } else if (Array.isArray(v)) { + this._images = v; + } else { + this._images = []; + } + } catch (e) { + this._images = []; + } + if (!Array.isArray(this._images)) { + this._images = []; + } + + // Preload images + this.preloadImages(); + + // Request redraw + if (this.node) { + this.node.setDirtyCanvas(true, true); + } + } + + get images() { + return this._images; + } + + /** + * Preload images into cache + */ + preloadImages() { + for (const filename of this._images) { + if (!this.imageCache.has(filename)) { + const img = new Image(); + img.src = `/view?filename=${encodeURIComponent(filename)}&type=input`; + img.onload = () => { + // Store with dimensions + this.imageCache.set(filename, { + img: img, + width: img.naturalWidth, + height: img.naturalHeight + }); + if (this.node) { + this.node.setDirtyCanvas(true, true); + } + }; + // Temporarily store just the image while loading + this.imageCache.set(filename, { img: img, width: 0, height: 0 }); + } + } + } + + /** + * Calculate slot height for one image entry + */ + getSlotHeight() { + return FILENAME_HEIGHT + IMAGE_HEIGHT + DIMENSIONS_HEIGHT + SLOT_PADDING * 2; + } + + /** + * Calculate widget height based on number of images + */ + computeSize() { + const numImages = this._images.length; + const height = (numImages * this.getSlotHeight()) + DROP_ZONE_HEIGHT + SLOT_PADDING + BUTTON_HEIGHT + SLOT_PADDING * 2; + return [this.node?.size?.[0] || 200, height]; + } + + /** + * Add an image to the list + */ + addImage(filename) { + if (filename && !this._images.includes(filename)) { + this._images.push(filename); + this.preloadImages(); + this.updateNodeSize(); + this.node?.setDirtyCanvas(true, true); + } + } + + /** + * Remove an image from the list + */ + removeImage(index) { + if (index >= 0 && index < this._images.length) { + const filename = this._images[index]; + this._images.splice(index, 1); + this.imageCache.delete(filename); + this.updateNodeSize(); + this.node?.setDirtyCanvas(true, true); + } + } + + /** + * Update node size to fit all images + */ + updateNodeSize() { + if (!this.node) return; + + // Calculate height needed for this widget + const [_, widgetHeight] = this.computeSize(); + + // Calculate space taken by other widgets above this one + let otherWidgetsHeight = 0; + const widgetIndex = this.node.widgets?.indexOf(this); + if (widgetIndex > 0) { + for (let i = 0; i < widgetIndex; i++) { + const w = this.node.widgets[i]; + // Standard widget height is ~20-26px, use 26 as safe default + otherWidgetsHeight += (w.computedHeight || 26) + 4; + } + } + + // Title bar + header space (~70px) + other widgets + this widget's height + const headerSpace = 70; + const idealHeight = headerSpace + otherWidgetsHeight + widgetHeight; + const minHeight = 200; + + this.node.size[1] = Math.max(minHeight, idealHeight); + this.node.setDirtyCanvas(true, true); + } + + /** + * Upload an image file + */ + async uploadImage(file) { + if (this.isUploading) return; + this.isUploading = true; + + try { + const formData = new FormData(); + formData.append("image", file); + formData.append("type", "input"); + formData.append("overwrite", "false"); + + const response = await api.fetchApi("/upload/image", { + method: "POST", + body: formData, + }); + + if (response.ok) { + const result = await response.json(); + let filename = result.name; + if (result.subfolder) { + filename = `${result.subfolder}/${result.name}`; + } + this.addImage(filename); + } else { + console.error("[ReferenceChain] Upload failed:", response.status); + } + } catch (e) { + console.error("[ReferenceChain] Upload error:", e); + } finally { + this.isUploading = false; + } + } + + /** + * Open file dialog to select images + */ + openFileDialog() { + const input = document.createElement("input"); + input.type = "file"; + input.accept = "image/*"; + input.multiple = true; + input.onchange = async (e) => { + const files = e.target.files; + for (const file of files) { + await this.uploadImage(file); + } + }; + input.click(); + } + + /** + * Truncate filename to fit width + */ + truncateFilename(ctx, filename, maxWidth) { + const ellipsis = "..."; + let text = filename; + + if (ctx.measureText(text).width <= maxWidth) { + return text; + } + + // Try to show end of filename (more useful) + while (text.length > 0 && ctx.measureText(ellipsis + text).width > maxWidth) { + text = text.slice(1); + } + + return ellipsis + text; + } + + /** + * Draw the widget - ComfyUI style + */ + draw(ctx, node, widgetWidth, y, widgetHeight) { + const margin = 15; + const x = margin; + const width = widgetWidth - margin * 2; + + ctx.save(); + + let currentY = y + SLOT_PADDING; + const slotHeight = this.getSlotHeight(); + + // Store positions for click handling + this._slotPositions = []; + + // Draw each image slot + for (let i = 0; i < this._images.length; i++) { + const filename = this._images[i]; + const slotY = currentY + i * slotHeight; + + // Get cached image data + const cached = this.imageCache.get(filename); + const img = cached?.img; + const imgWidth = cached?.width || 0; + const imgHeight = cached?.height || 0; + + // --- Filename bar (like LoadImage combo) --- + const filenameY = slotY; + + // Background for filename + ctx.fillStyle = "#232323"; + ctx.beginPath(); + ctx.roundRect(x, filenameY, width, FILENAME_HEIGHT, [4, 4, 0, 0]); + ctx.fill(); + + // Border + ctx.strokeStyle = "#3a3a3a"; + ctx.lineWidth = 1; + ctx.stroke(); + + // Filename text + ctx.fillStyle = "#cccccc"; + ctx.font = "11px Arial"; + ctx.textAlign = "left"; + ctx.textBaseline = "middle"; + + const displayFilename = this.truncateFilename(ctx, filename, width - REMOVE_BUTTON_SIZE - 20); + ctx.fillText(displayFilename, x + 8, filenameY + FILENAME_HEIGHT / 2); + + // --- Remove button (X) in filename bar --- + const removeX = x + width - REMOVE_BUTTON_SIZE - 4; + const removeY = filenameY + (FILENAME_HEIGHT - REMOVE_BUTTON_SIZE) / 2 + 1; + + // Hover highlight + if (this.hoverRemoveIndex === i) { + ctx.fillStyle = "#555555"; + ctx.beginPath(); + ctx.roundRect(removeX, removeY, REMOVE_BUTTON_SIZE, REMOVE_BUTTON_SIZE - 2, 2); + ctx.fill(); + } + + // Draw X + ctx.strokeStyle = this.hoverRemoveIndex === i ? "#ffffff" : "#888888"; + ctx.lineWidth = 1.5; + ctx.beginPath(); + const xPad = 4; + ctx.moveTo(removeX + xPad, removeY + xPad); + ctx.lineTo(removeX + REMOVE_BUTTON_SIZE - xPad, removeY + REMOVE_BUTTON_SIZE - 2 - xPad); + ctx.moveTo(removeX + REMOVE_BUTTON_SIZE - xPad, removeY + xPad); + ctx.lineTo(removeX + xPad, removeY + REMOVE_BUTTON_SIZE - 2 - xPad); + ctx.stroke(); + + // --- Image area --- + const imageY = filenameY + FILENAME_HEIGHT; + + // Image background + ctx.fillStyle = "#1a1a1a"; + ctx.fillRect(x, imageY, width, IMAGE_HEIGHT); + + // Draw image (fill width, maintain aspect ratio) + if (img && img.complete && imgWidth > 0) { + const imgAspect = imgWidth / imgHeight; + const areaAspect = width / IMAGE_HEIGHT; + + let drawW, drawH, drawX, drawY; + + if (imgAspect > areaAspect) { + // Image is wider - fit to width + drawW = width; + drawH = width / imgAspect; + drawX = x; + drawY = imageY + (IMAGE_HEIGHT - drawH) / 2; + } else { + // Image is taller - fit to height + drawH = IMAGE_HEIGHT; + drawW = IMAGE_HEIGHT * imgAspect; + drawX = x + (width - drawW) / 2; + drawY = imageY; + } + + ctx.drawImage(img, drawX, drawY, drawW, drawH); + } else { + // Loading placeholder + ctx.fillStyle = "#333333"; + ctx.font = "12px Arial"; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText("Loading...", x + width / 2, imageY + IMAGE_HEIGHT / 2); + } + + // --- Dimensions text --- + const dimY = imageY + IMAGE_HEIGHT; + + ctx.fillStyle = "#1a1a1a"; + ctx.beginPath(); + ctx.roundRect(x, dimY, width, DIMENSIONS_HEIGHT, [0, 0, 4, 4]); + ctx.fill(); + + if (imgWidth > 0 && imgHeight > 0) { + ctx.fillStyle = "#888888"; + ctx.font = "10px Arial"; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText(`${imgWidth} × ${imgHeight}`, x + width / 2, dimY + DIMENSIONS_HEIGHT / 2); + } + + // Store slot position for click handling + this._slotPositions.push({ + y: slotY, + height: slotHeight, + filenameY: filenameY, + filenameHeight: FILENAME_HEIGHT, + removeX, + removeY, + removeSize: REMOVE_BUTTON_SIZE, + }); + } + + // --- Drop zone area --- + const dropZoneY = currentY + this._images.length * slotHeight + SLOT_PADDING; + + // Drop zone background + ctx.fillStyle = this.dropZoneActive ? "#2a2a2a" : "#1a1a1a"; + ctx.beginPath(); + ctx.roundRect(x, dropZoneY, width, DROP_ZONE_HEIGHT, 4); + ctx.fill(); + + // Drop zone dashed border + ctx.strokeStyle = this.dropZoneActive ? "#666666" : "#444444"; + ctx.lineWidth = 1; + ctx.setLineDash([5, 5]); + ctx.stroke(); + ctx.setLineDash([]); // Reset dash + + // Drop zone text + ctx.fillStyle = this.dropZoneActive ? "#888888" : "#555555"; + ctx.font = "12px Arial"; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText("drop image here", x + width / 2, dropZoneY + DROP_ZONE_HEIGHT / 2); + + // --- "choose file to upload" button (ComfyUI style) --- + const buttonY = dropZoneY + DROP_ZONE_HEIGHT + SLOT_PADDING; + + // Button background + ctx.fillStyle = this.hoverAddButton ? "#3a3a3a" : "#2a2a2a"; + ctx.beginPath(); + ctx.roundRect(x, buttonY, width, BUTTON_HEIGHT, 4); + ctx.fill(); + + // Button border + ctx.strokeStyle = "#4a4a4a"; + ctx.lineWidth = 1; + ctx.stroke(); + + // Button text + ctx.fillStyle = "#cccccc"; + ctx.font = "12px Arial"; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText("choose file to upload", x + width / 2, buttonY + BUTTON_HEIGHT / 2); + + ctx.restore(); + + // Store positions + this._dropZoneY = dropZoneY; + this._dropZoneHeight = DROP_ZONE_HEIGHT; + this._buttonY = buttonY; + this._buttonHeight = BUTTON_HEIGHT; + this._x = x; + this._width = width; + } + + /** + * Handle mouse events + */ + mouse(event, pos, node) { + const [mouseX, mouseY] = pos; + + if (event.type === "pointermove") { + // Update hover state for remove buttons + let newHoverIndex = -1; + if (this._slotPositions) { + for (let i = 0; i < this._slotPositions.length; i++) { + const slot = this._slotPositions[i]; + if ( + mouseX >= slot.removeX && + mouseX <= slot.removeX + slot.removeSize && + mouseY >= slot.removeY && + mouseY <= slot.removeY + slot.removeSize + ) { + newHoverIndex = i; + break; + } + } + } + + // Update hover state for add button + const newHoverAddButton = ( + this._buttonY !== undefined && + mouseX >= this._x && + mouseX <= this._x + this._width && + mouseY >= this._buttonY && + mouseY <= this._buttonY + this._buttonHeight + ); + + if (newHoverIndex !== this.hoverRemoveIndex || newHoverAddButton !== this.hoverAddButton) { + this.hoverRemoveIndex = newHoverIndex; + this.hoverAddButton = newHoverAddButton; + node.setDirtyCanvas(true, true); + } + return false; + } + + if (event.type === "pointerdown") { + // Check if clicked on a remove button + if (this._slotPositions) { + for (let i = 0; i < this._slotPositions.length; i++) { + const slot = this._slotPositions[i]; + if ( + mouseX >= slot.removeX && + mouseX <= slot.removeX + slot.removeSize && + mouseY >= slot.removeY && + mouseY <= slot.removeY + slot.removeSize + ) { + this.removeImage(i); + return true; + } + } + } + + // Check if clicked on add button + if ( + this._buttonY !== undefined && + mouseX >= this._x && + mouseX <= this._x + this._width && + mouseY >= this._buttonY && + mouseY <= this._buttonY + this._buttonHeight + ) { + this.openFileDialog(); + return true; + } + } + + return false; + } + + /** + * Serialize widget value + */ + serializeValue() { + return this.value; + } +} + +/** + * Register the extension + */ +app.registerExtension({ + name: "Comfy.ReferenceChain", + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name !== "ReferenceChainConditioning") { + return; + } + + // Store original onNodeCreated + const onNodeCreated = nodeType.prototype.onNodeCreated; + + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated?.apply(this, arguments); + + // Find the images widget and replace it + const imagesWidget = this.widgets?.find((w) => w.name === "images"); + if (imagesWidget) { + // Get index + const index = this.widgets.indexOf(imagesWidget); + + // Create custom widget + const customWidget = new ImageListWidget( + this, + "images", + nodeData.input?.required?.images + ); + + // Replace widget + this.widgets[index] = customWidget; + + // Initialize size + customWidget.updateNodeSize(); + } + + return result; + }; + + // Handle drag and drop + const onDragOver = nodeType.prototype.onDragOver; + nodeType.prototype.onDragOver = function (e) { + const result = onDragOver?.apply(this, arguments); + + if (e.dataTransfer?.types?.includes("Files")) { + e.preventDefault(); + + // Update drop zone highlight + const widget = this.widgets?.find((w) => w.name === "images"); + if (widget && widget instanceof ImageListWidget) { + if (!widget.dropZoneActive) { + widget.dropZoneActive = true; + this.setDirtyCanvas(true, true); + } + } + + return true; + } + + return result; + }; + + const onDragLeave = nodeType.prototype.onDragLeave; + nodeType.prototype.onDragLeave = function (e) { + const result = onDragLeave?.apply(this, arguments); + + // Clear drop zone highlight + const widget = this.widgets?.find((w) => w.name === "images"); + if (widget && widget instanceof ImageListWidget) { + if (widget.dropZoneActive) { + widget.dropZoneActive = false; + this.setDirtyCanvas(true, true); + } + } + + return result; + }; + + const onDragDrop = nodeType.prototype.onDragDrop; + nodeType.prototype.onDragDrop = async function (e) { + // Clear drop zone highlight + const widget = this.widgets?.find((w) => w.name === "images"); + if (widget && widget instanceof ImageListWidget) { + widget.dropZoneActive = false; + this.setDirtyCanvas(true, true); + } + + if (e.dataTransfer?.files?.length > 0) { + if (widget && widget instanceof ImageListWidget) { + for (const file of e.dataTransfer.files) { + if (file.type.startsWith("image/")) { + await widget.uploadImage(file); + } + } + return true; + } + } + + return onDragDrop?.apply(this, arguments); + }; + + // Handle paste + const onPaste = nodeType.prototype.onPaste; + nodeType.prototype.onPaste = async function (e) { + const widget = this.widgets?.find((w) => w.name === "images"); + if (widget && widget instanceof ImageListWidget) { + const items = e.clipboardData?.items; + if (items) { + for (const item of items) { + if (item.type.startsWith("image/")) { + const file = item.getAsFile(); + if (file) { + await widget.uploadImage(file); + return true; + } + } + } + } + } + + return onPaste?.apply(this, arguments); + }; + }, +}); diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..4ff2a60 --- /dev/null +++ b/nodes.py @@ -0,0 +1,380 @@ +import json +import math +import hashlib +import os +import base64 +from io import BytesIO + +import numpy as np +import torch +from PIL import Image, ImageOps, ImageSequence + +import folder_paths +import node_helpers +import comfy.utils + + +class ReferenceChainConditioning: + """ + A node that takes multiple images as input through a dynamic list UI, + processes them (scale + VAE encode), and chains them through Reference Latent + conditioning - outputting a single Conditioning that contains all reference latents. + """ + + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "conditioning": ("CONDITIONING",), + "vae": ("VAE",), + "upscale_method": (cls.upscale_methods, {"default": "lanczos"}), + "scale_megapixels": ("FLOAT", { + "default": 1.0, + "min": 0.01, + "max": 16.0, + "step": 0.01, + }), + "images": ("STRING", { + "default": "[]", + "multiline": False, + "image_list_upload": True, # Custom flag for frontend + }), + } + } + + RETURN_TYPES = ("CONDITIONING", "IMAGE") + RETURN_NAMES = ("conditioning", "first_image_scaled") + FUNCTION = "process" + CATEGORY = "conditioning" + + def load_image(self, image_path): + """Load a single image from the input directory and convert to tensor.""" + img = node_helpers.pillow(Image.open, image_path) + + output_images = [] + w, h = None, None + + for i in ImageSequence.Iterator(img): + i = node_helpers.pillow(ImageOps.exif_transpose, i) + + if i.mode == 'I': + i = i.point(lambda x: x * (1 / 255)) + image = i.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + output_images.append(image) + + if img.format == "MPO": + break # ignore all frames except the first one for MPO format + + if len(output_images) > 1: + output_image = torch.cat(output_images, dim=0) + elif len(output_images) == 1: + output_image = output_images[0] + else: + raise ValueError(f"Could not load image from {image_path}") + + return output_image + + def scale_image(self, image, upscale_method, megapixels): + """Scale image to target megapixels.""" + # Image is in BHWC format, move to BCHW for upscaling + samples = image.movedim(-1, 1) + total = megapixels * 1024 * 1024 + + scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) + width = round(samples.shape[3] * scale_by) + height = round(samples.shape[2] * scale_by) + + # Ensure dimensions are at least 8 (VAE requirement) + width = max(8, width) + height = max(8, height) + + s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled") + # Move back to BHWC + s = s.movedim(1, -1) + return s + + def process(self, conditioning, vae, upscale_method, scale_megapixels, images): + """ + Process multiple images and chain them as reference latents. + + Args: + conditioning: Input conditioning to append reference latents to + vae: VAE model for encoding images to latents + upscale_method: Method for scaling images + scale_megapixels: Target megapixels for image scaling + images: JSON array of image filenames + + Returns: + Conditioning with all reference latents appended + """ + # Parse JSON array of image filenames + try: + images_list = json.loads(images) + except json.JSONDecodeError: + images_list = [] + + if not isinstance(images_list, list): + images_list = [] + + # Filter out empty strings + images_list = [img for img in images_list if img and isinstance(img, str)] + + # If no images, return conditioning unchanged with empty image + if not images_list: + # Return empty 1x1 black image as placeholder + empty_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32) + return (conditioning, empty_image) + + # Process each image + result_conditioning = conditioning + first_scaled_image = None + + for image_filename in images_list: + # Get full path using folder_paths + image_path = folder_paths.get_annotated_filepath(image_filename) + + if not os.path.exists(image_path): + print(f"[ReferenceChainConditioning] Warning: Image not found: {image_path}") + continue + + # Load image + image_tensor = self.load_image(image_path) + + # Scale to target megapixels + scaled_image = self.scale_image(image_tensor, upscale_method, scale_megapixels) + + # Store first scaled image for output + if first_scaled_image is None: + first_scaled_image = scaled_image + + # VAE encode + # VAE expects BHWC format with values 0-1 + latent_samples = vae.encode(scaled_image[:, :, :, :3]) + + # Append to conditioning using the same method as ReferenceLatent + result_conditioning = node_helpers.conditioning_set_values( + result_conditioning, + {"reference_latents": [latent_samples]}, + append=True + ) + + # Ensure we have a first image (fallback to empty if somehow none processed) + if first_scaled_image is None: + first_scaled_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32) + + return (result_conditioning, first_scaled_image) + + @classmethod + def IS_CHANGED(cls, conditioning, vae, upscale_method, scale_megapixels, images): + """Check if the node needs to be re-executed based on image changes.""" + try: + images_list = json.loads(images) + except json.JSONDecodeError: + return "" + + if not isinstance(images_list, list): + return "" + + # Create a hash of all image files + m = hashlib.sha256() + + for image_filename in images_list: + if not image_filename or not isinstance(image_filename, str): + continue + + image_path = folder_paths.get_annotated_filepath(image_filename) + + if os.path.exists(image_path): + with open(image_path, 'rb') as f: + m.update(f.read()) + + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(cls, conditioning, vae, upscale_method, scale_megapixels, images): + """Validate that all specified images exist.""" + try: + images_list = json.loads(images) + except json.JSONDecodeError: + return True # Empty/invalid JSON is okay, just means no images + + if not isinstance(images_list, list): + return True + + for image_filename in images_list: + if not image_filename or not isinstance(image_filename, str): + continue + + if not folder_paths.exists_annotated_filepath(image_filename): + return f"Invalid image file: {image_filename}" + + return True + + +class ReferenceChainConditioningBase64: + """ + A node that takes multiple base64-encoded images as input, + processes them (scale + VAE encode), and chains them through Reference Latent + conditioning - outputting a single Conditioning that contains all reference latents. + + This version is optimized for API usage where images are passed directly as base64 strings. + """ + + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "conditioning": ("CONDITIONING",), + "vae": ("VAE",), + "upscale_method": (cls.upscale_methods, {"default": "lanczos"}), + "scale_megapixels": ("FLOAT", { + "default": 1.0, + "min": 0.01, + "max": 16.0, + "step": 0.01, + }), + "images_base64": ("STRING", { + "default": "[]", + "multiline": True, + }), + } + } + + RETURN_TYPES = ("CONDITIONING", "IMAGE") + RETURN_NAMES = ("conditioning", "first_image_scaled") + FUNCTION = "process" + CATEGORY = "conditioning" + + def load_image_from_base64(self, base64_data): + """Load a single image from base64 string and convert to tensor.""" + # Handle data URL prefix if present (e.g., "data:image/png;base64,...") + if "," in base64_data: + base64_data = base64_data.split(",", 1)[1] + + # Decode base64 + image_data = base64.b64decode(base64_data) + + # Open with PIL + img = Image.open(BytesIO(image_data)) + + # Handle EXIF orientation + img = ImageOps.exif_transpose(img) + + # Convert to RGB + if img.mode == 'I': + img = img.point(lambda x: x * (1 / 255)) + image = img.convert("RGB") + + # Convert to tensor + image_np = np.array(image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_np)[None,] + + return image_tensor + + def scale_image(self, image, upscale_method, megapixels): + """Scale image to target megapixels.""" + # Image is in BHWC format, move to BCHW for upscaling + samples = image.movedim(-1, 1) + total = megapixels * 1024 * 1024 + + scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) + width = round(samples.shape[3] * scale_by) + height = round(samples.shape[2] * scale_by) + + # Ensure dimensions are at least 8 (VAE requirement) + width = max(8, width) + height = max(8, height) + + s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled") + # Move back to BHWC + s = s.movedim(1, -1) + return s + + def process(self, conditioning, vae, upscale_method, scale_megapixels, images_base64): + """ + Process multiple base64 images and chain them as reference latents. + + Args: + conditioning: Input conditioning to append reference latents to + vae: VAE model for encoding images to latents + upscale_method: Method for scaling images + scale_megapixels: Target megapixels for image scaling + images_base64: JSON array of base64-encoded image strings + + Returns: + Conditioning with all reference latents appended + """ + # Parse JSON array of base64 strings + try: + images_list = json.loads(images_base64) + except json.JSONDecodeError: + # If not valid JSON, treat as single base64 string + images_list = [images_base64] if images_base64.strip() else [] + + if not isinstance(images_list, list): + images_list = [images_list] if images_list else [] + + # Filter out empty strings + images_list = [img for img in images_list if img and isinstance(img, str) and img.strip()] + + # If no images, return conditioning unchanged with empty image + if not images_list: + empty_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32) + return (conditioning, empty_image) + + # Process each image + result_conditioning = conditioning + first_scaled_image = None + + for i, base64_data in enumerate(images_list): + try: + # Load image from base64 + image_tensor = self.load_image_from_base64(base64_data) + + # Scale to target megapixels + scaled_image = self.scale_image(image_tensor, upscale_method, scale_megapixels) + + # Store first scaled image for output + if first_scaled_image is None: + first_scaled_image = scaled_image + + # VAE encode + latent_samples = vae.encode(scaled_image[:, :, :, :3]) + + # Append to conditioning using the same method as ReferenceLatent + result_conditioning = node_helpers.conditioning_set_values( + result_conditioning, + {"reference_latents": [latent_samples]}, + append=True + ) + except Exception as e: + print(f"[ReferenceChainConditioningBase64] Warning: Failed to process image {i}: {e}") + continue + + # Ensure we have a first image (fallback to empty if somehow none processed) + if first_scaled_image is None: + first_scaled_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32) + + return (result_conditioning, first_scaled_image) + + @classmethod + def IS_CHANGED(cls, conditioning, vae, upscale_method, scale_megapixels, images_base64): + """Check if the node needs to be re-executed based on input changes.""" + # Hash the base64 data directly + m = hashlib.sha256() + m.update(images_base64.encode('utf-8')) + return m.digest().hex() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..150d666 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "comfyui-referencechain" +description = "A ComfyUI node that takes multiple images, scales them, VAE encodes them, and chains them as reference latents into a single conditioning output. Useful for edit models like Flux-2-klein that accept multiple reference images." +version = "1.0.0" +license = {text = "MIT"} +dependencies = [] + +[project.urls] +Repository = "https://github.com/remingtonspaz/ComfyUI-ReferenceChain" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "remingtonspaz" +DisplayName = "ComfyUI-ReferenceChain" +# Icon = "https://example.com/icon.png"