diff --git a/__init__.py b/__init__.py index 626725a..49d7a85 100644 --- a/__init__.py +++ b/__init__.py @@ -24,6 +24,7 @@ from .nodes import ( SimpleChatImageGrid, SimpleChatAnimaPromptXYMatrix, SimpleChatXYPlot, + SimpleChatXYCellPrefix, ) # Register API routes @@ -51,6 +52,7 @@ NODE_CLASS_MAPPINGS = { "SimpleChatImageGrid": SimpleChatImageGrid, "SimpleChatAnimaPromptXYMatrix": SimpleChatAnimaPromptXYMatrix, "SimpleChatXYPlot": SimpleChatXYPlot, + "SimpleChatXYCellPrefix": SimpleChatXYCellPrefix, } # Display names for the UI @@ -74,6 +76,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SimpleChatImageGrid": "Image Grid (Batch)", "SimpleChatAnimaPromptXYMatrix": "Anima XY Matrix (JSON List)", "SimpleChatXYPlot": "XY Plot (Labels)", + "SimpleChatXYCellPrefix": "XY Cell Prefix (Batch)", } # Expose web directory for JS extensions diff --git a/assets/ArtistBlendXYPlot_Test.json b/assets/ArtistBlendXYPlot_Test.json index ca7e133..5884a76 100644 --- a/assets/ArtistBlendXYPlot_Test.json +++ b/assets/ArtistBlendXYPlot_Test.json @@ -1,8 +1,8 @@ { "id": "simplechat-artist-blend-xyplot-test", "revision": 0, - "last_node_id": 13, - "last_link_id": 20, + "last_node_id": 14, + "last_link_id": 23, "nodes": [ { "id": 1, @@ -68,7 +68,8 @@ "combine", "newline", "\\n", - true + true, + "upper_triangle" ] }, { @@ -197,12 +198,31 @@ "order": 9, "mode": 0, "inputs": [ - { "name": "images", "type": "IMAGE", "link": 14 } + { "name": "images", "type": "IMAGE", "link": 14 }, + { "name": "filename_prefix", "type": "STRING", "link": 23 } ], "outputs": [], "properties": { "Node name for S&R": "SaveImage" }, "widgets_values": ["artist_blend_each"] }, + { + "id": 14, + "type": "SimpleChatXYCellPrefix", + "pos": [1460, 1020], + "size": [420, 220], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { "name": "x_labels", "type": "STRING", "link": 21 }, + { "name": "y_labels", "type": "STRING", "link": 22 } + ], + "outputs": [ + { "name": "filename_prefix", "type": "STRING", "links": [23] } + ], + "properties": { "aux_id": "Moeblack/ComfyUI-SimpleChat", "Node name for S&R": "SimpleChatXYCellPrefix" }, + "widgets_values": ["+", true, "", true, "upper_triangle"] + }, { "id": 10, "type": "SimpleChatXYPlot", @@ -273,7 +293,10 @@ [17, 3, 3, 10, 3, "STRING"], [18, 9, 0, 10, 0, "IMAGE"], [19, 10, 0, 11, 0, "IMAGE"], - [20, 10, 0, 12, 0, "IMAGE"] + [20, 10, 0, 12, 0, "IMAGE"], + [21, 3, 2, 14, 0, "STRING"], + [22, 3, 3, 14, 1, "STRING"], + [23, 14, 0, 13, 1, "STRING"] ], "groups": [], "config": {}, diff --git a/nodes/__init__.py b/nodes/__init__.py index b0dc8b5..5e96ac2 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -20,6 +20,7 @@ from .anima_prompt_router import SimpleChatAnimaPromptRouter from .anima_prompt_xy_matrix import SimpleChatAnimaPromptXYMatrix from .image_grid import SimpleChatImageGrid from .image_xy_plot import SimpleChatXYPlot +from .xy_cell_prefix import SimpleChatXYCellPrefix __all__ = [ "SimpleChatConfig", @@ -41,4 +42,5 @@ __all__ = [ "SimpleChatAnimaPromptXYMatrix", "SimpleChatImageGrid", "SimpleChatXYPlot", + "SimpleChatXYCellPrefix", ] diff --git a/nodes/anima_prompt_xy_matrix.py b/nodes/anima_prompt_xy_matrix.py index 7e35d13..4ac8d26 100644 --- a/nodes/anima_prompt_xy_matrix.py +++ b/nodes/anima_prompt_xy_matrix.py @@ -204,6 +204,14 @@ class SimpleChatAnimaPromptXYMatrix: "pair_join": (["newline", "comma", "space", "custom"], {"default": "newline"}), "custom_join": ("STRING", {"default": "\\n", "tooltip": "Used when pair_join=custom. Use \\n for newline."}), "diagonal_single": ("BOOLEAN", {"default": True, "tooltip": "When X==Y, use single value on diagonal (recommended)."}), + # Keep this as the last widget for backward compatibility with older workflows. + "matrix_mode": ( + ["full", "upper_triangle", "diagonal_only"], + { + "default": "full", + "tooltip": "When x_field==y_field (e.g. artist x artist), choose full matrix or run fewer cells.", + }, + ), } } @@ -233,6 +241,7 @@ class SimpleChatAnimaPromptXYMatrix: pair_join: str = "newline", custom_join: str = "\\n", diagonal_single: bool = True, + matrix_mode: str = "full", ): raw = _strip_code_fence(json_text) if strip_code_fence else json_text base = _try_parse_json(raw) if isinstance(raw, str) else raw @@ -311,11 +320,30 @@ class SimpleChatAnimaPromptXYMatrix: return json.dumps(obj, ensure_ascii=False, indent=2) + same_axis = (x_key and y_key and x_key == y_key) or (x_field == y_field) + # Row-major (y rows) out_list: List[str] = [] - for yv in ys: - for xv in xs: - out_list.append(_build_one(xv, yv)) + if matrix_mode != "full" and same_axis and len(xs) == len(ys): + n = len(xs) + if matrix_mode == "upper_triangle": + for r in range(n): + yv = ys[r] + for c in range(r, n): + xv = xs[c] + out_list.append(_build_one(xv, yv)) + elif matrix_mode == "diagonal_only": + for i in range(n): + out_list.append(_build_one(xs[i], ys[i])) + else: + # Unknown mode -> fallback + for yv in ys: + for xv in xs: + out_list.append(_build_one(xv, yv)) + else: + for yv in ys: + for xv in xs: + out_list.append(_build_one(xv, yv)) columns = len(xs) x_labels = "\n".join(xs) diff --git a/nodes/image_xy_plot.py b/nodes/image_xy_plot.py index f749aea..1bd01a9 100644 --- a/nodes/image_xy_plot.py +++ b/nodes/image_xy_plot.py @@ -8,13 +8,56 @@ This is meant for fast artist/background testing. from __future__ import annotations import math -from typing import List, Tuple +from typing import Any, List, Tuple import numpy as np import torch from PIL import Image, ImageDraw, ImageFont +def _first(v: Any, default: Any = None) -> Any: + if isinstance(v, (list, tuple)): + return v[0] if len(v) > 0 else default + return v + + +def _flatten_images(images: Any) -> torch.Tensor: + """ + Accept either: + - torch.Tensor: (B,H,W,C) or (H,W,C) + - list/tuple of torch.Tensor batches (common when upstream is OUTPUT_IS_LIST) + Return a single torch.Tensor (B,H,W,C). + """ + if isinstance(images, torch.Tensor): + t = images + if t.dim() == 3: + t = t.unsqueeze(0) + return t + + if isinstance(images, (list, tuple)): + batches: list[torch.Tensor] = [] + for item in images: + if item is None: + continue + if isinstance(item, torch.Tensor): + t = item + if t.dim() == 3: + t = t.unsqueeze(0) + if t.dim() != 4: + raise ValueError(f"Expected IMAGE tensor with 4 dims (B,H,W,C), got shape={tuple(t.shape)}") + batches.append(t) + continue + if isinstance(item, (list, tuple)): + batches.append(_flatten_images(item)) + continue + raise ValueError(f"Unsupported images list element type: {type(item)}") + if not batches: + raise ValueError("images list is empty.") + return torch.cat(batches, dim=0) + + raise ValueError(f"Unsupported images type: {type(images)}") + + def _split_labels(text: str) -> List[str]: if not isinstance(text, str): return [] @@ -43,6 +86,9 @@ def _to_tensor(img: Image.Image) -> torch.Tensor: class SimpleChatXYPlot: + # If upstream provides a list (from OUTPUT_IS_LIST), run ONCE with the whole list. + INPUT_IS_LIST = True + @classmethod def INPUT_TYPES(cls): return { @@ -70,7 +116,7 @@ class SimpleChatXYPlot: def plot( self, - images: torch.Tensor, + images: Any, columns: int = 4, x_labels: str = "", y_labels: str = "", @@ -80,26 +126,61 @@ class SimpleChatXYPlot: text_color: str = "white", title: str = "", ): + # When INPUT_IS_LIST=True, many inputs may arrive as lists; pick the first. + columns = _first(columns, 4) + x_labels = _first(x_labels, "") or "" + y_labels = _first(y_labels, "") or "" + padding = _first(padding, 8) + label_padding = _first(label_padding, 8) + bg_color = _first(bg_color, "black") + text_color = _first(text_color, "white") + title = _first(title, "") or "" + + images = _flatten_images(images) if images is None or not isinstance(images, torch.Tensor): - raise ValueError("images must be a torch Tensor (ComfyUI IMAGE).") + raise ValueError("images must be a torch Tensor (ComfyUI IMAGE) or a list of IMAGE tensors.") if images.dim() != 4: raise ValueError(f"Expected IMAGE tensor with 4 dims (B,H,W,C), got shape={tuple(images.shape)}") b, h, w, c = images.shape columns = max(1, int(columns)) - rows = int(math.ceil(b / columns)) if b > 0 else 1 padding = max(0, int(padding)) label_padding = max(0, int(label_padding)) xs = _split_labels(x_labels) ys = _split_labels(y_labels) + has_x_labels = len(xs) > 0 + has_y_labels = len(ys) > 0 - # Try to align labels length with grid + # Prefer labels to define grid shape (scientific-style XY plots). + if has_x_labels: + columns = max(1, len(xs)) + + rows = int(math.ceil(b / columns)) if b > 0 else 1 + if has_y_labels: + rows = max(1, len(ys)) + + # Ensure grid is large enough to hold all images. + if b > rows * columns: + rows = int(math.ceil(b / columns)) + + # Pad labels if too short. if len(xs) < columns: xs = xs + [""] * (columns - len(xs)) if len(ys) < rows: ys = ys + [""] * (rows - len(ys)) + # Detect triangular inputs for square grids: + # - upper_triangle: N(N+1)/2 images (row-major over cells where col>=row) + # - diagonal_only: N images (row-major over diag cells) + triangle_mode = None # "upper_triangle" | "diagonal_only" + if has_x_labels and has_y_labels and rows == columns: + n = columns + if b == (n * (n + 1)) // 2: + triangle_mode = "upper_triangle" + elif b == n: + triangle_mode = "diagonal_only" + # Font font = ImageFont.load_default() @@ -149,13 +230,53 @@ class SimpleChatXYPlot: draw.text((x0, y_cursor), label, fill=fg, font=font) y_cursor += x_h + label_padding - # Paste images + Y labels - for idx in range(b): - r = idx // columns - col = idx % columns - x = left_margin + col * (w + padding) - y = y_cursor + r * (h + padding) - canvas.paste(_to_pil(images[idx]), (x, y)) + # Paste images (+ fill masked cells for triangular layouts) + # Use a mid-gray mask for the unused half (common in scientific matrices). + mask_color = (64, 64, 64) if bg_color == "black" else (220, 220, 220) + mask = Image.new("RGB", (w, h), color=mask_color) + + if triangle_mode == "upper_triangle": + idx = 0 + for r in range(rows): + for col in range(columns): + x = left_margin + col * (w + padding) + y = y_cursor + r * (h + padding) + if col < r: + canvas.paste(mask, (x, y)) + continue + if idx < b: + canvas.paste(_to_pil(images[idx]), (x, y)) + idx += 1 + else: + canvas.paste(mask, (x, y)) + elif triangle_mode == "diagonal_only": + idx = 0 + for r in range(rows): + for col in range(columns): + x = left_margin + col * (w + padding) + y = y_cursor + r * (h + padding) + if col == r and idx < b: + canvas.paste(_to_pil(images[idx]), (x, y)) + idx += 1 + else: + canvas.paste(mask, (x, y)) + else: + # Standard row-major fill + max_cells = rows * columns + for idx in range(min(b, max_cells)): + r = idx // columns + col = idx % columns + x = left_margin + col * (w + padding) + y = y_cursor + r * (h + padding) + canvas.paste(_to_pil(images[idx]), (x, y)) + # Fill remaining cells (if any) for a tidy grid. + if b < max_cells: + for idx in range(b, max_cells): + r = idx // columns + col = idx % columns + x = left_margin + col * (w + padding) + y = y_cursor + r * (h + padding) + canvas.paste(mask, (x, y)) if y_w > 0: for r in range(rows): diff --git a/nodes/xy_cell_prefix.py b/nodes/xy_cell_prefix.py new file mode 100644 index 0000000..f148cd6 --- /dev/null +++ b/nodes/xy_cell_prefix.py @@ -0,0 +1,154 @@ +""" +XY Cell Prefix (Batch) + +Generate a filename_prefix list for each XY cell (row-major), +without changing the prompt itself. + +Typical use: +- Connect x_labels/y_labels from `Anima XY Matrix (JSON List)` +- Connect output `filename_prefix` to `SaveImage.filename_prefix` +""" + +from __future__ import annotations + +import re +from typing import Any, List + + +_ILLEGAL_FS_CHARS_RE = re.compile(r'[<>:"/\\\\|?*]+') +_WS_RE = re.compile(r"\s+") + + +def _first(v: Any, default: Any = None) -> Any: + if isinstance(v, (list, tuple)): + return v[0] if len(v) > 0 else default + return v + + +def _split_labels(text: str) -> List[str]: + if not isinstance(text, str): + return [] + raw = text.replace("\r\n", "\n").replace("\r", "\n") + out: List[str] = [] + for line in raw.split("\n"): + s = (line or "").strip() + if not s: + continue + if s.startswith("#") or s.startswith("//"): + continue + out.append(s) + return out + + +def _sanitize_prefix(s: str, max_len: int = 120) -> str: + s = (s or "").strip() + if not s: + return "" + s = _ILLEGAL_FS_CHARS_RE.sub("_", s) + s = s.replace("\n", "_").replace("\r", "_").replace("\t", "_") + s = _WS_RE.sub("_", s) + s = s.strip(" .") # Windows doesn't allow trailing dot/space + if len(s) > max_len: + s = s[:max_len].rstrip(" .") + return s + + +class SimpleChatXYCellPrefix: + """ + Build per-cell filename prefixes from x/y labels. + + - Row-major order (y rows, then x columns), matching `SimpleChatAnimaPromptXYMatrix`. + - Diagonal can be single name when x==y. + - Default joiner for display is '+' (purely for naming; doesn't affect prompts). + """ + + # In list pipelines, run ONCE with the whole list (avoid N-times expansion). + INPUT_IS_LIST = True + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "x_labels": ("STRING", {"multiline": True, "default": "", "forceInput": True}), + "y_labels": ("STRING", {"multiline": True, "default": "", "forceInput": True}), + }, + "optional": { + "joiner": ("STRING", {"default": "+"}), + "diagonal_single": ("BOOLEAN", {"default": True}), + "base_prefix": ("STRING", {"default": ""}), + "sanitize": ("BOOLEAN", {"default": True}), + # Keep this as the last widget for future backward compatibility. + "matrix_mode": (["full", "upper_triangle", "diagonal_only"], {"default": "full"}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename_prefix",) + OUTPUT_IS_LIST = (True,) + FUNCTION = "build" + CATEGORY = "SimpleChat/Utils" + DESCRIPTION = "Build per-cell filename_prefix list from x/y labels (row-major), for SaveImage." + + def build( + self, + x_labels: Any, + y_labels: Any, + joiner: Any = "+", + diagonal_single: Any = True, + base_prefix: Any = "", + sanitize: Any = True, + matrix_mode: Any = "full", + ): + x_labels = _first(x_labels, "") or "" + y_labels = _first(y_labels, "") or "" + joiner = str(_first(joiner, "+") or "+") + diagonal_single = bool(_first(diagonal_single, True)) + base_prefix = str(_first(base_prefix, "") or "") + sanitize = bool(_first(sanitize, True)) + matrix_mode = str(_first(matrix_mode, "full") or "full") + + xs = _split_labels(x_labels) + ys = _split_labels(y_labels) + if not xs: + xs = [""] + if not ys: + ys = [""] + + def _one(x: str, y: str) -> str: + xa = (x or "").strip() + ya = (y or "").strip() + if diagonal_single and xa and ya and xa == ya: + name = xa + else: + if xa and ya: + name = f"{xa}{joiner}{ya}" + else: + name = xa or ya or "" + + if base_prefix: + name = f"{base_prefix}{name}" + if sanitize: + name = _sanitize_prefix(name) + return name or "cell" + + out: List[str] = [] + if matrix_mode != "full" and len(xs) == len(ys): + n = len(xs) + if matrix_mode == "upper_triangle": + for r in range(n): + for c in range(r, n): + out.append(_one(xs[c], ys[r])) + elif matrix_mode == "diagonal_only": + for i in range(n): + out.append(_one(xs[i], ys[i])) + else: + for y in ys: + for x in xs: + out.append(_one(x, y)) + else: + for y in ys: + for x in xs: + out.append(_one(x, y)) + + return (out,) +