feat: expand ImagePrep to support up to 8 image inputs
- Generate optional image slots via loop instead of hardcoding - Use **kwargs to collect dynamic image inputs - Widget inputs (format, quality) ordered before connection slots to prevent widgets_values serialization drift - Fully backward compatible with existing 4-image workflows Closes #26 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
shanfang
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
7d1fe48969
commit
f9a0c9535f
+59
-76
@@ -5,24 +5,29 @@ from PIL import Image
|
||||
from typing import Optional, Union
|
||||
import torch
|
||||
|
||||
# ── Config ──────────────────────────────────────────────────────────────
|
||||
MAX_IMAGES = 8
|
||||
|
||||
|
||||
class ImagePrep:
|
||||
"""
|
||||
Custom node for preprocessing images before passing them to LLMs.
|
||||
Preprocess up to {MAX_IMAGES} images into base64 data-URIs for LLM vision input.
|
||||
Only connected slots are processed; unused slots cost nothing.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
optional = {
|
||||
"format": (["PNG", "JPEG", "WebP", "GIF", "BMP", "TIFF"], {"default": "PNG"}),
|
||||
"quality": (["High", "Medium", "Low"], {"default": "High"}),
|
||||
}
|
||||
for i in range(2, MAX_IMAGES + 1):
|
||||
optional[f"image_{i}"] = ("IMAGE", {"default": None})
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", {"default": None}),
|
||||
},
|
||||
"optional": {
|
||||
"image_2": ("IMAGE", {"default": None}),
|
||||
"image_3": ("IMAGE", {"default": None}),
|
||||
"image_4": ("IMAGE", {"default": None}),
|
||||
"format": (["PNG", "JPEG", "WebP", "GIF", "BMP", "TIFF"], {"default": "PNG"}),
|
||||
"quality": (["High", "Medium", "Low"], {"default": "High"}),
|
||||
}
|
||||
"optional": optional,
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
@@ -30,83 +35,61 @@ class ImagePrep:
|
||||
FUNCTION = "preprocess"
|
||||
CATEGORY = "🚦ComfyUI_LLMs_Toolkit/Image"
|
||||
|
||||
def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
|
||||
"""Convert a single [H, W, C] tensor to PIL Image."""
|
||||
numpy_image = tensor.cpu().numpy()
|
||||
if len(numpy_image.shape) == 3 and numpy_image.shape[2] == 1:
|
||||
numpy_image = numpy_image.squeeze(-1)
|
||||
numpy_image = (numpy_image * 255).astype('uint8')
|
||||
return Image.fromarray(numpy_image)
|
||||
# ── Internal ────────────────────────────────────────────────────────
|
||||
|
||||
def preprocess(self, image: Optional[Union[str, Image.Image, torch.Tensor]] = None,
|
||||
image_2: Optional[Union[str, Image.Image, torch.Tensor]] = None,
|
||||
image_3: Optional[Union[str, Image.Image, torch.Tensor]] = None,
|
||||
image_4: Optional[Union[str, Image.Image, torch.Tensor]] = None,
|
||||
format: str = "PNG", quality: str = "High"):
|
||||
quality_map = {"High": 95, "Medium": 75, "Low": 50}
|
||||
quality_str = quality
|
||||
quality_val = quality_map.get(quality_str, 95)
|
||||
|
||||
images_to_process = [img for img in [image, image_2, image_3, image_4] if img is not None]
|
||||
|
||||
if not images_to_process:
|
||||
def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
|
||||
arr = tensor.cpu().numpy()
|
||||
if len(arr.shape) == 3 and arr.shape[2] == 1:
|
||||
arr = arr.squeeze(-1)
|
||||
return Image.fromarray((arr * 255).astype("uint8"))
|
||||
|
||||
def _encode(self, image: Image.Image, fmt: str, quality_val: int) -> str:
|
||||
size_map = {"High": 1024, "Medium": 768, "Low": 512}
|
||||
max_size = size_map.get(fmt, 1024)
|
||||
if max(image.size) > max_size:
|
||||
image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
|
||||
|
||||
buf = io.BytesIO()
|
||||
kw = {"format": fmt}
|
||||
if fmt in ("JPEG", "WebP"):
|
||||
kw["quality"] = quality_val
|
||||
image.save(buf, **kw)
|
||||
|
||||
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
print(f"[LLMs_Toolkit] encoded={buf.tell()/1024:.1f}KB {fmt} ({image.width}x{image.height})")
|
||||
return f"data:image/{fmt.lower()};base64,{b64}"
|
||||
|
||||
def _process_tensor(self, t: torch.Tensor, fmt, qval):
|
||||
urls = []
|
||||
if len(t.shape) == 4:
|
||||
for i in range(t.shape[0]):
|
||||
urls.append(self._encode(self._tensor_to_pil(t[i]), fmt, qval))
|
||||
else:
|
||||
urls.append(self._encode(self._tensor_to_pil(t), fmt, qval))
|
||||
return urls
|
||||
|
||||
# ── Entry ───────────────────────────────────────────────────────────
|
||||
|
||||
def preprocess(self, image=None, format="PNG", quality="High", **kwargs):
|
||||
quality_val = {"High": 95, "Medium": 75, "Low": 50}.get(quality, 95)
|
||||
|
||||
slots = [image] + [kwargs.get(f"image_{i}") for i in range(2, MAX_IMAGES + 1)]
|
||||
imgs = [s for s in slots if s is not None]
|
||||
|
||||
if not imgs:
|
||||
raise ValueError("At least one image input must be provided.")
|
||||
|
||||
image_urls = []
|
||||
|
||||
for img in images_to_process:
|
||||
urls = []
|
||||
for img in imgs:
|
||||
if isinstance(img, torch.Tensor):
|
||||
if len(img.shape) == 4:
|
||||
for i in range(img.shape[0]):
|
||||
pil_image = self._tensor_to_pil(img[i])
|
||||
url = self._process_single_image(pil_image, format, quality_str, quality_val)
|
||||
image_urls.append(url)
|
||||
else:
|
||||
pil_image = self._tensor_to_pil(img)
|
||||
url = self._process_single_image(pil_image, format, quality_str, quality_val)
|
||||
image_urls.append(url)
|
||||
|
||||
urls.extend(self._process_tensor(img, format, quality_val))
|
||||
elif isinstance(img, Image.Image):
|
||||
# Single PIL image
|
||||
url = self._process_single_image(img, format, quality_str, quality_val)
|
||||
image_urls.append(url)
|
||||
|
||||
urls.append(self._encode(img, format, quality_val))
|
||||
else:
|
||||
raise ValueError("Unsupported image type. Expected torch.Tensor or PIL.Image.")
|
||||
|
||||
# If only one image, return as string for backward compatibility?
|
||||
# Plan says: "ImagePreprocessor output will change from str to List[str]"
|
||||
# To be safe for ComfyUI string passing, let's return a list
|
||||
# We use JSON serialization to avoid ComfyUI auto-batching the list
|
||||
return (json.dumps(image_urls),)
|
||||
|
||||
def _process_single_image(self, image: Image.Image, format: str, quality_str: str, quality_val: int) -> str:
|
||||
# Resize image
|
||||
size_map = {"High": 1024, "Medium": 768, "Low": 512}
|
||||
max_size = size_map.get(quality_str, 1024)
|
||||
|
||||
w, h = image.size
|
||||
# Only resize if larger than max_size
|
||||
if max(w, h) > max_size:
|
||||
image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
|
||||
# print(f"[LLMs_Toolkit] resize={max_size}px")
|
||||
|
||||
# Convert to base64
|
||||
buffered = io.BytesIO()
|
||||
save_kwargs = {"format": format}
|
||||
if format in ["JPEG", "WebP"]:
|
||||
save_kwargs["quality"] = quality_val
|
||||
|
||||
image.save(buffered, **save_kwargs)
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||||
image_url = f"data:image/{format.lower()};base64,{img_str}"
|
||||
|
||||
size_kb = buffered.tell() / 1024
|
||||
print(f"[LLMs_Toolkit] encoded={size_kb:.1f}KB {format} ({image.width}x{image.height})")
|
||||
|
||||
return image_url
|
||||
return (json.dumps(urls),)
|
||||
|
||||
|
||||
# Register the node with ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {"ImagePrep": ImagePrep}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"ImagePrep": "Image Prep"}
|
||||
|
||||
Reference in New Issue
Block a user