From f9a0c9535f697d87ea54f4e1294b66893e2114cd Mon Sep 17 00:00:00 2001 From: shanfang Date: Thu, 23 Apr 2026 13:12:54 +0800 Subject: [PATCH] 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) --- nodes/image_prep.py | 135 +++++++++++++++++++------------------------- 1 file changed, 59 insertions(+), 76 deletions(-) diff --git a/nodes/image_prep.py b/nodes/image_prep.py index 8ec23be..3e1a02a 100644 --- a/nodes/image_prep.py +++ b/nodes/image_prep.py @@ -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"}