Satrio DewantonoandClaude Opus 4.6
9047141549
Allows connecting IMAGE outputs from other nodes as reference latent sources. Uses the KJNodes ImageBatchMulti pattern: an image_input_node_count widget + "Update image inputs" button to dynamically add/remove input slots. Node inputs are processed before uploaded images in the reference chain. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
402 lines
15 KiB
Python
402 lines
15 KiB
Python
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",),
|
|
"neg_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
|
|
}),
|
|
"image_input_node_count": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 1000,
|
|
"step": 1,
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image_input_1": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "IMAGE")
|
|
RETURN_NAMES = ("conditioning", "neg_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 _encode_and_append(self, image_tensor, vae, upscale_method, scale_megapixels,
|
|
result_conditioning, result_neg_conditioning, first_scaled_image):
|
|
"""Scale, VAE-encode a single image tensor and append to conditioning."""
|
|
scaled_image = self.scale_image(image_tensor, upscale_method, scale_megapixels)
|
|
|
|
if first_scaled_image is None:
|
|
first_scaled_image = scaled_image
|
|
|
|
latent_samples = vae.encode(scaled_image[:, :, :, :3])
|
|
|
|
result_conditioning = node_helpers.conditioning_set_values(
|
|
result_conditioning,
|
|
{"reference_latents": [latent_samples]},
|
|
append=True
|
|
)
|
|
result_neg_conditioning = node_helpers.conditioning_set_values(
|
|
result_neg_conditioning,
|
|
{"reference_latents": [latent_samples]},
|
|
append=True
|
|
)
|
|
|
|
return result_conditioning, result_neg_conditioning, first_scaled_image
|
|
|
|
def process(self, conditioning, neg_conditioning, vae, upscale_method, scale_megapixels, image_input_node_count, images, **kwargs):
|
|
"""
|
|
Process multiple images and chain them as reference latents.
|
|
Images come from two sources: uploaded files (via the widget) and
|
|
IMAGE node inputs (image_input_1, image_input_2, ...).
|
|
"""
|
|
result_conditioning = conditioning
|
|
result_neg_conditioning = neg_conditioning
|
|
first_scaled_image = None
|
|
|
|
# 1. Process IMAGE node inputs first (image_input_1, image_input_2, ...)
|
|
for i in range(1, image_input_node_count + 1):
|
|
image_tensor = kwargs.get(f"image_input_{i}")
|
|
if image_tensor is None:
|
|
continue
|
|
|
|
result_conditioning, result_neg_conditioning, first_scaled_image = \
|
|
self._encode_and_append(image_tensor, vae, upscale_method, scale_megapixels,
|
|
result_conditioning, result_neg_conditioning, first_scaled_image)
|
|
|
|
# 2. Process uploaded images from the drag-and-drop widget
|
|
try:
|
|
images_list = json.loads(images)
|
|
except json.JSONDecodeError:
|
|
images_list = []
|
|
|
|
if not isinstance(images_list, list):
|
|
images_list = []
|
|
|
|
images_list = [img for img in images_list if img and isinstance(img, str)]
|
|
|
|
for image_filename in images_list:
|
|
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
|
|
|
|
image_tensor = self.load_image(image_path)
|
|
result_conditioning, result_neg_conditioning, first_scaled_image = \
|
|
self._encode_and_append(image_tensor, vae, upscale_method, scale_megapixels,
|
|
result_conditioning, result_neg_conditioning, first_scaled_image)
|
|
|
|
# Fallback if no images were processed
|
|
if first_scaled_image is None:
|
|
first_scaled_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
|
|
|
|
return (result_conditioning, result_neg_conditioning, first_scaled_image)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, conditioning, neg_conditioning, vae, upscale_method, scale_megapixels, image_input_node_count, images, **kwargs):
|
|
"""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, neg_conditioning, vae, upscale_method, scale_megapixels, image_input_node_count, images, **kwargs):
|
|
"""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",),
|
|
"neg_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", "CONDITIONING", "IMAGE")
|
|
RETURN_NAMES = ("conditioning", "neg_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, neg_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
|
|
neg_conditioning: Input negative 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, neg_conditioning, empty_image)
|
|
|
|
# Process each image
|
|
result_conditioning = conditioning
|
|
result_neg_conditioning = neg_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
|
|
)
|
|
|
|
# Append to negative conditioning as well
|
|
result_neg_conditioning = node_helpers.conditioning_set_values(
|
|
result_neg_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, result_neg_conditioning, first_scaled_image)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, conditioning, neg_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()
|