From b2935e1e49cf6f1b16e11db47f5bf87fd0972819 Mon Sep 17 00:00:00 2001 From: Zar4X Date: Mon, 31 Aug 2026 18:15:38 +0800 Subject: [PATCH] Batch processing nodes with Asset Filter Asset Filter drops stale asset slots kept from previous runs via filename-change detection with per-workflow persisted state; auto-growing 0-based image/video inputs capped at the output counts; positional outputs that emit None for filtered slots so reference consumers (MiniMax H3) run once with only the survivors. --- .github/workflows/publish.yml | 28 + .gitignore | 3 + README.md | 155 ++++ __init__.py | 53 ++ nodes/any_batch_process.py | 674 +++++++++++++++ nodes/image_batch_process.py | 1358 ++++++++++++++++++++++++++++++ nodes/lora_batch_process.py | 335 ++++++++ nodes/mask_batch_process.py | 135 +++ nodes/model3d_batch_process.py | 327 +++++++ nodes/runtime_filename_format.py | 140 +++ nodes/text_batch_process.py | 495 +++++++++++ nodes/video_batch_process.py | 650 ++++++++++++++ pyproject.toml | 14 + requirements.txt | 3 + tests/asset_filter_js.test.mjs | 383 +++++++++ tests/test_asset_filter.py | 377 +++++++++ web/asset_filter.js | 220 +++++ 17 files changed, 5350 insertions(+) create mode 100644 .github/workflows/publish.yml create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes/any_batch_process.py create mode 100644 nodes/image_batch_process.py create mode 100644 nodes/lora_batch_process.py create mode 100644 nodes/mask_batch_process.py create mode 100644 nodes/model3d_batch_process.py create mode 100644 nodes/runtime_filename_format.py create mode 100644 nodes/text_batch_process.py create mode 100644 nodes/video_batch_process.py create mode 100644 pyproject.toml create mode 100644 requirements.txt create mode 100644 tests/asset_filter_js.test.mjs create mode 100644 tests/test_asset_filter.py create mode 100644 web/asset_filter.js diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..0c24043 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,28 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + - master + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'Zar4X' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + submodules: true + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a9342a2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.py[cod] +.tracking diff --git a/README.md b/README.md new file mode 100644 index 0000000..a074715 --- /dev/null +++ b/README.md @@ -0,0 +1,155 @@ +# ComfyUI Batch Process + +A collection of ComfyUI nodes for batch processing text, images, videos, LoRAs, masks, and 3D models. + +## Core Features + +### Text Processing + +- Load TXT files sequentially or randomly +- Filter by filename, prefix, suffix, or file contents +- Select a 1-based file range +- Modify complete strings, prefixes, or suffixes +- Use the selected TXT filename in save filename templates + +### Image Processing + +- Load images sequentially or randomly +- Filter by filename and optionally scan subfolders +- Return one image, a combined image batch, or an image list +- Center-pad different resolutions when building a batch +- Save PNG, JPG, WebP, BMP, TIFF, and GIF files +- Save matching TXT captions and append animation frames + +### Video Processing + +- Save ComfyUI `VIDEO` objects in supported containers and codecs +- Preserve available source FPS +- Save matching TXT captions and workflow metadata + +### LoRA Processing + +- Load LoRAs sequentially or randomly +- Filter by filename, prefix, or suffix +- Apply separate model and CLIP strengths + +### General, Mask, and 3D Processing + +- Group compatible images, latents, values, `MESH`, and `File3D` objects +- Filter out stale asset slots left over from previous runs +- Select, repeat, copy, and combine mask batches +- Save and preview batches of `MESH`, `File3D`, and supported 3D files +- Save optional TXT captions next to 3D models + +## Installation + +Clone the repository into the ComfyUI `custom_nodes` directory: + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/Zar4X/ComfyUI-Batch-Process.git batch-process +pip install -r batch-process/requirements.txt +``` + +Restart ComfyUI after installing or updating the package. + +## Nodes Overview + +### TXT Batch Loader + +Loads UTF-8 TXT files with filename/content filtering, natural sorting, index ranges, incremental loading, seeded random selection, and reset control. + +**Outputs:** text content and filename without the extension. + +### Text Modify Tool + +Searches, replaces, or removes complete text, filename prefixes, or filename suffixes using a custom or automatically detected delimiter. + +### Image Batch Loader + +Loads PNG, JPG/JPEG, WebP, BMP, and GIF files with filtering, natural sorting, recursive search, index ranges, and sequential/random modes. + +**Outputs:** selected image, filename, image count, combined image batch, and image list. + +Images with different resolutions are center-padded with black borders for the combined batch. The image list retains individual dimensions. + +### Image Batch Saver + +Saves images and optional TXT captions with configurable output paths, prefixes, suffixes, delimiters, counters, formats, workflow metadata, frame ranges, and animation frame appending. + +### Video Batch Saver + +Saves ComfyUI video objects and optional TXT captions with configurable paths, naming, containers, codecs, counters, and workflow metadata. + +### LoRA Batch Loader + +Loads `.safetensors`, `.ckpt`, `.pt`, and `.bin` LoRAs with filename filtering, incremental/random modes, and separate model/CLIP strengths. + +**Outputs:** patched model, patched CLIP, and selected filename. + +### Any Batch Group + +Combines compatible values into a batch. Supports image tensors, tensor lists, latents, strings, numbers, `MESH`, and `File3D` objects. Different image sizes are center-padded automatically. + +### Asset Filter + +Removes stale asset slots that ComfyUI keeps from previous runs. Connect reference images (`IMAGE`) and videos (any payload, e.g. `VHS_LoadVideo` output) — the input list grows automatically when the last free slot is connected, with image slots grouped above video slots. Each zone stops growing once it has one input per matching output (9 images, 3 videos), since an input past the last output would have nowhere to go. + +Each slot is traced upstream to its source filename. A slot passes when its filename changed since this node last ran; unchanged slots are filtered out. Special cases: the very first run (no saved state) falls back to file modification time within `freshness_hours` (default 24), a re-run where nothing changed reuses the previous run's selection, and slots that are not file-backed (produced by the current run) always pass. If every connected asset would be filtered, all of them pass through as a fallback and a warning is printed. + +**Outputs:** `image_0`..`image_8`, `video_0`..`video_2`, plus `image_count` and `video_count`. Slot indices are 0-based so they line up one-to-one with consumers like MiniMax H3 (`image_0` -> `ref_image_0`, and so on). + +Outputs mirror the inputs positionally: input `image_N` always leaves on output `image_N`, so a slot's identity never shifts and a preview wired to one slot cannot start showing another slot's image. Each output carries a single asset at its original resolution — nothing is batched, so no black-border padding is ever added. Outputs whose input was filtered out (or is unconnected) return `None`. + +Wire `image_0`..`image_8` into a reference consumer's fixed slots once (for example MiniMax H3's `ref_image_0..8` / `ref_video_0..2`) and leave them wired. Such consumers skip `None` slots, so the consumer runs exactly once with only the surviving assets no matter how many survive, and gaps in the middle are harmless. + +`None` will, however, crash consumers that do not tolerate it — `PreviewImage` and `SaveImage` among them. Feed those from a slot you know is populated, or gate them on `image_count`. + +### 3D Model Batch Saver + +Saves `MESH`, `File3D`, supported file paths, and lists of 3D models. `MESH` is written as GLB; file-based models retain their original format. + +### 3D Model Batch Preview + +Creates ComfyUI 3D previews for individual models or model batches. + +### Mask Nodes + +- **Mask From Batch:** selects a range from a mask batch +- **Mask Repeat Batch:** repeats masks along the batch dimension +- **Mask Batch Copy:** creates two independent mask copies +- **Mask Batch Composite:** concatenates two mask batches + +## Runtime TXT Filename Templates + +ComfyUI normally resolves filename placeholders from input widgets only. This package also supports the runtime `filename` output from TXT Batch Loader. + +Example Save Image prefix: + +```text +D_%TXTBatchLoader.filename%_%EasySeed.seed% +``` + +For `forest_scene.txt` and seed `197152849854180`, the resolved prefix is: + +```text +D_forest_scene_197152849854180 +``` + +The placeholder uses the node's `Node name for S&R` value. Custom names are supported, for example `%Text Source.filename%`. The TXT Batch Loader must participate in the current execution path before the saver runs. + +## Common Index Rules + +- `start_index = 0`: start from the beginning +- `end_index = 0`: continue to the final item +- Positive start/end values are 1-based; the end is inclusive +- Toggle `reset_on_queue` to restart an incremental loader from its first filtered item + +## Requirements + +- ComfyUI +- Pillow 10 or newer +- PyTorch 2.0 or newer +- Codecs required by the selected video format + +Current package version: `1.2.0` diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..58c9c82 --- /dev/null +++ b/__init__.py @@ -0,0 +1,53 @@ +from .nodes.image_batch_process import ImageBatchLoader, ImageBatchSaver +from .nodes.text_batch_process import TXTBatchLoader, TextModifyTool +from .nodes.mask_batch_process import ( + MaskFromBatch, + MaskRepeatBatch, + MaskBatchCopy, + MaskBatchComposite, +) +from .nodes.lora_batch_process import LoraBatchLoader +from .nodes.video_batch_process import VideoBatchSaver # type: ignore +from .nodes.any_batch_process import AnyBatchGroup, AssetFilter +from .nodes.model3d_batch_process import Model3DBatchSaver, Model3DBatchPreview +from .nodes.runtime_filename_format import install_runtime_filename_substitution + +install_runtime_filename_substitution() + +WEB_DIRECTORY = "./web" + +NODE_CLASS_MAPPINGS = { + "ImageBatchLoader": ImageBatchLoader, + "ImageBatchSaver": ImageBatchSaver, + "TextModifyTool": TextModifyTool, + "TXTBatchLoader": TXTBatchLoader, + "LoraBatchLoader": LoraBatchLoader, + "VideoBatchSaver": VideoBatchSaver, + "AnyBatchGroup": AnyBatchGroup, + "AssetFilter": AssetFilter, + "Model3DBatchSaver": Model3DBatchSaver, + "Model3DBatchPreview": Model3DBatchPreview, + "MaskFromBatch": MaskFromBatch, + "MaskRepeatBatch": MaskRepeatBatch, + "MaskBatchCopy": MaskBatchCopy, + "MaskBatchComposite": MaskBatchComposite, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageBatchLoader": "Image Batch Loader", + "ImageBatchSaver": "Image Batch Saver", + "TextModifyTool": "Text Modify Tool", + "TXTBatchLoader": "TXT Batch Loader", + "LoraBatchLoader": "LoRA Batch Loader", + "VideoBatchSaver": "Video Batch Saver", + "AnyBatchGroup": "Any Batch Group", + "AssetFilter": "Asset Filter", + "Model3DBatchSaver": "3D Model Batch Saver", + "Model3DBatchPreview": "3D Model Batch Preview", + "MaskFromBatch": "Mask From Batch", + "MaskRepeatBatch": "Mask Repeat Batch", + "MaskBatchCopy": "Mask Batch Copy", + "MaskBatchComposite": "Mask Batch Composite", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/nodes/any_batch_process.py b/nodes/any_batch_process.py new file mode 100644 index 0000000..ad880a0 --- /dev/null +++ b/nodes/any_batch_process.py @@ -0,0 +1,674 @@ +import json +import os +import re +import time + +import torch +import torch.nn.functional as F + +try: + import comfy.utils +except Exception: + comfy = None + +try: + import folder_paths +except Exception: + folder_paths = None + +try: + from comfy_api.latest._util.geometry_types import MESH as _MESH_TYPE, File3D as _FILE3D_TYPE +except Exception: + _MESH_TYPE = None + _FILE3D_TYPE = None + +try: + from comfy_extras.nodes_save_3d import pack_variable_mesh_batch as _pack_mesh_batch, get_mesh_batch_item as _get_mesh_batch_item +except Exception: + _pack_mesh_batch = None + _get_mesh_batch_item = None + + +class _AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + + +any_type = _AnyType("*") + + +def _to_bhwc(tensor): + if tensor.dim() == 3: + # HWC -> BHWC + return tensor.unsqueeze(0) + if tensor.dim() != 4: + raise ValueError( + f"Only 3D/4D image tensors are supported, got shape {tuple(tensor.shape)}" + ) + # If looks like BCHW, convert to BHWC + if tensor.shape[1] in (1, 3, 4) and tensor.shape[-1] not in (1, 3, 4): + return tensor.movedim(1, -1) + return tensor + + +def _pad_to_size_center(tensor, target_h, target_w): + # tensor: BHWC + b, h, w, c = tensor.shape + if h == target_h and w == target_w: + return tensor + + out = torch.zeros( + (b, target_h, target_w, c), dtype=tensor.dtype, device=tensor.device + ) + top = (target_h - h) // 2 + left = (target_w - w) // 2 + out[:, top : top + h, left : left + w, :] = tensor + return out + + +def _batch_tensors_with_padding(tensors): + converted = [_to_bhwc(t) for t in tensors] + channels = converted[0].shape[-1] + if any(t.shape[-1] != channels for t in converted): + raise ValueError("Channel mismatch in tensor list") + + max_h = max(t.shape[1] for t in converted) + max_w = max(t.shape[2] for t in converted) + padded = [_pad_to_size_center(t, max_h, max_w) for t in converted] + return torch.cat(padded, dim=0) + + +class AnyBatchGroup: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "any_1": (any_type, {}), + "any_2": (any_type, {}), + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("batch",) + FUNCTION = "batch" + CATEGORY = "Batch Process" + + def _common_upscale(self, tensor, width, height): + if comfy is not None: + return comfy.utils.common_upscale( + tensor, width, height, "bilinear", "center" + ) + # Fallback for non-Comfy test environments. + return F.interpolate(tensor, size=(height, width), mode="bilinear") + + def latent_batch(self, any_1, any_2): + samples_out = any_1.copy() + s1 = any_1["samples"] + s2 = any_2["samples"] + + if s1.shape[1:] != s2.shape[1:]: + s2 = self._common_upscale(s2, s1.shape[3], s1.shape[2]) + s = torch.cat((s1, s2), dim=0) + samples_out["samples"] = s + samples_out["batch_index"] = any_1.get( + "batch_index", [x for x in range(0, s1.shape[0])] + ) + any_2.get("batch_index", [x for x in range(0, s2.shape[0])]) + + return samples_out + + def _to_bhwc(self, tensor): + return _to_bhwc(tensor) + + def _pad_to_size_center(self, tensor, target_h, target_w): + return _pad_to_size_center(tensor, target_h, target_w) + + def _is_list_of_tensors(self, value): + return isinstance(value, list) and len(value) > 0 and all( + isinstance(x, torch.Tensor) for x in value + ) + + def _batch_tensors_with_padding(self, tensors): + return _batch_tensors_with_padding(tensors) + + def _is_mesh(self, value): + return _MESH_TYPE is not None and isinstance(value, _MESH_TYPE) + + def _is_file3d(self, value): + return _FILE3D_TYPE is not None and isinstance(value, _FILE3D_TYPE) + + def _contains_file3d(self, value): + if self._is_file3d(value): + return True + if isinstance(value, list): + return any(self._is_file3d(x) for x in value) + return False + + def _contains_mesh(self, value): + if self._is_mesh(value): + return True + if isinstance(value, list): + return any(self._is_mesh(x) for x in value) + return False + + def _to_file3d_list(self, value): + if self._is_file3d(value): + return [value] + if isinstance(value, list): + return [x for x in value if self._is_file3d(x)] + return [] + + def _mesh_to_list(self, mesh): + # Unpack a (possibly variable-length) MESH into a list of single-item MESH objects. + out = [] + b = mesh.vertices.shape[0] + tex = getattr(mesh, "texture", None) + for i in range(b): + v, f, c, u = _get_mesh_batch_item(mesh, i) + t_i = tex[i:i+1] if tex is not None else None + out.append(_MESH_TYPE(v.unsqueeze(0), f.unsqueeze(0), + uvs=u.unsqueeze(0) if u is not None else None, + vertex_colors=c.unsqueeze(0) if c is not None else None, + texture=t_i)) + return out + + def _concat_meshes(self, a, b): + if _pack_mesh_batch is None or _get_mesh_batch_item is None: + raise RuntimeError("AnyBatchGroup: MESH batching requires comfy_extras.nodes_save_3d helpers") + vertices, faces, colors, uvs, textures = [], [], [], [], [] + has_colors = has_uvs = has_tex = False + for m in (a, b): + mb = m.vertices.shape[0] + tex = getattr(m, "texture", None) + for i in range(mb): + v, f, c, u = _get_mesh_batch_item(m, i) + vertices.append(v) + faces.append(f) + colors.append(c) + uvs.append(u) + if tex is not None: + textures.append(tex[i]) + has_tex = True + if c is not None: + has_colors = True + if u is not None: + has_uvs = True + + if has_colors and any(c is None for c in colors): + raise ValueError("AnyBatchGroup: cannot batch meshes — some have vertex_colors and some don't") + if has_uvs and any(u is None for u in uvs): + raise ValueError("AnyBatchGroup: cannot batch meshes — some have uvs and some don't") + if has_tex and len(textures) != len(vertices): + raise ValueError("AnyBatchGroup: cannot batch meshes — texture coverage is partial") + + packed_tex = torch.stack(textures, dim=0) if has_tex else None + return _pack_mesh_batch( + vertices, faces, + colors=colors if has_colors else None, + uvs=uvs if has_uvs else None, + texture=packed_tex, + ) + + def batch(self, any_1, any_2): + # 3D model batching: MESH ↔ MESH, File3D ↔ File3D / list[File3D]. Mixed raises. + a_is_mesh, b_is_mesh = self._contains_mesh(any_1), self._contains_mesh(any_2) + a_is_f3d, b_is_f3d = self._contains_file3d(any_1), self._contains_file3d(any_2) + if (a_is_mesh and b_is_f3d) or (a_is_f3d and b_is_mesh): + raise ValueError( + "AnyBatchGroup: cannot mix MESH and File3D in one batch — convert to a single type first" + ) + if a_is_mesh or b_is_mesh: + if any_1 is None: + return (any_2,) + if any_2 is None: + return (any_1,) + meshes_1 = [any_1] if self._is_mesh(any_1) else [m for m in any_1 if self._is_mesh(m)] + meshes_2 = [any_2] if self._is_mesh(any_2) else [m for m in any_2 if self._is_mesh(m)] + merged = meshes_1[0] + for m in meshes_1[1:] + meshes_2: + merged = self._concat_meshes(merged, m) + return (merged,) + if a_is_f3d or b_is_f3d: + if any_1 is None: + return (self._to_file3d_list(any_2) or any_2,) + if any_2 is None: + return (self._to_file3d_list(any_1) or any_1,) + return (self._to_file3d_list(any_1) + self._to_file3d_list(any_2),) + + # Some any-type chains may pass list-of-tensors; normalize and batch consistently. + if self._is_list_of_tensors(any_1) or self._is_list_of_tensors(any_2): + list_1 = any_1 if self._is_list_of_tensors(any_1) else [any_1] + list_2 = any_2 if self._is_list_of_tensors(any_2) else [any_2] + tensors = [t for t in (list_1 + list_2) if isinstance(t, torch.Tensor)] + if len(tensors) == 0: + return (any_1 if any_2 is None else any_2,) + return (self._batch_tensors_with_padding(tensors),) + + if isinstance(any_1, torch.Tensor) or isinstance(any_2, torch.Tensor): + if any_1 is None: + return (any_2,) + if any_2 is None: + return (any_1,) + + t1 = self._to_bhwc(any_1) + t2 = self._to_bhwc(any_2) + + if t1.shape[-1] != t2.shape[-1]: + raise ValueError( + f"Channel mismatch: {t1.shape[-1]} vs {t2.shape[-1]}" + ) + + # Same size: normal concat. + if t1.shape[1:] == t2.shape[1:]: + return (torch.cat((t1, t2), 0),) + + # Different size: pad to max resolution with black borders, then concat. + max_h = max(t1.shape[1], t2.shape[1]) + max_w = max(t1.shape[2], t2.shape[2]) + t1 = self._pad_to_size_center(t1, max_h, max_w) + t2 = self._pad_to_size_center(t2, max_h, max_w) + return (torch.cat((t1, t2), 0),) + elif isinstance(any_1, (str, float, int)): + if any_2 is None: + return (any_1,) + if isinstance(any_2, tuple): + return (any_2 + (any_1,),) + if isinstance(any_2, list): + return (any_2 + [any_1],) + return ([any_1, any_2],) + elif isinstance(any_2, (str, float, int)): + if any_1 is None: + return (any_2,) + if isinstance(any_1, tuple): + return (any_1 + (any_2,),) + if isinstance(any_1, list): + return (any_1 + [any_2],) + return ([any_2, any_1],) + elif isinstance(any_1, dict) and "samples" in any_1: + if any_2 is None: + return (any_1,) + if isinstance(any_2, dict) and "samples" in any_2: + return (self.latent_batch(any_1, any_2),) + elif isinstance(any_2, dict) and "samples" in any_2: + if any_1 is None: + return (any_2,) + if isinstance(any_1, dict) and "samples" in any_1: + return (self.latent_batch(any_2, any_1),) + else: + if any_1 is None: + return (any_2,) + if any_2 is None: + return (any_1,) + return (any_1 + any_2,) + + +_ASSET_INPUT_RE = re.compile(r"^(image|video)_(\d+)$") + +# Extensions accepted when tracing an input back to a source file. Keeps +# arbitrary STRING widgets (captions etc.) from being mistaken for assets. +_MEDIA_EXTENSIONS = { + ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tiff", ".tif", + ".mp4", ".mov", ".avi", ".webm", ".mkv", ".m4v", ".mpg", ".mpeg", +} + + +class _AssetFilterInputs(dict): + """Flexible optional inputs for AssetFilter. + + ComfyUI asks this dict whether a queued input name is valid and which + type it expects. ``image_N`` slots accept IMAGE, ``video_N`` slots accept + anything (video payloads may be IMAGE frame batches or VIDEO objects, + depending on the loader). The frontend grows/shrinks the input list, so + no fixed input names are declared. + """ + + def __contains__(self, key): + return _ASSET_INPUT_RE.match(key) is not None + + def __getitem__(self, key): + match = _ASSET_INPUT_RE.match(key) + if match is None: + raise KeyError(key) + return ("IMAGE",) if match.group(1) == "image" else ("*",) + + def __iter__(self): + return iter(()) + + def keys(self): + return iter(()) + + +def _has_media_extension(path): + return os.path.splitext(path)[1].lower() in _MEDIA_EXTENSIONS + + +def _resolve_source_file(value): + """Return the absolute path if a widget string points at an existing media file.""" + if not isinstance(value, str) or not value.strip(): + return None + value = value.strip() + + candidates = [] + if folder_paths is not None: + # Handles ComfyUI's annotated widget values, e.g. "clip.mp4 [input]". + try: + candidates.append(folder_paths.get_annotated_filepath(value)) + except Exception: + pass + candidates.append(value) + + roots = [] + if folder_paths is not None: + for getter in ( + folder_paths.get_input_directory, + folder_paths.get_output_directory, + folder_paths.get_temp_directory, + ): + try: + roots.append(getter()) + except Exception: + pass + + for candidate in candidates: + if os.path.isabs(candidate): + if os.path.isfile(candidate) and _has_media_extension(candidate): + return candidate + continue + for root in roots: + full = os.path.join(root, candidate) + if os.path.isfile(full) and _has_media_extension(full): + return full + return None + + +def _trace_source_file(prompt, node_id, input_name, max_depth=8): + """Walk upstream from an input until some node exposes a widget value that + resolves to an existing media file. Returns ``(resolved_path, widget_value)``; + both are None when the input is not file-backed (produced by the run).""" + if not isinstance(prompt, dict) or node_id is None or input_name is None: + return None, None + + node_id = str(node_id) + visited = set() + queue = [(node_id, input_name, 0)] + while queue: + current_id, current_input, depth = queue.pop(0) + if depth > max_depth: + continue + node = prompt.get(current_id) + if not isinstance(node, dict): + continue + link = node.get("inputs", {}).get(current_input) + if not isinstance(link, (list, tuple)) or len(link) < 1: + continue + upstream_id = str(link[0]) + if upstream_id in visited: + continue + visited.add(upstream_id) + upstream = prompt.get(upstream_id) + if not isinstance(upstream, dict): + continue + inputs = upstream.get("inputs", {}) + # Check this node's own widget values before going deeper. + for value in inputs.values(): + if isinstance(value, str): + resolved = _resolve_source_file(value) + if resolved is not None: + return resolved, value + for key, value in inputs.items(): + if isinstance(value, (list, tuple)): + queue.append((upstream_id, key, depth + 1)) + return None, None + + +_STATE_FILENAME = "zar4x_asset_filter_state.json" +_STATE_MAX_ENTRIES = 64 +_RUNTIME_ASSET = "__runtime__" +_memory_state = {} + + +def _state_path(): + if folder_paths is None: + return None + for getter in (folder_paths.get_user_directory, folder_paths.get_temp_directory): + try: + return os.path.join(getter(), _STATE_FILENAME) + except Exception: + continue + return None + + +def _load_state(): + path = _state_path() + if path is None: + return _memory_state + try: + with open(path, "r", encoding="utf-8") as handle: + loaded = json.load(handle) + return loaded if isinstance(loaded, dict) else {} + except Exception: + return {} + + +def _save_state(state): + if not isinstance(state, dict): + return + if len(state) > _STATE_MAX_ENTRIES: + # Drop oldest entries; dict order is insertion order. + for key in list(state)[: len(state) - _STATE_MAX_ENTRIES]: + state.pop(key, None) + path = _state_path() + if path is None: + # state may be _memory_state itself; only copy when it is not. + if state is not _memory_state: + _memory_state.clear() + _memory_state.update(state) + return + try: + with open(path, "w", encoding="utf-8") as handle: + json.dump(state, handle, ensure_ascii=False) + except Exception as error: + print(f"[AssetFilter] warning: could not persist state to {path}: {error}") + + +def _state_key(extra_pnginfo, node_id): + # ComfyUI sends extra_pnginfo as {"workflow": {...}}; the stable workflow + # uuid lives inside that nested dict. + workflow_id = None + if isinstance(extra_pnginfo, dict): + workflow = extra_pnginfo.get("workflow") + if isinstance(workflow, dict): + workflow_id = workflow.get("id") + if workflow_id is None: + workflow_id = extra_pnginfo.get("id") + return f"{workflow_id}:{node_id}" + + +_MAX_IMAGE_OUTPUTS = 9 +_MAX_VIDEO_OUTPUTS = 3 + + +class AssetFilter: + """Drops stale asset slots left over from previous runs. + + Connect reference images/videos (e.g. LoadImage, VHS_LoadVideo) to the + growing input list. A slot is considered new when its source filename + changed since this node last ran; unchanged slots are filtered out. + On the very first run (no saved state) the file modification time within + `freshness_hours` decides instead, and re-running without any change + reuses the previous run's selection. + + Outputs mirror the inputs positionally: input image_N always leaves on + output image_N, so a slot's identity never shifts and previews cannot swap + images between runs. Outputs whose input was filtered (or is unconnected) + return None. Reference consumers such as MiniMax H3 skip None slots, so + wiring image_0..8 into ref_image_0..8 once and leaving it wired works + regardless of how many assets survive — the consumer still runs exactly + once with only the survivors, and gaps in the middle are harmless. + + Note that None will crash consumers that do not tolerate it (PreviewImage, + SaveImage). Feed those from a specific slot you know is populated, or gate + them on image_count. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "freshness_hours": ( + "FLOAT", + {"default": 24.0, "min": 0.1, "max": 720.0, "step": 0.5}, + ), + }, + "optional": _AssetFilterInputs(), + "hidden": { + "prompt": "PROMPT", + "node_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + RETURN_TYPES = ( + tuple(["IMAGE"] * _MAX_IMAGE_OUTPUTS) + + tuple([any_type] * _MAX_VIDEO_OUTPUTS) + + ("INT", "INT") + ) + RETURN_NAMES = ( + tuple(f"image_{i}" for i in range(_MAX_IMAGE_OUTPUTS)) + + tuple(f"video_{i}" for i in range(_MAX_VIDEO_OUTPUTS)) + + ("image_count", "video_count") + ) + FUNCTION = "filter_assets" + CATEGORY = "Batch Process" + + def filter_assets(self, freshness_hours=24.0, prompt=None, node_id=None, + extra_pnginfo=None, **kwargs): + assets = [] + for name, value in kwargs.items(): + match = _ASSET_INPUT_RE.match(name) + if match is None: + continue + assets.append((match.group(1), int(match.group(2)), value)) + assets.sort(key=lambda item: (item[0], item[1])) + + window = float(freshness_hours) * 3600.0 + now = time.time() + + state = _load_state() + key = _state_key(extra_pnginfo, node_id) + prev = state.get(key) + if not isinstance(prev, dict): + prev = None + + entries = [] + for kind, index, value in assets: + slot = f"{kind}_{index}" + resolved, widget = _trace_source_file(prompt, node_id, slot) + # The widget value is the filename the user sees; use it as the + # change-detection identity. + name = widget if isinstance(widget, str) and widget.strip() else ( + resolved or _RUNTIME_ASSET + ) + entries.append({ + "slot": slot, + "kind": kind, + "index": index, + "value": value, + "resolved": resolved, + "name": name, + "kept": None, + "reason": None, + }) + + all_unchanged = prev is not None and set(prev.keys()) == { + e["slot"] for e in entries + } and all(prev[e["slot"]].get("name") == e["name"] for e in entries) + + for entry in entries: + slot, resolved, name = entry["slot"], entry["resolved"], entry["name"] + runtime = resolved is None + if prev is None: + # First run with no saved state: fall back to mtime. + kept = runtime + if not kept: + try: + kept = (now - os.path.getmtime(resolved)) <= window + except OSError: + kept = True + reason = "bootstrap-fresh" if kept else "bootstrap-stale" + elif all_unchanged: + # Pure re-run: reuse the previous selection. + kept = bool(prev[slot].get("kept")) + reason = "reused" + else: + prev_name = prev.get(slot, {}).get("name") + kept = runtime or name != prev_name + reason = "changed" if kept else "unchanged" + entry["kept"] = kept + entry["reason"] = reason + + if entries and not any(e["kept"] for e in entries): + print( + "[AssetFilter] warning: every connected asset looks stale — " + "passing all of them through as a fallback." + ) + for entry in entries: + entry["kept"] = True + + images = { + e["index"]: e["value"] + for e in entries + if e["kind"] == "image" and e["kept"] + } + videos = { + e["index"]: e["value"] + for e in entries + if e["kind"] == "video" and e["kept"] + } + + for entry in entries: + print( + f"[AssetFilter] {entry['slot']}: {entry['reason']} " + f"({'keep' if entry['kept'] else 'drop'}) " + f"<- {entry['resolved'] or entry['name']}" + ) + + new_entry = { + e["slot"]: {"name": e["name"], "kept": e["kept"]} for e in entries + } + if prev is None or not all_unchanged: + state[key] = new_entry + _save_state(state) + + total_images = sum(1 for e in entries if e["kind"] == "image") + total_videos = sum(1 for e in entries if e["kind"] == "video") + print( + f"[AssetFilter] kept {len(images)}/{total_images} images, " + f"{len(videos)}/{total_videos} videos" + ) + self._warn_out_of_range(images, _MAX_IMAGE_OUTPUTS, "image") + self._warn_out_of_range(videos, _MAX_VIDEO_OUTPUTS, "video") + # Mirror inputs positionally: input image_N leaves on output image_N so + # a slot's identity never shifts between runs. Filtered and unconnected + # slots emit None; reference consumers skip None, and a gap in the + # middle is harmless to them. An ExecutionBlocker would instead prune + # the whole downstream node, not just that one slot. + return ( + self._place(images, _MAX_IMAGE_OUTPUTS) + + self._place(videos, _MAX_VIDEO_OUTPUTS) + + (len(images), len(videos)) + ) + + @staticmethod + def _place(by_index, count): + """Put each survivor on the output matching its 0-based input index.""" + return tuple(by_index.get(i) for i in range(count)) + + @staticmethod + def _warn_out_of_range(by_index, count, kind): + extra = sorted(i for i in by_index if i >= count) + if extra: + slots = ", ".join(f"{kind}_{i}" for i in extra) + print( + f"[AssetFilter] warning: {slots} survived but only " + f"{count} {kind} outputs exist — dropping them." + ) diff --git a/nodes/image_batch_process.py b/nodes/image_batch_process.py new file mode 100644 index 0000000..9c7db53 --- /dev/null +++ b/nodes/image_batch_process.py @@ -0,0 +1,1358 @@ +import os +import random +import json +import hashlib +import re +import glob + +import numpy as np +import torch +from PIL import Image, ImageOps, ImageSequence +from PIL.PngImagePlugin import PngInfo + +import folder_paths +import node_helpers +from server import PromptServer + + +class ImageBatchLoader: + RETURN_TYPES = ("IMAGE", "STRING", "STRING", "IMAGE", "IMAGE") + RETURN_NAMES = ("image", "filename", "image_count", "image_batch", "image_list") + OUTPUT_IS_LIST = (False, False, False, False, True) + FUNCTION = "load_batch_images" + CATEGORY = "Batch Process" + + SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"} + + def __init__(self): + self.image_states = {} + self.current_directory = "" + self.images = [] + self.search_states = {} + self._last_scan_key = None + self._all_image_paths = [] + self._last_reset_on_queue = {} + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "directory": ("STRING",), + "search_title": ("STRING", {"default": ""}), + "delimiter": ("STRING", {"default": ""}), + "mode": ( + ["single_image", "incremental_image", "random"], + {"default": "incremental_image"}, + ), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "filename_option": ( + [ + "filename", + "prefix", + "suffix", + "prefix & suffix", + "prefix nor suffix", + ], + ), + "image_batch": ( + "BOOLEAN", + { + "label_on": "yes", + "label_off": "no", + "default": False, + "defaultInput": False, + "tooltip": "Load all matching/filtered images in the directory as a single batched IMAGE tensor (B,H,W,C). Images of differing sizes are center-padded to the max resolution so they can be concatenated.", + }, + ), + "image_list": ( + "BOOLEAN", + { + "label_on": "yes", + "label_off": "no", + "default": False, + "defaultInput": False, + }, + ), + "subfolder": ( + "BOOLEAN", + { + "label_on": "yes", + "label_off": "no", + "default": False, + "defaultInput": False, + }, + ), + "start_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "tooltip": "Start index (1-based). Use 0 to start from beginning, or set 1 for first image, 2 for second image, etc.", + }, + ), + "end_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "tooltip": "End index (inclusive, 1-based). Use 0 to include all remaining images, or set >0 to limit range (1=first, 2=second, etc.).", + }, + ), + }, + "optional": { + "reset_on_queue": ( + "INT", + { + "default": 1, + "min": 0, + "max": 1, + "step": 1, + "tooltip": "Any change to this value (0↔1) resets the read index to start from the first image again. Keep unchanged during a queue run to advance normally.", + }, + ), + }, + "hidden": {"node_id": "UNIQUE_ID"}, + } + + def set_directory( + self, + directory, + filename_option="filename", + search_title="", + delimiter="", + subfolder=False, + ): + scan_key = (directory, filename_option, search_title, delimiter, subfolder) + if scan_key != self._last_scan_key: + if not os.path.isdir(directory): + raise ValueError( + f"The provided path '{directory}' is not a valid directory." + ) + + # Fast list of all images (no PIL verification) + all_image_paths = self.list_images(directory, subfolder) + self._all_image_paths = all_image_paths + + # Extract just the filenames for filtering + all_images = [os.path.basename(path) for path in all_image_paths] + + filtered_images = self.filter_images( + directory, all_images, filename_option, search_title, delimiter + ) + + # Filter full paths in O(n) by basename membership + allowed_names = set(filtered_images) + self.images = [ + p for p in all_image_paths if os.path.basename(p) in allowed_names + ] + self.images = sorted( + self.images, + key=lambda p: [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", os.path.basename(p)) + ], + ) + + self.current_directory = directory + self._last_scan_key = scan_key + + if scan_key not in self.search_states: + self.search_states[scan_key] = 0 + + if not self.images: + print("No matching image files found in the provided directory.") + + def load_images(self, directory): + if not os.path.isdir(directory): + raise ValueError(f"Invalid directory: {directory}") + + all_images = [ + f + for f in os.listdir(directory) + if any(f.endswith(ext) for ext in self.SUPPORTED_EXTENSIONS) + ] + paths = [os.path.join(directory, f) for f in all_images] + return sorted( + paths, + key=lambda p: [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", os.path.basename(p)) + ], + ) + + def filter_images(self, directory, files, filename_option, search_title, delimiter): + def get_prefix(filename): + if delimiter: + return filename.split(delimiter)[0] + else: + return re.split(r"[^a-zA-Z0-9]", filename)[0] + + def get_suffix(filename): + name_without_ext = os.path.splitext(filename)[0] + if delimiter: + return name_without_ext.split(delimiter)[-1] + else: + return re.split(r"[^a-zA-Z0-9]", name_without_ext)[-1] + + filtered_files = files + + if search_title: + if filename_option == "filename": + filtered_files = [f for f in filtered_files if search_title in f] + elif filename_option == "prefix": + search_prefix = get_prefix(search_title) + filtered_files = [ + f for f in filtered_files if get_prefix(f) == search_prefix + ] + elif filename_option == "suffix": + search_suffix = get_suffix(search_title) + filtered_files = [ + f for f in filtered_files if get_suffix(f) == search_suffix + ] + elif filename_option == "prefix & suffix": + search_prefix = get_prefix(search_title) + search_suffix = get_suffix(search_title) + filtered_files = [ + f + for f in filtered_files + if get_prefix(f) == search_prefix or get_suffix(f) == search_suffix + ] + elif filename_option == "prefix nor suffix": + search_prefix = get_prefix(search_title) + search_suffix = get_suffix(search_title) + filtered_files = [ + f + for f in filtered_files + if get_prefix(f) != search_prefix and get_suffix(f) != search_suffix + ] + + return filtered_files + + @classmethod + def list_images(cls, path: str, subfolder: bool = False, verify: bool = False): + images = [] + + if os.path.isfile(path): + files = [path] + else: + if subfolder: + files = [] + for root, _, filenames in os.walk(path): + for name in filenames: + files.append(os.path.join(root, name)) + else: + try: + files = [ + entry.path for entry in os.scandir(path) if entry.is_file() + ] + except FileNotFoundError: + files = [] + + valid_extensions = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"} + + candidate_files = [ + f for f in files if os.path.splitext(f)[1].lower() in valid_extensions + ] + + if not verify: + return sorted( + candidate_files, + key=lambda p: [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", os.path.basename(p)) + ], + ) + + for filename in candidate_files: + try: + with Image.open(filename) as img: + img.verify() + images.append(filename) + except Exception: + continue + + return sorted( + images, + key=lambda p: [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", os.path.basename(p)) + ], + ) + + def load_batch_images( + self, + directory, + search_title="", + delimiter="", + mode="incremental_image", + seed=0, + filename_option="filename", + image_batch=False, + image_list=False, + subfolder=False, + start_index=0, + end_index=-1, + reset_on_queue=1, + node_id=None, + ): + # Ensure directory is scanned once and cached + self.set_directory( + directory, filename_option, search_title, delimiter, subfolder + ) + + # Use cached listing for count and names + all_images_in_dir = self._all_image_paths + image_count = str(len(all_images_in_dir)) + + # Empty batch placeholder (1x64x64x3 black) used when no batch is requested/available + empty_batch = torch.zeros(1, 64, 64, 3) + + # Determine which image list to use for index filtering + images_to_filter = self.images if self.images else all_images_in_dir + + # Apply index filtering + filtered_images = images_to_filter + if filtered_images: + # Determine actual start and end indices + # start_index is 1-indexed: 1 = first image, 2 = second image, etc. + # start_index=0 means start from beginning (index 0) + # end_index=0 means end at last image (no limit, use all remaining) + # Only non-zero values apply as limits + # Convert 1-indexed to 0-indexed: subtract 1 if > 0 + actual_start = (start_index - 1) if start_index > 0 else 0 + actual_end = ( + (end_index - 1) if end_index > 0 else (len(filtered_images) - 1) + ) + + # Validate indices + if actual_start < 0: + actual_start = 0 + if actual_start >= len(filtered_images): + actual_start = ( + len(filtered_images) - 1 if len(filtered_images) > 0 else 0 + ) + + if actual_end < 0: + actual_end = 0 + if actual_end >= len(filtered_images): + actual_end = len(filtered_images) - 1 + if actual_end < actual_start: + actual_end = actual_start + + # Only apply slicing if we have valid indices and images + if len(filtered_images) > 0 and actual_start <= actual_end: + # Slice the images list (actual_end+1 because slice is exclusive on end) + filtered_images = filtered_images[actual_start : actual_end + 1] + else: + filtered_images = [] + + # Load all filtered images once if either batch or list output is requested. + all_loaded_images = [] + if image_batch or image_list: + if filtered_images: + all_loaded_images = self.load_all_images( + path=directory, + subfolder=subfolder, + node_id=node_id, + filepaths=filtered_images, + ) + + # Build the batched IMAGE tensor (B,H,W,C) for the image_batch output. + batch_out = empty_batch + if image_batch and all_loaded_images: + batch_out = self.build_image_batch(all_loaded_images) + + # The image_list output is the per-image list (each entry a (1,H,W,C) tensor). + list_out = all_loaded_images if image_list else [] + + if not filtered_images: + return ( + empty_batch, + "no_images_found", + image_count, + empty_batch, + [], + ) + + # First "image" output: a single representative image (per mode). + if image_list and all_loaded_images: + # In list mode keep the legacy behavior: first image as the single output. + single_image = all_loaded_images[0] + single_filename = os.path.splitext( + os.path.basename(filtered_images[0]) + )[0] + return ( + single_image, + single_filename, + image_count, + batch_out, + list_out, + ) + + search_key = ( + directory, + filename_option, + search_title, + delimiter, + subfolder, + start_index, + end_index, + ) + + if self._last_reset_on_queue.get(search_key) != reset_on_queue: + self.search_states[search_key] = 0 + self._last_reset_on_queue[search_key] = reset_on_queue + + if mode == "single_image": + image, filename = self.load_image_by_index(search_key, filtered_images) + return image, filename, image_count, batch_out, list_out + elif mode == "incremental_image": + image, filename = self.load_image_by_index(search_key, filtered_images) + return image, filename, image_count, batch_out, list_out + elif mode == "random": + random.seed(seed) + rnd_index = random.randint(0, len(filtered_images) - 1) + image, filename = self.load_image_by_path(filtered_images[rnd_index]) + return image, filename, image_count, batch_out, list_out + else: + raise ValueError(f"Unknown mode: {mode}") + + def build_image_batch(self, image_tensors): + """Concatenate a list of (1,H,W,C) tensors into a single (B,H,W,C) batch. + + Images of differing sizes are center-padded to the max H/W with black + borders so they can be concatenated along the batch dimension. + """ + if not image_tensors: + return torch.zeros(1, 64, 64, 3) + if len(image_tensors) == 1: + return image_tensors[0] + + max_h = max(t.shape[1] for t in image_tensors) + max_w = max(t.shape[2] for t in image_tensors) + channels = image_tensors[0].shape[3] + + padded = [] + for t in image_tensors: + if t.shape[3] != channels: + print( + f"Skipping image with mismatched channels: {t.shape[3]} vs {channels}" + ) + continue + b, h, w, c = t.shape + if h == max_h and w == max_w: + padded.append(t) + continue + out = torch.zeros((b, max_h, max_w, c), dtype=t.dtype, device=t.device) + top = (max_h - h) // 2 + left = (max_w - w) // 2 + out[:, top : top + h, left : left + w, :] = t + padded.append(out) + + return torch.cat(padded, dim=0) + + def load_all_images( + self, + path: str = None, + subfolder: bool = False, + node_id: str = None, + filepaths: list = None, + ): + """Load all images for the image_list output""" + images = [] + if filepaths is None: + # Fallback to listing if not provided + filepaths = self.list_images(path, subfolder) + + for index, image_path in enumerate(filepaths): + try: + img = node_helpers.pillow(Image.open, image_path) + img = node_helpers.pillow(ImageOps.exif_transpose, img) + + if img.mode == "I": + img = img.point(lambda i: i * (1 / 255)) + img = img.convert("RGB") + + image_np = np.array(img).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_np)[None, ...] + images.append(image_tensor) + + if node_id: + PromptServer.instance.send_sync( + "progress", + {"node": node_id, "max": len(filepaths), "value": index}, + ) + except Exception as e: + print(f"Error loading image {image_path}: {str(e)}") + continue + + return images + + def load_image_by_index(self, search_key, filtered_images): + if not filtered_images: + print("No images loaded.") + return None, None + + if search_key not in self.search_states: + self.search_states[search_key] = 0 + + current_index = self.search_states[search_key] + if current_index >= len(filtered_images): + current_index = 0 + + file_path = filtered_images[current_index] + self.search_states[search_key] = (current_index + 1) % len(filtered_images) + + return self.load_image_by_path(file_path) + + def load_image_by_path(self, path): + try: + image = Image.open(path) + image = ImageOps.exif_transpose(image).convert("RGB") + filename = os.path.basename(path) + # 去除文件扩展名 + filename = os.path.splitext(filename)[0] + return self.pil2tensor(image), filename + except Exception as e: + print(f"Error loading image: {str(e)}") + return (torch.zeros(1, 64, 64, 3)), "error" + + def pil2tensor(self, image): + image_np = np.array(image).astype(np.float32) / 255.0 + if len(image_np.shape) == 2: + image_np = np.expand_dims(image_np, axis=-1) + image_np = np.expand_dims(image_np, axis=0) + return torch.from_numpy(image_np) + + @classmethod + def IS_CHANGED(cls, directory, **kwargs): + if not os.path.exists(directory): + return "" + try: + loader = cls() + paths = loader.load_images(directory) + return hashlib.sha256(",".join(paths).encode()).hexdigest() + except Exception as e: + print(f"Error checking for changes: {str(e)}") + return "" + + +class ImageBatchSaver: + INPUT_IS_LIST = True + FUNCTION = "save" + CATEGORY = "Batch Process" + RETURN_TYPES = () + RETURN_NAMES = () + OUTPUT_NODE = True + + ALLOWED_EXT = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".gif"] + + @classmethod + def IS_CHANGED(cls, **kwargs): + # Always execute - return a unique value each time + import time + + return time.time() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "images": ("IMAGE",), + "contents": ("STRING", {"forceInput": True}), + "output_path": ("STRING", {"default": ""}), + "filename_prefix": ("STRING", {"default": "IMG"}), + "filename_delimiter": ("STRING", {"default": "_"}), + "filename_suffix": ("STRING", {"default": ""}), + "extension": ( + ["png", "jpg", "jpeg", "webp", "bmp", "tiff", "gif"], + {"default": "png"}, + ), + "filename_number_padding": ( + "INT", + {"default": 4, "min": 1, "max": 9, "step": 1}, + ), + "filename_number": ( + ["off", "start", "end"], + {"default": "end"}, + ), + "embeded_workflow": ("BOOLEAN", {"default": True}), + "append_frames": ( + "BOOLEAN", + { + "label_on": "true", + "label_off": "false", + "default": False, + "defaultInput": False, + "tooltip": "Append new frames to existing file with same base name. Useful for accumulating frames across loop iterations to reduce memory usage.", + }, + ), + "start_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "tooltip": "Start index (1-based) for frames to save. Use 0 to save all frames, or set 1 for first frame, 2 for second frame, etc.", + }, + ), + "end_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "tooltip": "End index (inclusive, 1-based) for frames to save. Use 0 to save all remaining frames, or set >0 to limit range.", + }, + ), + }, + "hidden": { + "node_id": "UNIQUE_ID", + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + def save( + self, + images=None, + contents=None, + output_path="", + filename_prefix="IMG", + filename_delimiter="_", + filename_suffix="", + extension="png", + filename_number_padding=4, + filename_number="end", + append_frames=False, + start_index=0, + end_index=0, + embeded_workflow=True, + node_id=None, + prompt=None, + extra_pnginfo=None, + ): + # 输入验证和处理 + extension = self._get_first_or_default(extension, "png") + filename_prefix = self._get_first_or_default(filename_prefix, "IMG") + output_path = self._get_first_or_default(output_path, "") + filename_number_padding = self._get_first_or_default(filename_number_padding, 4) + filename_delimiter = self._get_first_or_default(filename_delimiter, "_") + filename_suffix = self._get_first_or_default(filename_suffix, "") + filename_suffix = filename_suffix.strip("'[]") + filename_number = self._get_first_or_default(filename_number, "end") + append_frames = self._get_first_or_default(append_frames, False) + start_index = self._get_first_or_default(start_index, 0) + end_index = self._get_first_or_default(end_index, 0) + embeded_workflow = self._get_first_or_default(embeded_workflow, True) + + # 解析数字位置选项 + counter_start = filename_number in ["start"] + counter_end = filename_number in ["end"] + + if extension not in ["png", "jpg", "jpeg", "webp", "bmp", "tiff", "gif"]: + raise ValueError(f"Invalid extension: {extension}") + + if filename_number_padding < 1: + raise ValueError( + f"filename_number_padding must be at least 1, got {filename_number_padding}" + ) + + # 处理 images 输入 + if images is not None: + processed_images = [] + for img in images if isinstance(images, list) else [images]: + if ( + isinstance(img, torch.Tensor) + and img.dim() == 4 + and img.shape[0] > 1 + ): + batch_size = img.shape[0] + for i in range(batch_size): + processed_images.append(img[i].unsqueeze(0)) + else: + processed_images.append(img) + images = processed_images + image_count = len(images) + else: + images = [] + image_count = len(contents) if contents is not None else 1 + + # 处理 output_path + output_path = self._normalize_input(output_path, image_count) + + # 处理 filename_prefix + filename_prefix = self._normalize_input(filename_prefix, image_count) + + # 处理文本输入 + contents = ( + self._normalize_input(contents, image_count) + if contents is not None + else None + ) + + # 准备输出目录 + output_dir = folder_paths.get_output_directory() + + # 批量保存 + saved_files = [] + for idx, (prefix, path) in enumerate(zip(filename_prefix, output_path)): + try: + final_output_path = self._get_output_path(output_dir, path) + os.makedirs(final_output_path, exist_ok=True) + + original_filename = os.path.splitext(os.path.basename(prefix))[0] + base_filename = self._generate_filename( + prefix=original_filename, + suffix=filename_suffix, + padding=filename_number_padding, + counter_start=counter_start, + counter_end=counter_end, + delimiter=filename_delimiter, + final_output_path=final_output_path, + ) + + # 保存图片 + if images and idx < len(images): + image_to_save = images[idx] + + # Apply start_index and end_index filtering (1-based) + if ( + isinstance(image_to_save, torch.Tensor) + and image_to_save.dim() == 4 + ): + # Multi-frame tensor: (batch, height, width, channels) + num_frames = image_to_save.shape[0] + actual_start = (start_index - 1) if start_index > 0 else 0 + actual_end = ( + (end_index - 1) if end_index > 0 else (num_frames - 1) + ) + + # Validate and clamp indices + actual_start = max(0, min(actual_start, num_frames - 1)) + actual_end = max(actual_start, min(actual_end, num_frames - 1)) + + if actual_start <= actual_end and num_frames > 0: + image_to_save = image_to_save[actual_start : actual_end + 1] + else: + image_to_save = image_to_save[0:1] # At least one frame + + # Handle append mode + if append_frames: + # Auto-convert to GIF if current format doesn't support multiple frames + original_extension = extension + if extension.lower() not in ["gif", "tiff", "webp"]: + # Automatically convert to GIF for append_frames (best for animations) + extension = "gif" + print( + f"Info: {original_extension} format doesn't support multiple frames. " + f"Auto-converting to GIF format for append_frames feature." + ) + + # Find existing file with same base name (try both original and converted extensions) + base_name_no_counter = self._get_base_name_without_counter( + original_filename, filename_suffix, filename_delimiter + ) + existing_path = self._find_existing_file( + final_output_path, + base_name_no_counter, + extension, + counter_start, + counter_end, + filename_delimiter, + ) + + # If not found with converted extension, try original extension + if not existing_path and original_extension != extension: + existing_path = self._find_existing_file( + final_output_path, + base_name_no_counter, + original_extension, + counter_start, + counter_end, + filename_delimiter, + ) + # If found with original extension, we'll convert it to new format + if existing_path: + print( + f"Info: Found existing file with {original_extension} extension. " + f"Will convert to {extension} format for multi-frame support." + ) + + if existing_path and os.path.exists(existing_path): + # Load existing frames and append new ones + existing_frames = self._load_image_frames(existing_path) + if existing_frames is not None: + # Concatenate: existing frames + new frames + if image_to_save.dim() == 3: + image_to_save = image_to_save.unsqueeze(0) + if existing_frames.dim() == 3: + existing_frames = existing_frames.unsqueeze(0) + image_to_save = torch.cat( + [existing_frames, image_to_save], dim=0 + ) + + # Update path extension if format was converted + if original_extension != extension: + # Change extension in path + base_path = os.path.splitext(existing_path)[0] + full_path = f"{base_path}.{extension}" + else: + full_path = existing_path + + num_new_frames = ( + image_to_save.shape[0] - existing_frames.shape[0] + ) + print( + f"Appending {num_new_frames} frame(s) to file: {full_path} " + f"(total: {image_to_save.shape[0]} frames)" + ) + else: + # Failed to load, create new file with correct extension + full_path = os.path.join( + final_output_path, f"{base_filename}.{extension}" + ) + else: + # No existing file, create new one with correct extension + full_path = os.path.join( + final_output_path, f"{base_filename}.{extension}" + ) + else: + # Normal mode: always create new file + full_path = os.path.join( + final_output_path, f"{base_filename}.{extension}" + ) + + self._save_image_tensor( + image_to_save, + full_path, + embed_workflow=embeded_workflow, + prompt=prompt, + extra_pnginfo=extra_pnginfo, + extension=extension, + ) + saved_files.append(full_path) + print(f"Saved image: {full_path}") + elif not images: + # 如果没有图像，仍然生成文件路径（可能用于其他用途） + full_path = os.path.join( + final_output_path, f"{base_filename}.{extension}" + ) + saved_files.append(full_path) + print(f"Generated path: {full_path}") + + # 保存文本 + if contents and idx < len(contents): + content_path = os.path.join( + final_output_path, f"{base_filename}.txt" + ) + with open(content_path, "w", encoding="utf-8") as f: + f.write(str(contents[idx]).strip()) + saved_files.append(content_path) + print(f"Saved content: {content_path}") + + except Exception as e: + print(f"Error saving file {idx+1}: {str(e)}") + + self._update_progress(node_id, idx + 1, image_count) + + return () + + def _save_image_tensor( + self, tensor, path, embed_workflow, prompt, extra_pnginfo, extension + ): + try: + # Ensure tensor is on CPU + if tensor.is_cuda: + tensor = tensor.cpu() + + # Handle multi-frame tensors (4D: batch, height, width, channels) + if tensor.dim() == 4 and tensor.shape[0] > 1: + self._save_multi_frame_tensor( + tensor, path, embed_workflow, prompt, extra_pnginfo, extension + ) + return + + # Handle single frame - normalize to 3D (H, W, C) + if tensor.dim() == 4: + tensor = tensor.squeeze(0) + + # Convert to numpy array + if tensor.dtype == torch.uint8: + # Already in uint8 format + img_array = tensor.numpy() + else: + # Convert from float [0,1] to uint8 [0,255] + img_array = np.clip(255.0 * tensor.numpy(), 0, 255).astype(np.uint8) + + # Handle different tensor shapes + if img_array.ndim == 3: + # 3D array: could be (C, H, W) or (H, W, C) + if img_array.shape[0] == 3 or img_array.shape[0] == 1: + # (C, H, W) format - transpose to (H, W, C) + img_array = np.transpose(img_array, (1, 2, 0)) + elif img_array.shape[2] == 3 or img_array.shape[2] == 1: + # Already in (H, W, C) format + pass + else: + # Unusual shape - try to reshape + # If it's a very wide/narrow image, handle it + if img_array.shape[0] == 1 and img_array.shape[1] == 1: + # Shape like (1, 1, width) - reshape to (1, width, 1) then add channel + img_array = img_array.reshape(img_array.shape[2], 1, 1) + img_array = np.transpose(img_array, (1, 0, 2)) + elif img_array.shape[1] == 1: + # Shape like (height, 1, width) - transpose + img_array = np.transpose(img_array, (0, 2, 1)) + + # Ensure we have valid image dimensions + if img_array.ndim == 2: + # Grayscale 2D image - add channel dimension + img_array = np.expand_dims(img_array, axis=2) + elif img_array.ndim == 1: + # 1D array - try to reshape to 2D + # Assume it's a square or try to infer dimensions + size = int(np.sqrt(len(img_array))) + if size * size == len(img_array): + img_array = img_array.reshape(size, size, 1) + else: + raise ValueError( + f"Cannot reshape 1D array of length {len(img_array)} to image" + ) + + # Ensure valid channel count (1 for grayscale, 3 for RGB) + if img_array.shape[2] == 1: + # Grayscale - keep as is, PIL will handle it + img_array = img_array.squeeze( + 2 + ) # Remove channel dimension for grayscale + elif img_array.shape[2] > 3: + # More than 3 channels - take first 3 + img_array = img_array[:, :, :3] + elif img_array.shape[2] == 2: + # 2 channels - duplicate to make 3 + img_array = np.concatenate([img_array, img_array[:, :, :1]], axis=2) + + # Create PIL Image + if img_array.ndim == 2: + # Grayscale + img = Image.fromarray(img_array, mode="L") + elif img_array.ndim == 3 and img_array.shape[2] == 3: + # RGB + img = Image.fromarray(img_array, mode="RGB") + else: + # Fallback - convert to RGB + if img_array.ndim == 2: + img = Image.fromarray(img_array, mode="L").convert("RGB") + else: + img = Image.fromarray(img_array[:, :, :3], mode="RGB") + + except Exception as e: + print(f"Error processing tensor with shape {tensor.shape}: {str(e)}") + # Fallback: try basic conversion + try: + tensor_np = tensor.cpu().numpy() + if tensor_np.dtype != np.uint8: + tensor_np = np.clip(255.0 * tensor_np, 0, 255).astype(np.uint8) + + # Try to reshape to valid image + while tensor_np.ndim < 2: + tensor_np = np.expand_dims(tensor_np, axis=0) + while tensor_np.ndim > 3: + tensor_np = tensor_np.squeeze(0) + + if tensor_np.ndim == 2: + img = Image.fromarray(tensor_np, mode="L").convert("RGB") + elif tensor_np.ndim == 3: + if tensor_np.shape[0] == 3 or tensor_np.shape[0] == 1: + tensor_np = np.transpose(tensor_np, (1, 2, 0)) + if tensor_np.shape[2] == 1: + tensor_np = tensor_np.squeeze(2) + img = Image.fromarray(tensor_np, mode="L").convert("RGB") + else: + img = Image.fromarray(tensor_np[:, :, :3], mode="RGB") + else: + raise ValueError(f"Cannot convert tensor with shape {tensor.shape}") + except Exception as e2: + raise ValueError( + f"Failed to save image: {str(e)}, fallback also failed: {str(e2)}" + ) + + if embed_workflow: + if extension.lower() == "webp": + exif_data = img.getexif() + if prompt is not None: + prompt_str = json.dumps(prompt) + exif_data[0x010F] = "Prompt:" + prompt_str + if extra_pnginfo is not None: + if isinstance(extra_pnginfo, list): + merged_extra = {} + for item in extra_pnginfo: + if isinstance(item, dict): + merged_extra.update(item) + extra_pnginfo = merged_extra + if isinstance(extra_pnginfo, dict): + workflow_metadata = json.dumps(extra_pnginfo) + exif_data[0x010E] = "Workflow:" + workflow_metadata + exif_bytes = exif_data.tobytes() + img.save(path, exif=exif_bytes) + else: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + if isinstance(extra_pnginfo, list): + merged_extra = {} + for item in extra_pnginfo: + if isinstance(item, dict): + merged_extra.update(item) + extra_pnginfo = merged_extra + if isinstance(extra_pnginfo, dict): + for key in extra_pnginfo: + metadata.add_text(key, json.dumps(extra_pnginfo[key])) + img.save(path, pnginfo=metadata) + else: + img.save(path) + + def _save_multi_frame_tensor( + self, tensor, path, embed_workflow, prompt, extra_pnginfo, extension + ): + """Save a 4D tensor (batch, height, width, channels) as multi-frame image""" + frames = [] + for i in range(tensor.shape[0]): + frame = tensor[i] + + # Ensure frame is on CPU + if frame.is_cuda: + frame = frame.cpu() + + # Convert to numpy + if frame.dtype == torch.uint8: + frame_np = frame.numpy() + else: + frame_np = np.clip(255.0 * frame.numpy(), 0, 255).astype(np.uint8) + + # Handle different frame shapes + if frame_np.ndim == 3: + # 3D: could be (C, H, W) or (H, W, C) + if frame_np.shape[0] == 3 or frame_np.shape[0] == 1: + # (C, H, W) format - transpose to (H, W, C) + frame_np = np.transpose(frame_np, (1, 2, 0)) + elif frame_np.shape[2] == 3 or frame_np.shape[2] == 1: + # Already in (H, W, C) format + pass + else: + # Try to fix unusual shapes + if frame_np.shape[0] == 1: + frame_np = np.transpose(frame_np, (1, 2, 0)) + elif frame_np.ndim == 2: + # 2D grayscale - add channel dimension + frame_np = np.expand_dims(frame_np, axis=2) + + # Ensure valid format for PIL + if frame_np.ndim == 2: + # Grayscale + pil_img = Image.fromarray(frame_np, mode="L") + elif frame_np.ndim == 3: + if frame_np.shape[2] == 1: + # Single channel - convert to grayscale + pil_img = Image.fromarray(frame_np.squeeze(2), mode="L") + elif frame_np.shape[2] == 3: + # RGB + pil_img = Image.fromarray(frame_np, mode="RGB") + else: + # More channels - take first 3 + pil_img = Image.fromarray(frame_np[:, :, :3], mode="RGB") + else: + raise ValueError( + f"Unsupported frame shape: {frame.shape} -> {frame_np.shape}" + ) + + frames.append(pil_img) + + # Save metadata only on first frame + first_frame = frames[0] + if embed_workflow: + if extension.lower() == "webp": + exif_data = first_frame.getexif() + if prompt is not None: + prompt_str = json.dumps(prompt) + exif_data[0x010F] = "Prompt:" + prompt_str + if extra_pnginfo is not None: + if isinstance(extra_pnginfo, list): + merged_extra = {} + for item in extra_pnginfo: + if isinstance(item, dict): + merged_extra.update(item) + extra_pnginfo = merged_extra + if isinstance(extra_pnginfo, dict): + workflow_metadata = json.dumps(extra_pnginfo) + exif_data[0x010E] = "Workflow:" + workflow_metadata + exif_bytes = exif_data.tobytes() + first_frame.save( + path, + save_all=True, + append_images=frames[1:] if len(frames) > 1 else [], + duration=100, # 100ms per frame + loop=0, + exif=exif_bytes, + ) + elif extension.lower() == "gif": + first_frame.save( + path, + save_all=True, + append_images=frames[1:] if len(frames) > 1 else [], + duration=100, + loop=0, + ) + elif extension.lower() == "tiff": + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + if isinstance(extra_pnginfo, list): + merged_extra = {} + for item in extra_pnginfo: + if isinstance(item, dict): + merged_extra.update(item) + extra_pnginfo = merged_extra + if isinstance(extra_pnginfo, dict): + for key in extra_pnginfo: + metadata.add_text(key, json.dumps(extra_pnginfo[key])) + first_frame.save( + path, + save_all=True, + append_images=frames[1:] if len(frames) > 1 else [], + pnginfo=metadata, + ) + else: + # For formats that don't support multiple frames, save only first frame + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + if isinstance(extra_pnginfo, list): + merged_extra = {} + for item in extra_pnginfo: + if isinstance(item, dict): + merged_extra.update(item) + extra_pnginfo = merged_extra + if isinstance(extra_pnginfo, dict): + for key in extra_pnginfo: + metadata.add_text(key, json.dumps(extra_pnginfo[key])) + first_frame.save(path, pnginfo=metadata) + print( + f"Warning: {extension} doesn't support multiple frames. Only first frame saved." + ) + else: + if extension.lower() in ["gif", "webp"]: + first_frame.save( + path, + save_all=True, + append_images=frames[1:] if len(frames) > 1 else [], + duration=100, + loop=0, + ) + elif extension.lower() == "tiff": + first_frame.save( + path, + save_all=True, + append_images=frames[1:] if len(frames) > 1 else [], + ) + else: + first_frame.save(path) + print( + f"Warning: {extension} doesn't support multiple frames. Only first frame saved." + ) + + def _get_base_name_without_counter(self, prefix, suffix, delimiter): + """Generate base filename without counter""" + parts = [str(prefix)] + if suffix and suffix.strip(): + parts.append(str(suffix)) + return delimiter.join(parts) + + def _find_existing_file( + self, directory, base_name, extension, counter_start, counter_end, delimiter + ): + """Find existing file with same base name (ignoring counter)""" + if not os.path.exists(directory): + return None + + existing_files = os.listdir(directory) + base_pattern = re.escape(base_name) + + # Build patterns to match files with this base name + if counter_start and counter_end: + # Counter can be at start or end + pattern_start = re.compile( + rf"^(\d+){re.escape(delimiter)}{base_pattern}\.{re.escape(extension)}$" + ) + pattern_end = re.compile( + rf"^{base_pattern}{re.escape(delimiter)}(\d+)\.{re.escape(extension)}$" + ) + elif counter_start: + pattern_start = re.compile( + rf"^(\d+){re.escape(delimiter)}{base_pattern}\.{re.escape(extension)}$" + ) + pattern_end = None + elif counter_end: + pattern_start = None + pattern_end = re.compile( + rf"^{base_pattern}{re.escape(delimiter)}(\d+)\.{re.escape(extension)}$" + ) + else: + # No counter, exact match + exact_name = f"{base_name}.{extension}" + exact_path = os.path.join(directory, exact_name) + if os.path.exists(exact_path): + return exact_path + return None + + # Find matching files + matching_files = [] + for file in existing_files: + if not file.endswith(f".{extension}"): + continue + file_base = os.path.splitext(file)[0] + + match = None + if pattern_start: + match = pattern_start.match(file_base) + if not match and pattern_end: + match = pattern_end.match(file_base) + + if match: + matching_files.append((file, int(match.group(1)))) + + if matching_files: + # Return the file with the highest counter (most recent) + matching_files.sort(key=lambda x: x[1], reverse=True) + return os.path.join(directory, matching_files[0][0]) + + return None + + def _load_image_frames(self, path): + """Load an image file and return all frames as a tensor""" + try: + img = Image.open(path) + frames = [] + + # Handle multi-frame images (GIF, TIFF, WebP) + try: + for frame in ImageSequence.Iterator(img): + frame = frame.convert("RGB") + img_array = np.array(frame).astype(np.float32) / 255.0 + frames.append(torch.from_numpy(img_array)) + except Exception: + # Single frame image + img = img.convert("RGB") + img_array = np.array(img).astype(np.float32) / 255.0 + frames.append(torch.from_numpy(img_array)) + + if len(frames) == 0: + return None + + if len(frames) == 1: + return frames[0] + else: + # Stack frames along batch dimension: (num_frames, H, W, C) + return torch.stack(frames, dim=0) + except Exception as e: + print(f"Error loading image frames from {path}: {str(e)}") + return None + + def _get_first_or_default(self, value, default): + if isinstance(value, list): + for v in value: + if isinstance(v, type(default)): + return v + return default + return value if value is not None else default + + def _normalize_input(self, input_data, count): + if input_data is None: + return [None] * count + if not isinstance(input_data, list): + return [input_data] * count + return input_data + [input_data[-1]] * (count - len(input_data)) + + def _get_output_path(self, base_dir, user_path): + if not user_path or str(user_path).lower() in ["none", "."]: + return base_dir + return ( + os.path.join(base_dir, user_path) + if not os.path.isabs(str(user_path)) + else user_path + ) + + def _generate_filename( + self, + prefix, + suffix, + padding, + counter_start, + counter_end, + delimiter, + final_output_path, + ): + existing_files = os.listdir(final_output_path) + + if suffix.strip(): + pattern_start = re.compile( + rf"^(\d+){re.escape(delimiter)}{re.escape(prefix)}{re.escape(delimiter)}{re.escape(suffix)}$" + ) + pattern_end = re.compile( + rf"^{re.escape(prefix)}{re.escape(delimiter)}{re.escape(suffix)}{re.escape(delimiter)}(\d+)$" + ) + else: + pattern_start = re.compile( + rf"^(\d+){re.escape(delimiter)}{re.escape(prefix)}$" + ) + pattern_end = re.compile( + rf"^{re.escape(prefix)}{re.escape(delimiter)}(\d+)$" + ) + + numbers = [] + for file in existing_files: + file_base = os.path.splitext(file)[0] + if counter_start and not counter_end: + match = pattern_start.match(file_base) + elif counter_end and not counter_start: + match = pattern_end.match(file_base) + elif counter_start and counter_end: + match_start = pattern_start.match(file_base) + match_end = pattern_end.match(file_base) + match = match_start or match_end + if match_start and match_end: + numbers.extend([int(match_start.group(1)), int(match_end.group(1))]) + continue + else: + match = None + if match: + numbers.append(int(match.group(1))) + + counter = max(numbers) + 1 if numbers else 1 + parts = [] + if counter_start: + parts.append(f"{counter:0{padding}d}") + parts.append(str(prefix)) + if suffix.strip(): + parts.append(str(suffix)) + if counter_end: + parts.append(f"{counter:0{padding}d}") + + return delimiter.join(parts) + + def _update_progress(self, node_id, current, total): + if node_id: + PromptServer.instance.send_sync( + "progress", {"node": node_id, "value": current, "max": total} + ) diff --git a/nodes/lora_batch_process.py b/nodes/lora_batch_process.py new file mode 100644 index 0000000..59f3f64 --- /dev/null +++ b/nodes/lora_batch_process.py @@ -0,0 +1,335 @@ +import os +import random +import hashlib +import re +import sys +import io +import logging +import warnings +from contextlib import contextmanager +import folder_paths +import comfy.utils +import comfy.lora + + +@contextmanager +def suppress_comfy_logs(): + """Temporarily suppress stdout, stderr, logging, and warnings to hide verbose ComfyUI LoRA loading messages""" + # Save original states + old_stdout = sys.stdout + old_stderr = sys.stderr + old_logging_level = logging.root.level + + # Also save file descriptors for low-level redirection + old_stdout_fd = None + old_stderr_fd = None + devnull_fd = None + + try: + # Suppress Python-level output + logging.root.setLevel(logging.CRITICAL + 1) + warnings.filterwarnings("ignore") + + # Create a devnull-like object for Python-level redirection + devnull = io.StringIO() + sys.stdout = devnull + sys.stderr = devnull + + # Also redirect OS-level file descriptors (for C extensions, PyTorch, etc.) + # This catches output that bypasses Python's sys.stdout/stderr + try: + sys.stdout.flush() + sys.stderr.flush() + + # Duplicate the current stdout/stderr file descriptors + old_stdout_fd = os.dup(1) + old_stderr_fd = os.dup(2) + + # Open devnull and redirect stdout/stderr to it + devnull_fd = os.open(os.devnull, os.O_WRONLY) + os.dup2(devnull_fd, 1) # Redirect stdout (fd 1) + os.dup2(devnull_fd, 2) # Redirect stderr (fd 2) + except (OSError, AttributeError): + # If OS-level redirection fails, continue with Python-level only + pass + + yield + + finally: + # Restore Python-level output + sys.stdout = old_stdout + sys.stderr = old_stderr + logging.root.setLevel(old_logging_level) + warnings.filterwarnings("default") + + # Restore OS-level file descriptors + if old_stdout_fd is not None: + try: + sys.stdout.flush() + sys.stderr.flush() + os.dup2(old_stdout_fd, 1) + os.dup2(old_stderr_fd, 2) + os.close(old_stdout_fd) + os.close(old_stderr_fd) + if devnull_fd is not None: + os.close(devnull_fd) + except (OSError, AttributeError): + pass + + +class LoraBatchLoader: + RETURN_TYPES = ("MODEL", "CLIP", "STRING") + RETURN_NAMES = ("model", "clip", "filename") + FUNCTION = "load_batch_loras" + CATEGORY = "Batch Process" + + SUPPORTED_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".bin"} + + def __init__(self): + self.lora_states = {} + self.current_directory = "" + self.loras = [] + self.search_states = {} + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "directory": ("STRING",), + "search_title": ("STRING", {"default": ""}), + "delimiter": ("STRING", {"default": ""}), + "mode": ( + ["incremental", "random"], + {"default": "incremental"}, + ), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "filename_option": ( + [ + "filename", + "prefix", + "suffix", + "prefix & suffix", + "prefix nor suffix", + ], + ), + "strength_model": ( + "FLOAT", + {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}, + ), + "strength_clip": ( + "FLOAT", + {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}, + ), + }, + } + + def set_directory( + self, directory, filename_option="filename", search_title="", delimiter="" + ): + if ( + directory != self.current_directory + or filename_option + or search_title + or delimiter + ): + if not os.path.isdir(directory): + raise ValueError( + f"The provided path '{directory}' is not a valid directory." + ) + + all_loras = [ + f + for f in os.listdir(directory) + if any(f.lower().endswith(ext) for ext in self.SUPPORTED_EXTENSIONS) + ] + filtered_loras = self.filter_loras( + directory, all_loras, filename_option, search_title, delimiter + ) + + self.loras = sorted([os.path.join(directory, f) for f in filtered_loras]) + self.current_directory = directory + + search_key = (directory, filename_option, search_title, delimiter) + if search_key not in self.search_states: + self.search_states[search_key] = 0 + + if not self.loras: + print("No matching LoRA files found in the provided directory.") + else: + print(f"Found {len(self.loras)} LoRA files in directory.") + + def load_loras(self, directory): + if not os.path.isdir(directory): + raise ValueError(f"Invalid directory: {directory}") + + all_loras = [ + f + for f in os.listdir(directory) + if any(f.lower().endswith(ext) for ext in self.SUPPORTED_EXTENSIONS) + ] + return sorted([os.path.join(directory, f) for f in all_loras]) + + def filter_loras(self, directory, files, filename_option, search_title, delimiter): + def get_prefix(filename): + if delimiter: + return filename.split(delimiter)[0] + else: + return re.split(r"[^a-zA-Z0-9]", filename)[0] + + def get_suffix(filename): + name_without_ext = os.path.splitext(filename)[0] + if delimiter: + return name_without_ext.split(delimiter)[-1] + else: + return re.split(r"[^a-zA-Z0-9]", name_without_ext)[-1] + + filtered_files = files + + if search_title: + if filename_option == "filename": + filtered_files = [f for f in filtered_files if search_title in f] + elif filename_option == "prefix": + search_prefix = get_prefix(search_title) + filtered_files = [ + f for f in filtered_files if get_prefix(f) == search_prefix + ] + elif filename_option == "suffix": + search_suffix = get_suffix(search_title) + filtered_files = [ + f for f in filtered_files if get_suffix(f) == search_suffix + ] + elif filename_option == "prefix & suffix": + search_prefix = get_prefix(search_title) + search_suffix = get_suffix(search_title) + filtered_files = [ + f + for f in filtered_files + if get_prefix(f) == search_prefix or get_suffix(f) == search_suffix + ] + elif filename_option == "prefix nor suffix": + search_prefix = get_prefix(search_title) + search_suffix = get_suffix(search_title) + filtered_files = [ + f + for f in filtered_files + if get_prefix(f) != search_prefix and get_suffix(f) != search_suffix + ] + + return filtered_files + + def load_batch_loras( + self, + model, + clip, + directory, + search_title="", + delimiter="", + mode="incremental", + seed=0, + filename_option="filename", + strength_model=1.0, + strength_clip=1.0, + ): + self.set_directory(directory, filename_option, search_title, delimiter) + + if not self.loras: + print("No LoRA files found, returning original model and clip.") + return (model, clip, "no_loras_found") + + search_key = (directory, filename_option, search_title, delimiter) + + if mode == "incremental": + return self.load_lora_by_index( + model, clip, search_key, strength_model, strength_clip + ) + elif mode == "random": + random.seed(seed) + rnd_index = random.randint(0, len(self.loras) - 1) + print( + f"[LoRA Batch Loader] Random mode - Index: {rnd_index + 1}/{len(self.loras)}" + ) + return self.load_lora_by_path( + model, clip, self.loras[rnd_index], strength_model, strength_clip + ) + else: + raise ValueError(f"Unknown mode: {mode}") + + def load_lora_by_index( + self, model, clip, search_key, strength_model, strength_clip + ): + if not self.loras: + print("No LoRAs loaded.") + return model, clip, "no_loras" + + current_index = self.search_states[search_key] + if current_index >= len(self.loras): + current_index = 0 + + file_path = self.loras[current_index] + self.search_states[search_key] = (current_index + 1) % len(self.loras) + + # Print index info for better tracking + print(f"[LoRA Batch Loader] Index: {current_index + 1}/{len(self.loras)}") + return self.load_lora_by_path( + model, clip, file_path, strength_model, strength_clip + ) + + def load_lora_by_path(self, model, clip, path, strength_model, strength_clip): + try: + filename = os.path.basename(path) + # Remove file extension for cleaner filename output + filename_clean = os.path.splitext(filename)[0] + + # Print LoRA name every time it's loaded + print(f"[LoRA Batch Loader] Loading LoRA: {filename_clean}") + + # Suppress verbose ComfyUI LoRA loading messages + with suppress_comfy_logs(): + # Load the LoRA using ComfyUI's built-in LoRA loading functionality + lora = comfy.utils.load_torch_file(path, safe_load=True) + + # Create key mapping for the LoRA + model_lora_keys = comfy.lora.model_lora_keys_unet(model.model) + clip_lora_keys = comfy.lora.model_lora_keys_clip(clip.cond_stage_model) + + # Combine the key mappings + key_map = {} + key_map.update(model_lora_keys) + key_map.update(clip_lora_keys) + + # Load the LoRA patches (this prints all the verbose messages) + loaded = comfy.lora.load_lora(lora, key_map) + + # Apply LoRA to model + new_modelpatcher = model.clone() + k = {} + for x in loaded: + k[x] = loaded[x] + new_modelpatcher.add_patches(k, strength_model) + + # Apply LoRA to clip + new_clip = clip.clone() + k = {} + for x in loaded: + k[x] = loaded[x] + new_clip.add_patches(k, strength_clip) + + return (new_modelpatcher, new_clip, filename_clean) + + except Exception as e: + print(f"Error loading LoRA {path}: {str(e)}") + return (model, clip, "error_loading_lora") + + @classmethod + def IS_CHANGED(cls, directory, **kwargs): + if not os.path.exists(directory): + return "" + try: + loader = cls() + paths = loader.load_loras(directory) + return hashlib.sha256(",".join(paths).encode()).hexdigest() + except Exception as e: + print(f"Error checking for changes: {str(e)}") + return "" diff --git a/nodes/mask_batch_process.py b/nodes/mask_batch_process.py new file mode 100644 index 0000000..4a612d1 --- /dev/null +++ b/nodes/mask_batch_process.py @@ -0,0 +1,135 @@ +import torch + + +class MaskFromBatch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "masks": ("MASK",), + "batch_index": ("INT", { + "default": 0, + "min": 0, + "max": 63, + "step": 1 + }), + "length": ("INT", { + "default": 1, + "min": 1, + "max": 64, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("selected_masks",) + FUNCTION = "batch_select" + CATEGORY = "Batch Process/Mask" + + def batch_select(self, masks, batch_index, length): + # 确保mask是正确的形状 + if isinstance(masks, torch.Tensor): + if len(masks.shape) == 2: + masks = masks.unsqueeze(0) + + # 获取批次大小 + batch_size = masks.shape[0] + + # 确保索引和长度在有效范围内 + batch_index = min(batch_size - 1, batch_index) + length = min(batch_size - batch_index, length) + + # 选择指定的批次 + selected_masks = masks[batch_index:batch_index + length].clone() + + return (selected_masks,) + + +class MaskRepeatBatch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "masks": ("MASK",), + "amount": ("INT", { + "default": 1, + "min": 1, + "max": 64, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("repeated_masks",) + FUNCTION = "repeat" + CATEGORY = "Batch Process/Mask" + + def repeat(self, masks, amount): + # 确保mask是正确的形状 + if isinstance(masks, torch.Tensor): + if len(masks.shape) == 2: + masks = masks.unsqueeze(0) + + # 重复指定次数 + repeated_masks = masks.repeat((amount, 1, 1)) + + return (repeated_masks,) + + +class MaskBatchCopy: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "masks": ("MASK",), + } + } + + RETURN_TYPES = ("MASK", "MASK") + RETURN_NAMES = ("masks_1", "masks_2") + FUNCTION = "copy" + CATEGORY = "Batch Process/Mask" + + def copy(self, masks): + # 确保mask是正确的形状 + if isinstance(masks, torch.Tensor): + if len(masks.shape) == 2: + masks = masks.unsqueeze(0) + + # 创建两个独立的副本 + copy1 = masks.clone() + copy2 = masks.clone() + + return (copy1, copy2) + + +class MaskBatchComposite: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "masks_1": ("MASK",), + "masks_2": ("MASK",), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("combined_masks",) + FUNCTION = "combine" + CATEGORY = "Batch Process/Mask" + + def combine(self, masks_1, masks_2): + # 确保两个mask都是正确的形状 + if isinstance(masks_1, torch.Tensor): + if len(masks_1.shape) == 2: + masks_1 = masks_1.unsqueeze(0) + if isinstance(masks_2, torch.Tensor): + if len(masks_2.shape) == 2: + masks_2 = masks_2.unsqueeze(0) + + # 合并mask批次 + combined_masks = torch.cat([masks_1, masks_2], dim=0) + + return (combined_masks,) diff --git a/nodes/model3d_batch_process.py b/nodes/model3d_batch_process.py new file mode 100644 index 0000000..18cc593 --- /dev/null +++ b/nodes/model3d_batch_process.py @@ -0,0 +1,327 @@ +import os +import re +import uuid +import logging + +import numpy as np +import torch +from PIL import Image + +import folder_paths +from server import PromptServer + +try: + from comfy_api.latest._util.geometry_types import MESH as _MESH_TYPE, File3D as _FILE3D_TYPE +except Exception: + _MESH_TYPE = None + _FILE3D_TYPE = None + +try: + from comfy_extras.nodes_save_3d import save_glb as _save_glb, get_mesh_batch_item as _get_mesh_batch_item +except Exception: + _save_glb = None + _get_mesh_batch_item = None + + +class _AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + + +any_type = _AnyType("*") + +MODEL_3D_EXTS = {".glb", ".gltf", ".obj", ".fbx", ".stl", ".ply", ".usdz", ".splat", ".spz", ".ksplat"} + + +def _is_mesh(v): + return _MESH_TYPE is not None and isinstance(v, _MESH_TYPE) + + +def _is_file3d(v): + return _FILE3D_TYPE is not None and isinstance(v, _FILE3D_TYPE) + + +def _flatten_models(models): + """Normalize any model input shape to a flat list of items. + + Each item is one of: (kind, payload) where kind ∈ {'mesh_item', 'file3d', 'path'}. + Mesh batches are split per-row; lists are flattened. + """ + out = [] + + def push(v): + if v is None: + return + if isinstance(v, (list, tuple)): + for x in v: + push(x) + return + if _is_mesh(v): + if _get_mesh_batch_item is None: + raise RuntimeError("Model3D batch IO: MESH support requires comfy_extras.nodes_save_3d") + b = v.vertices.shape[0] + tex = getattr(v, "texture", None) + for i in range(b): + verts, faces, colors, uvs = _get_mesh_batch_item(v, i) + tex_i = tex[i] if tex is not None else None + out.append(("mesh_item", (verts, faces, colors, uvs, tex_i))) + return + if _is_file3d(v): + out.append(("file3d", v)) + return + if isinstance(v, str): + if os.path.splitext(v)[1].lower() in MODEL_3D_EXTS: + out.append(("path", v)) + return + raise ValueError(f"Model3D batch IO: string '{v}' does not look like a 3D model path") + raise ValueError(f"Model3D batch IO: unsupported input type {type(v).__name__}") + + push(models) + return out + + +def _save_one(item, dest_path, metadata=None): + """Write a single normalized model item to dest_path. Returns the extension actually written.""" + kind, payload = item + requested_ext = os.path.splitext(dest_path)[1].lstrip(".").lower() + + if kind == "mesh_item": + if _save_glb is None: + raise RuntimeError("Model3D batch IO: GLB writer unavailable") + verts, faces, colors, uvs, tex = payload + if verts.shape[0] == 0 or faces.shape[0] == 0: + raise ValueError("empty mesh") + tex_img = None + if tex is not None: + arr = (tex.clamp(0.0, 1.0).cpu().numpy() * 255).astype(np.uint8) + tex_img = Image.fromarray(arr, mode="RGB") + # MESH path is always written as .glb regardless of the requested extension. + if requested_ext != "glb": + dest_path = os.path.splitext(dest_path)[0] + ".glb" + _save_glb(verts, faces, dest_path, metadata=metadata, + uvs=uvs, vertex_colors=colors, texture_image=tex_img) + return "glb", dest_path + + if kind == "file3d": + file3d = payload + src_ext = (file3d.format or "glb").lower() + if requested_ext != src_ext: + dest_path = os.path.splitext(dest_path)[0] + "." + src_ext + file3d.save_to(dest_path) + return src_ext, dest_path + + if kind == "path": + src = payload + src_ext = os.path.splitext(src)[1].lstrip(".").lower() or "glb" + if requested_ext != src_ext: + dest_path = os.path.splitext(dest_path)[0] + "." + src_ext + if os.path.abspath(src) != os.path.abspath(dest_path): + import shutil + shutil.copy2(src, dest_path) + return src_ext, dest_path + + raise ValueError(f"unknown item kind: {kind}") + + +def _normalize_list(value, count): + if value is None: + return [None] * count + if not isinstance(value, list): + return [value] * count + if len(value) >= count: + return value[:count] + return value + [value[-1]] * (count - len(value)) + + +def _resolve_output_path(base_dir, user_path): + if not user_path or str(user_path).lower() in ("none", "."): + return base_dir + return user_path if os.path.isabs(str(user_path)) else os.path.join(base_dir, user_path) + + +def _generate_filename(prefix, suffix, padding, counter_start, counter_end, delimiter, out_dir, ext): + try: + existing = [f for f in os.listdir(out_dir) if f.endswith("." + ext) or f.endswith(".txt")] + except FileNotFoundError: + existing = [] + + if suffix.strip(): + pat_start = re.compile(rf"^(\d+){re.escape(delimiter)}{re.escape(prefix)}{re.escape(delimiter)}{re.escape(suffix)}$") + pat_end = re.compile(rf"^{re.escape(prefix)}{re.escape(delimiter)}{re.escape(suffix)}{re.escape(delimiter)}(\d+)$") + else: + pat_start = re.compile(rf"^(\d+){re.escape(delimiter)}{re.escape(prefix)}$") + pat_end = re.compile(rf"^{re.escape(prefix)}{re.escape(delimiter)}(\d+)$") + + nums = [] + for f in existing: + base = os.path.splitext(f)[0] + match = None + if counter_start and not counter_end: + match = pat_start.match(base) + elif counter_end and not counter_start: + match = pat_end.match(base) + elif counter_start and counter_end: + ms, me = pat_start.match(base), pat_end.match(base) + if ms: + nums.append(int(ms.group(1))) + if me: + nums.append(int(me.group(1))) + continue + if match: + nums.append(int(match.group(1))) + + counter = max(nums) + 1 if nums else 1 + parts = [] + if counter_start: + parts.append(f"{counter:0{padding}d}") + parts.append(str(prefix)) + if suffix.strip(): + parts.append(str(suffix)) + if counter_end: + parts.append(f"{counter:0{padding}d}") + return delimiter.join(parts) + + +def _send_progress(node_id, current, total): + if node_id and PromptServer.instance: + PromptServer.instance.send_sync("progress", {"node": node_id, "value": current, "max": total}) + + +class Model3DBatchSaver: + FUNCTION = "save" + CATEGORY = "Batch Process" + RETURN_TYPES = () + RETURN_NAMES = () + OUTPUT_NODE = True + + @classmethod + def IS_CHANGED(cls, **kwargs): + import time + return time.time() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "models": (any_type, {"tooltip": "MESH batch, File3D, or list of either / paths (e.g. AnyBatchGroup output)"}), + "contents": ("STRING", {"forceInput": True, "tooltip": "Optional caption written next to each model as .txt"}), + "output_path": ("STRING", {"default": ""}), + "filename_prefix": ("STRING", {"default": "M3D"}), + "filename_delimiter": ("STRING", {"default": "_"}), + "filename_suffix": ("STRING", {"default": ""}), + "filename_number_padding": ("INT", {"default": 4, "min": 1, "max": 9, "step": 1}), + "filename_number": (["off", "start", "end"], {"default": "end"}), + "default_extension": (["glb", "gltf", "obj", "fbx", "stl", "ply", "usdz"], { + "default": "glb", + "tooltip": "Used only for naming/collision check. File3D items keep their own extension; MESH always writes as .glb", + }), + }, + "hidden": { + "node_id": "UNIQUE_ID", + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + def save(self, models=None, contents=None, output_path="", filename_prefix="M3D", + filename_delimiter="_", filename_suffix="", filename_number_padding=4, + filename_number="end", default_extension="glb", + node_id=None, prompt=None, extra_pnginfo=None): + if models is None: + logging.warning("Model3DBatchSaver: no models provided") + return () + + items = _flatten_models(models) + if not items: + logging.warning("Model3DBatchSaver: model input flattened to zero items") + return () + + counter_start = filename_number == "start" + counter_end = filename_number == "end" + out_dir = _resolve_output_path(folder_paths.get_output_directory(), output_path) + os.makedirs(out_dir, exist_ok=True) + + captions = _normalize_list(contents, len(items)) if contents is not None else [None] * len(items) + original_prefix = os.path.splitext(os.path.basename(filename_prefix))[0] + + for idx, (item, caption) in enumerate(zip(items, captions)): + try: + base_name = _generate_filename( + prefix=original_prefix, + suffix=filename_suffix.strip("'[]"), + padding=filename_number_padding, + counter_start=counter_start, + counter_end=counter_end, + delimiter=filename_delimiter, + out_dir=out_dir, + ext=default_extension, + ) + dest = os.path.join(out_dir, f"{base_name}.{default_extension}") + actual_ext, actual_dest = _save_one(item, dest) + logging.info(f"Model3DBatchSaver: wrote {actual_dest}") + + if caption is not None and str(caption).strip(): + txt_path = os.path.splitext(actual_dest)[0] + ".txt" + with open(txt_path, "w", encoding="utf-8") as f: + f.write(str(caption).strip()) + except Exception as e: + logging.exception(f"Model3DBatchSaver: failed item {idx}: {e}") + + _send_progress(node_id, idx + 1, len(items)) + + return () + + +class Model3DBatchPreview: + FUNCTION = "preview" + CATEGORY = "Batch Process" + RETURN_TYPES = () + RETURN_NAMES = () + OUTPUT_NODE = True + + @classmethod + def IS_CHANGED(cls, **kwargs): + import time + return time.time() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "models": (any_type, {"tooltip": "MESH batch, File3D, or list of either / paths"}), + }, + "hidden": {"node_id": "UNIQUE_ID"}, + } + + def preview(self, models=None, node_id=None): + if models is None: + return {"ui": {"3d": []}} + + items = _flatten_models(models) + if not items: + return {"ui": {"3d": []}} + + out_dir = folder_paths.get_output_directory() + subfolder = f"preview3d_batch_{uuid.uuid4().hex[:8]}" + full_dir = os.path.join(out_dir, subfolder) + os.makedirs(full_dir, exist_ok=True) + + results = [] + for idx, item in enumerate(items): + try: + filename = f"preview_{idx:04d}.glb" + dest = os.path.join(full_dir, filename) + actual_ext, actual_dest = _save_one(item, dest) + actual_filename = os.path.basename(actual_dest) + results.append({ + "filename": actual_filename, + "subfolder": subfolder, + "type": "output", + }) + except Exception as e: + logging.exception(f"Model3DBatchPreview: failed item {idx}: {e}") + _send_progress(node_id, idx + 1, len(items)) + + return {"ui": {"3d": results}} diff --git a/nodes/runtime_filename_format.py b/nodes/runtime_filename_format.py new file mode 100644 index 0000000..2ed3604 --- /dev/null +++ b/nodes/runtime_filename_format.py @@ -0,0 +1,140 @@ +"""Runtime output substitutions for ComfyUI save filename templates. + +ComfyUI's built-in ``%Node.widget%`` formatter only has access to input +widgets while the prompt is being queued. Batch loader filenames are runtime +outputs, so unresolved placeholders have to be expanded after the loader has +actually selected a file and before the saver creates its output path. +""" + +from collections import OrderedDict +from functools import wraps +import re +import threading + + +_PLACEHOLDER_RE = re.compile(r"%([^%]+)\.([^%.]+)%") +_INVALID_FILENAME_CHARS_RE = re.compile(r'[/?<>\\:*|"\x00-\x1f\x7f]') +_MAX_PROMPTS = 128 +_LOCK = threading.RLock() +_VALUES_BY_PROMPT = OrderedDict() +_LATEST_VALUES = {} + + +def _current_prompt_id(): + try: + from comfy_execution.utils import get_executing_context + + context = get_executing_context() + return context.prompt_id if context is not None else None + except (ImportError, AttributeError): + # Older ComfyUI versions do not expose the execution context helper. + return None + + +def _workflow_aliases(unique_id, extra_pnginfo): + aliases = set() + if unique_id is None or not isinstance(extra_pnginfo, dict): + return aliases + + workflow = extra_pnginfo.get("workflow") + if not isinstance(workflow, dict): + return aliases + + for node in workflow.get("nodes", []): + if not isinstance(node, dict) or str(node.get("id")) != str(unique_id): + continue + + for key in ("title", "type"): + value = node.get(key) + if value: + aliases.add(str(value)) + + properties = node.get("properties") + if isinstance(properties, dict): + value = properties.get("Node name for S&R") + if value: + aliases.add(str(value)) + break + + return aliases + + +def record_runtime_output( + node_type, + field, + value, + *, + unique_id=None, + extra_pnginfo=None, + display_name=None, + prompt_id=None, +): + """Record a node output so a saver can reference it in the same prompt.""" + aliases = {str(node_type)} + if display_name: + aliases.add(str(display_name)) + aliases.update(_workflow_aliases(unique_id, extra_pnginfo)) + + resolved_prompt_id = prompt_id if prompt_id is not None else _current_prompt_id() + with _LOCK: + if resolved_prompt_id is not None: + values = _VALUES_BY_PROMPT.setdefault(str(resolved_prompt_id), {}) + _VALUES_BY_PROMPT.move_to_end(str(resolved_prompt_id)) + while len(_VALUES_BY_PROMPT) > _MAX_PROMPTS: + _VALUES_BY_PROMPT.popitem(last=False) + else: + values = _LATEST_VALUES + + for alias in aliases: + values.setdefault(alias, {})[str(field)] = value + # ComfyUI can reuse a loader's cached output in a later prompt, + # which means the Python function is not called again for that + # prompt. Retaining the most recently executed value makes the + # saver resolve to the same filename as that cached output. + _LATEST_VALUES.setdefault(alias, {})[str(field)] = value + + +def replace_runtime_placeholders(text, prompt_id=None): + """Expand recorded ``%Node.output%`` values and preserve unknown tokens.""" + if not isinstance(text, str) or "%" not in text: + return text + + resolved_prompt_id = prompt_id if prompt_id is not None else _current_prompt_id() + with _LOCK: + values = {} + if resolved_prompt_id is not None: + values.update(_VALUES_BY_PROMPT.get(str(resolved_prompt_id), {})) + # This fallback keeps the feature usable on older, serial-execution + # ComfyUI versions that do not expose the current prompt id. + for alias, fields in _LATEST_VALUES.items(): + values.setdefault(alias, fields) + + def substitute(match): + alias, field = match.groups() + if alias not in values or field not in values[alias]: + return match.group(0) + value = str(values[alias][field] or "") + return _INVALID_FILENAME_CHARS_RE.sub("_", value) + + return _PLACEHOLDER_RE.sub(substitute, text) + + +def install_runtime_filename_substitution(): + """Install the lightweight hook used by all ComfyUI saver nodes.""" + try: + import folder_paths + except ImportError: + return False + + original = folder_paths.get_save_image_path + if getattr(original, "_batch_process_runtime_format", False): + return True + + @wraps(original) + def get_save_image_path(filename_prefix, *args, **kwargs): + filename_prefix = replace_runtime_placeholders(filename_prefix) + return original(filename_prefix, *args, **kwargs) + + get_save_image_path._batch_process_runtime_format = True + folder_paths.get_save_image_path = get_save_image_path + return True diff --git a/nodes/text_batch_process.py b/nodes/text_batch_process.py new file mode 100644 index 0000000..be44b35 --- /dev/null +++ b/nodes/text_batch_process.py @@ -0,0 +1,495 @@ +import os +import random +import re + +from .runtime_filename_format import record_runtime_output + + +class TXTBatchLoader: + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("content", "filename") + FUNCTION = "load_batch_texts" + CATEGORY = "Batch Process" + + def __init__(self): + self.txt_files = [] + self.current_index = 0 + self.current_directory = "" + self.search_states = {} + # When reset_on_queue value changes (any change), we reset index to 0 for that search_key + self._last_reset_on_queue = {} + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "directory": ("STRING",), + "search_title": ("STRING", {"default": ""}), + "search_content": ("STRING", {"default": ""}), + "delimiter": ("STRING", {"default": ""}), + "mode": ( + ["single_text", "incremental_text", "random"], + {"default": "incremental_text"}, + ), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "filename_option": ( + [ + "filename", + "prefix", + "suffix", + "prefix & suffix", + "prefix nor suffix", + ], + ), + "start_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "tooltip": "Start index (1-based). Use 0 to start from beginning, or set 1 for first text file, 2 for second text file, etc.", + }, + ), + "end_index": ( + "INT", + { + "default": 0, + "min": 0, + "max": 999999, + "step": 1, + "tooltip": "End index (inclusive, 1-based). Use 0 to include all remaining text files, or set >0 to limit range (1=first, 2=second, etc.).", + }, + ), + }, + "optional": { + "reset_on_queue": ( + "INT", + { + "default": 1, + "min": 0, + "max": 1, + "step": 1, + "tooltip": "Any change to this value (0↔1) resets the read index to start from the first file again. Keep unchanged during a queue run to advance normally.", + }, + ), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + def set_directory( + self, + directory, + filename_option="filename", + search_title="", + search_content="", + delimiter="", + ): + if ( + directory != self.current_directory + or filename_option + or search_title + or search_content + or delimiter + ): + print(f"Received directory: {directory}") + + if not os.path.isdir(directory): + raise ValueError( + f"The provided path '{directory}' is not a valid directory." + ) + + all_files = [f for f in os.listdir(directory) if f.endswith(".txt")] + + filtered_files = self.filter_files( + directory, + all_files, + filename_option, + search_title, + search_content, + delimiter, + ) + + # Natural sort: 1, 2, 3, ..., 10 instead of 1, 10, 2, 3, ... + def _natural_sort_key(path): + name = os.path.basename(path) + return [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", name) + ] + + self.txt_files = sorted( + [os.path.join(directory, f) for f in filtered_files], + key=_natural_sort_key, + ) + + self.current_directory = directory + + search_key = ( + directory, + filename_option, + search_title, + search_content, + delimiter, + ) + if search_key not in self.search_states: + self.search_states[search_key] = 0 + + if not self.txt_files: + print("No matching TXT files found in the provided directory.") + else: + # print(f'Matching TXT files found: {self.txt_files}') + pass + + def filter_files( + self, directory, files, filename_option, search_title, search_content, delimiter + ): + def get_prefix(filename): + if delimiter: + return filename.split(delimiter)[0] + else: + return re.split(r"[^a-zA-Z0-9]", filename)[0] + + def get_suffix(filename): + name_without_ext = os.path.splitext(filename)[0] + if delimiter: + return name_without_ext.split(delimiter)[-1] + else: + return re.split(r"[^a-zA-Z0-9]", name_without_ext)[-1] + + filtered_files = files + + if search_title: + if filename_option == "filename": + filtered_files = [f for f in filtered_files if search_title in f] + elif filename_option == "prefix": + search_prefix = get_prefix(search_title) + filtered_files = [ + f for f in filtered_files if get_prefix(f) == search_prefix + ] + elif filename_option == "suffix": + search_suffix = get_suffix(search_title) + filtered_files = [ + f for f in filtered_files if get_suffix(f) == search_suffix + ] + elif filename_option == "prefix & suffix": + search_prefix = get_prefix(search_title) + search_suffix = get_suffix(search_title) + filtered_files = [ + f + for f in filtered_files + if get_prefix(f) == search_prefix or get_suffix(f) == search_suffix + ] + elif filename_option == "prefix nor suffix": + search_prefix = get_prefix(search_title) + search_suffix = get_suffix(search_title) + filtered_files = [ + f + for f in filtered_files + if get_prefix(f) != search_prefix and get_suffix(f) != search_suffix + ] + + if search_content: + content_filtered_files = [] + for f in filtered_files: + file_path = os.path.join(directory, f) + with open(file_path, "r", encoding="utf-8") as file: + content = file.read() + pattern = ( + r"\b(?= len(self.txt_files): + current_index = 0 + + file_path = self.txt_files[current_index] + self.search_states[search_key] = (current_index + 1) % len(self.txt_files) + + with open(file_path, "r", encoding="utf-8") as file: + content = file.read() + + print(f"Current index: {current_index}, Reading file: {file_path}") + return content, os.path.basename(file_path) + + def load_batch_texts( + self, + directory, + search_title="", + search_content="", + delimiter="", + mode="incremental_text", + seed=0, + filename_option="filename", + start_index=0, + end_index=0, + reset_on_queue=1, + unique_id=None, + extra_pnginfo=None, + ): + self.set_directory( + directory, filename_option, search_title, search_content, delimiter + ) + + if not self.txt_files: + return ("", "") + + # Apply index filtering (1-based) + filtered_files = self.txt_files + if filtered_files: + # Convert 1-indexed to 0-indexed: subtract 1 if > 0 + actual_start = (start_index - 1) if start_index > 0 else 0 + actual_end = ( + (end_index - 1) if end_index > 0 else (len(filtered_files) - 1) + ) + + # Validate and clamp indices + if actual_start < 0: + actual_start = 0 + if actual_start >= len(filtered_files): + actual_start = ( + len(filtered_files) - 1 if len(filtered_files) > 0 else 0 + ) + + if actual_end < 0: + actual_end = 0 + if actual_end >= len(filtered_files): + actual_end = len(filtered_files) - 1 + if actual_end < actual_start: + actual_end = actual_start + + # Only apply slicing if we have valid indices and files + if len(filtered_files) > 0 and actual_start <= actual_end: + # Slice the files list (actual_end+1 because slice is exclusive on end) + filtered_files = filtered_files[actual_start : actual_end + 1] + else: + filtered_files = [] + + if not filtered_files: + return ("", "") + + search_key = ( + directory, + filename_option, + search_title, + search_content, + delimiter, + start_index, + end_index, + ) + + # When reset_on_queue value changes (any change), reset index + if self._last_reset_on_queue.get(search_key) != reset_on_queue: + self.search_states[search_key] = 0 + self._last_reset_on_queue[search_key] = reset_on_queue + + # Initialize search state for this key if it doesn't exist + if search_key not in self.search_states: + self.search_states[search_key] = 0 + + # Use filtered_files instead of self.txt_files for operations + # Temporarily replace self.txt_files with filtered list for get_next_text + original_files = self.txt_files + self.txt_files = filtered_files + + if mode == "single_text": + content, filename = self.get_next_text(search_key) + elif mode == "incremental_text": + content, filename = self.get_next_text(search_key) + elif mode == "random": + random.seed(seed) + rnd_index = random.randint(0, len(filtered_files) - 1) + file_path = filtered_files[rnd_index] + with open(file_path, "r", encoding="utf-8") as file: + content = file.read() + filename = os.path.basename(file_path) + else: + # Restore original files list before raising error + self.txt_files = original_files + raise ValueError(f"Unknown mode: {mode}") + + # Restore original files list + self.txt_files = original_files + + if content: + print(f"Final output content: {content}") + print(f"Final output filename: {filename}") + else: + print("Failed to read any content.") + + filename_without_extension = os.path.splitext(filename)[0] + record_runtime_output( + "TXTBatchLoader", + "filename", + filename_without_extension, + unique_id=unique_id, + extra_pnginfo=extra_pnginfo, + display_name="TXT Batch Loader", + ) + return (content, filename_without_extension) + + +class TextModifyTool: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_text": ("STRING",), + }, + "optional": { + "search": ("STRING", {"default": ""}), + "change_to": ("STRING", {"default": ""}), + "delimiter": ("STRING", {"default": ""}), + "option": ( + ["whole texts", "prefix", "suffix"], + {"default": "whole texts"}, + ), + "delete": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "modify_text" + CATEGORY = "text_processing" + + def find_separator(self, text, find_last=False): + """查找第一个或最后一个分隔符位置""" + sep_chars = ["-", "_", "——", "#"] + target_pos = -1 + found_sep = "" + for sep in sep_chars: + pos = text.rfind(sep) if find_last else text.find(sep) + if pos == -1: + continue + if (find_last and pos > target_pos) or ( + not find_last and (target_pos == -1 or pos < target_pos) + ): + target_pos = pos + found_sep = sep + return target_pos, found_sep + + def modify_text( + self, + input_text, + search="", + change_to="", + delimiter="", + option="whole texts", + delete=False, + ): + # 优先处理search功能 + if search: + replacement = "" if delete else change_to + + if option == "whole texts": + # 分离文件名和扩展名，仅处理文件名部分 + filename, ext = os.path.splitext(input_text) + processed = filename.replace(search, replacement) + return (processed,) + + elif option == "prefix": + if delimiter: + parts = input_text.split(delimiter, 1) + sep = delimiter + else: + # 自动检测第一个分隔符 + pos, sep = self.find_separator(input_text) + if pos != -1: + parts = [input_text[:pos], input_text[pos + len(sep) :]] + else: + parts = [input_text] + + # 处理分割结果 + if len(parts) > 1: + prefix, rest = parts[0], sep + parts[1] + else: + prefix, rest = input_text, "" + + # 执行前缀替换 + new_prefix = prefix.replace(search, replacement) + return (new_prefix + rest,) + + elif option == "suffix": + if delimiter: + parts = input_text.rsplit(delimiter, 1) + sep = delimiter + else: + # 自动检测最后一个分隔符 + pos, sep = self.find_separator(input_text, find_last=True) + if pos != -1: + parts = [input_text[:pos], input_text[pos + len(sep) :]] + else: + parts = [input_text] + + # 处理分割结果 + if len(parts) > 1: + rest, suffix = parts[0] + sep, parts[1] + else: + rest, suffix = "", input_text + + # 执行后缀替换 + new_suffix = suffix.replace(search, replacement) + return (rest + new_suffix,) + + # 处理没有search的情况 + if option == "whole texts": + # 分离文件名和扩展名，仅处理文件名部分 + filename, ext = os.path.splitext(input_text) + output_text = change_to if delete or change_to else filename + elif option == "prefix": + if delimiter: + parts = input_text.split(delimiter, 1) + else: + parts = re.split(r"[-_——#]", input_text, 1) + if delete: + output_text = parts[1] if len(parts) > 1 else input_text + else: + output_text = ( + f"{change_to}{delimiter if delimiter else '_'}{parts[1]}" + if len(parts) > 1 + else change_to + ) + elif option == "suffix": + if delimiter: + parts = input_text.rsplit(delimiter, 1) + else: + parts = re.split(r"[-_——#]", input_text) + if len(parts) > 1: + last_delimiter_index = max( + input_text.rfind(d) for d in ["-", "_", "——", "#"] + ) + parts = [ + input_text[:last_delimiter_index], + input_text[last_delimiter_index + 1 :], + ] + if delete: + output_text = parts[0] if len(parts) > 1 else input_text + else: + output_text = ( + f"{parts[0]}{delimiter if delimiter else '_'}{change_to}" + if len(parts) > 1 + else change_to + ) + else: + output_text = input_text + + return (output_text,) diff --git a/nodes/video_batch_process.py b/nodes/video_batch_process.py new file mode 100644 index 0000000..390eaab --- /dev/null +++ b/nodes/video_batch_process.py @@ -0,0 +1,650 @@ +import os +import torch +import numpy as np +import json +import folder_paths +from server import PromptServer +import re +from comfy_api.latest._util.video_types import VideoContainer, VideoCodec + + +class VideoBatchSaver: + INPUT_IS_LIST = True + FUNCTION = "save" + CATEGORY = "Batch Process" + RETURN_TYPES = () + RETURN_NAMES = () + OUTPUT_NODE = True + + SUPPORTED_FORMATS = VideoContainer.as_input() + SUPPORTED_CODECS = VideoCodec.as_input() + + @classmethod + def IS_CHANGED(cls, **kwargs): + # Always execute - return a unique value each time + import time + + return time.time() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "videos": ("VIDEO",), + "contents": ("STRING", {"forceInput": True}), + "output_path": ("STRING", {"default": ""}), + "filename_prefix": ("STRING", {"default": "VID"}), + "filename_delimiter": ("STRING", {"default": "_"}), + "filename_suffix": ("STRING", {"default": ""}), + "format": ( + cls.SUPPORTED_FORMATS, + {"default": VideoContainer.MP4.value}, + ), + "codec": ( + cls.SUPPORTED_CODECS, + {"default": VideoCodec.H264.value}, + ), + "filename_number_padding": ( + "INT", + {"default": 4, "min": 1, "max": 9, "step": 1}, + ), + "filename_number": ( + ["off", "start", "end"], + {"default": "end"}, + ), + "embeded_workflow": ("BOOLEAN", {"default": True}), + }, + "hidden": { + "node_id": "UNIQUE_ID", + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + def save( + self, + videos=None, + contents=None, + output_path="", + filename_prefix="VID", + filename_delimiter="_", + filename_suffix="", + format=VideoContainer.MP4.value, + codec=VideoCodec.H264.value, # 默认使用 H264 + filename_number_padding=4, + filename_number="end", + embeded_workflow=True, + node_id=None, + prompt=None, + extra_pnginfo=None, + ): + # 输入验证和处理 + format_value = self._get_first_or_default(format, VideoContainer.MP4.value) + codec_value = self._get_first_or_default(codec, VideoCodec.H264.value) + video_format = self._parse_video_container(format_value) + video_codec = self._parse_video_codec(codec_value) + extension = VideoContainer.get_extension(video_format) + filename_prefix = self._get_first_or_default(filename_prefix, "VID") + output_path = self._get_first_or_default(output_path, "") + filename_number_padding = self._get_first_or_default(filename_number_padding, 4) + filename_delimiter = self._get_first_or_default(filename_delimiter, "_") + filename_suffix = self._get_first_or_default(filename_suffix, "") + filename_suffix = filename_suffix.strip("'[]") + filename_number = self._get_first_or_default(filename_number, "end") + embeded_workflow = self._get_first_or_default(embeded_workflow, True) + + # 解析数字位置选项 + counter_start = filename_number in ["start"] + counter_end = filename_number in ["end"] + + if filename_number_padding < 1: + raise ValueError( + f"filename_number_padding must be at least 1, got {filename_number_padding}" + ) + + # 处理 videos 输入 + if videos is not None: + # 确保 videos 是列表 + if not isinstance(videos, list): + videos = [videos] + video_count = len(videos) + else: + videos = [] + video_count = len(contents) if contents is not None else 1 + + # 处理 output_path + output_path = self._normalize_input(output_path, video_count) + + # 处理 filename_prefix + filename_prefix = self._normalize_input(filename_prefix, video_count) + + # 处理文本输入 + contents = ( + self._normalize_input(contents, video_count) + if contents is not None + else None + ) + + # 准备输出目录 + output_dir = folder_paths.get_output_directory() + + # 批量保存 + saved_files = [] + for idx, (prefix, path) in enumerate(zip(filename_prefix, output_path)): + try: + final_output_path = self._get_output_path(output_dir, path) + os.makedirs(final_output_path, exist_ok=True) + + original_filename = os.path.splitext(os.path.basename(prefix))[0] + base_filename = self._generate_filename( + prefix=original_filename, + suffix=filename_suffix, + padding=filename_number_padding, + counter_start=counter_start, + counter_end=counter_end, + delimiter=filename_delimiter, + final_output_path=final_output_path, + extension=extension, + ) + + # 保存视频 + if videos and idx < len(videos): + full_path = os.path.join( + final_output_path, f"{base_filename}.{extension}" + ) + + # 使用 ComfyUI 的视频保存方法 + self._save_video_comfyui( + videos[idx], + full_path, + video_format=video_format, + codec=video_codec, + extension=extension, + embed_workflow=embeded_workflow, + prompt=prompt, + extra_pnginfo=extra_pnginfo, + ) + saved_files.append(full_path) + print(f"Saved video: {full_path}") + elif not videos: + # 如果没有视频，仍然生成文件路径（可能用于其他用途） + full_path = os.path.join( + final_output_path, f"{base_filename}.{extension}" + ) + saved_files.append(full_path) + print(f"Generated path: {full_path}") + + # 保存文本 + if contents and idx < len(contents): + content_path = os.path.join( + final_output_path, f"{base_filename}.txt" + ) + with open(content_path, "w", encoding="utf-8") as f: + f.write(str(contents[idx]).strip()) + saved_files.append(content_path) + print(f"Saved content: {content_path}") + + except Exception as e: + print(f"Error saving file {idx+1}: {str(e)}") + import traceback + + traceback.print_exc() + + self._update_progress(node_id, idx + 1, video_count) + + return () + + def _get_video_fps(self, video_obj): + """ + 从视频对象中提取 FPS，如果不可用则返回默认值 24 + """ + # 尝试从视频对象中获取 FPS + if hasattr(video_obj, "fps"): + fps = video_obj.fps + if fps is not None and fps > 0: + return float(fps) + if hasattr(video_obj, "get_fps"): + try: + fps = video_obj.get_fps() + if fps is not None and fps > 0: + return float(fps) + except: + pass + # 尝试从元数据或属性中获取 + if hasattr(video_obj, "__dict__"): + for attr in ["fps", "frame_rate", "rate"]: + if hasattr(video_obj, attr): + fps = getattr(video_obj, attr) + if fps is not None and fps > 0: + return float(fps) + # 默认返回 24 + return 24.0 + + def _save_video_comfyui( + self, + video_obj, + path, + video_format, + codec, + extension, + embed_workflow, + prompt, + extra_pnginfo, + ): + """ + 使用 ComfyUI 的视频保存功能 + """ + try: + # 准备元数据 + metadata = None + if embed_workflow: + metadata = {} + if prompt is not None: + metadata["prompt"] = prompt + if extra_pnginfo is not None: + if isinstance(extra_pnginfo, list): + merged_extra = {} + for item in extra_pnginfo: + if isinstance(item, dict): + merged_extra.update(item) + extra_pnginfo = merged_extra + if isinstance(extra_pnginfo, dict): + metadata.update(extra_pnginfo) + + # 检查视频对象是否有 save_to 方法 + if hasattr(video_obj, "save_to"): + try: + # 使用解析后的容器和编码设置（当前 ComfyUI 仅支持 MP4/H264） + video_obj.save_to( + path, + format=video_format, + codec=codec, + metadata=metadata, + ) + return # 成功保存，直接返回 + except Exception as save_error: + error_msg = str(save_error).lower() + # 检查是否是音频提取错误 + if "audio" in error_msg or "vhs failed" in error_msg: + print( + f"Warning: Audio extraction failed, attempting to save video without audio: {str(save_error)[:200]}" + ) + # 尝试直接提取视频帧并保存 + fps = self._get_video_fps(video_obj) + self._save_video_as_tensor(video_obj, path, extension, fps) + return + else: + # 其他错误，继续抛出 + raise + elif hasattr(video_obj, "save"): + try: + # 有些视频对象可能使用 save 方法而不是 save_to + video_obj.save(path) + return + except Exception as save_error: + error_msg = str(save_error).lower() + if "audio" in error_msg or "vhs failed" in error_msg: + print( + f"Warning: Audio extraction failed, attempting to save video without audio: {str(save_error)[:200]}" + ) + fps = self._get_video_fps(video_obj) + self._save_video_as_tensor(video_obj, path, extension, fps) + return + else: + raise + else: + # 如果没有保存方法，尝试转换为张量并保存 + fps = self._get_video_fps(video_obj) + self._save_video_as_tensor(video_obj, path, extension, fps) + + except Exception as e: + print(f"Error in ComfyUI video save: {str(e)[:200]}") + # 尝试使用更通用的保存方法 + try: + fps = self._get_video_fps(video_obj) + self._save_video_as_tensor(video_obj, path, extension, fps) + except Exception as tensor_error: + print(f"Error in tensor-based save: {str(tensor_error)[:200]}") + # 最后尝试通用方法 + fps = self._get_video_fps(video_obj) + self._save_video_generic(video_obj, path, extension, fps) + + def _save_video_as_tensor(self, video_obj, path, extension, fps): + """ + 将视频对象转换为张量并保存 + """ + try: + tensor = None + + # 尝试多种方式获取视频张量 + # 1. 优先访问 __components.video（ComfyUI 视频容器，避免触发音频提取） + if hasattr(video_obj, "__components"): + try: + components = video_obj.__components + if hasattr(components, "video"): + video_comp = components.video + # 尝试从 video 组件中获取张量 + if hasattr(video_comp, "tensor"): + tensor = video_comp.tensor + elif hasattr(video_comp, "get_tensor"): + tensor = video_comp.get_tensor() + elif hasattr(video_comp, "frames"): + tensor = video_comp.frames + elif isinstance(video_comp, torch.Tensor): + tensor = video_comp + elif hasattr(components, "frames"): + tensor = components.frames + except Exception as comp_error: + print( + f"Warning: Could not access __components: {str(comp_error)[:100]}" + ) + + # 2. 直接访问 tensor 属性 + if tensor is None and hasattr(video_obj, "tensor"): + tensor = video_obj.tensor + # 3. 调用 get_tensor 方法 + elif tensor is None and hasattr(video_obj, "get_tensor"): + tensor = video_obj.get_tensor() + # 4. 访问 frames 属性 + elif tensor is None and hasattr(video_obj, "frames"): + tensor = video_obj.frames + # 5. 访问 video 属性（ComfyUI 视频对象可能使用） + elif tensor is None and hasattr(video_obj, "video"): + tensor = video_obj.video + # 6. 尝试从 __dict__ 中查找 + elif tensor is None and hasattr(video_obj, "__dict__"): + video_dict = video_obj.__dict__ + for key in ["tensor", "frames", "video", "data"]: + if key in video_dict: + value = video_dict[key] + if isinstance(value, torch.Tensor): + tensor = value + break + # 7. 如果对象本身就是张量 + elif tensor is None and isinstance(video_obj, torch.Tensor): + tensor = video_obj + + # 如果找到了张量，保存它 + if tensor is not None and isinstance(tensor, torch.Tensor): + self._save_tensor_video(tensor, path, extension, fps) + else: + raise Exception( + f"Could not extract tensor from video object. Type: {type(video_obj)}, Available attributes: {[attr for attr in dir(video_obj) if not attr.startswith('__')][:15]}" + ) + + except Exception as e: + print(f"Error saving video as tensor: {str(e)}") + raise e + + def _save_tensor_video(self, tensor, path, extension, fps): + """ + 保存张量视频 + """ + try: + import imageio + + # 确保张量在 CPU 上 + if tensor.is_cuda: + tensor = tensor.cpu() + + # 处理不同的张量形状 + if tensor.dim() == 5: # (batch, frames, channels, height, width) + tensor = tensor[0] # 取第一个批次 + elif tensor.dim() == 4: # (frames, channels, height, width) + pass + else: + raise ValueError(f"Unsupported tensor shape: {tensor.shape}") + + # 转换为 numpy 数组 + frames = [] + for i in range(tensor.shape[0]): + frame = tensor[i] + + # 转换为 (height, width, channels) 格式 + if frame.shape[0] == 3: # (channels, height, width) + frame = frame.permute(1, 2, 0) + + # 确保是 uint8 类型 + if frame.dtype != torch.uint8: + frame = torch.clamp(frame * 255, 0, 255).to(torch.uint8) + + frame_np = frame.numpy() + frames.append(frame_np) + + # 保存视频 + imageio.mimwrite(path, frames, fps=fps, codec="libx264") + + except ImportError: + # 如果 imageio 不可用，尝试使用 OpenCV + self._save_tensor_video_opencv(tensor, path, extension, fps) + + def _save_tensor_video_opencv(self, tensor, path, extension, fps): + """ + 使用 OpenCV 保存张量视频 + """ + try: + import cv2 + + # 确保张量在 CPU 上 + if tensor.is_cuda: + tensor = tensor.cpu() + + # 处理不同的张量形状 + if tensor.dim() == 5: # (batch, frames, channels, height, width) + tensor = tensor[0] # 取第一个批次 + elif tensor.dim() == 4: # (frames, channels, height, width) + pass + else: + raise ValueError(f"Unsupported tensor shape: {tensor.shape}") + + # 获取视频尺寸 + if tensor.shape[1] == 3: # (frames, channels, height, width) + height, width = tensor.shape[2], tensor.shape[3] + else: + height, width = tensor.shape[1], tensor.shape[2] + + # 设置 FourCC 编码器 (使用 H264) + fourcc = cv2.VideoWriter_fourcc(*"avc1") # H264 编码 + out = cv2.VideoWriter(path, fourcc, fps, (width, height)) + + for i in range(tensor.shape[0]): + frame = tensor[i] + + # 转换为 (height, width, channels) 格式 + if frame.shape[0] == 3: # (channels, height, width) + frame = frame.permute(1, 2, 0) + + # 确保是 uint8 类型 + if frame.dtype != torch.uint8: + frame = torch.clamp(frame * 255, 0, 255).to(torch.uint8) + + frame_np = frame.numpy() + + # 转换为 BGR (OpenCV 使用 BGR) + if len(frame_np.shape) == 3 and frame_np.shape[2] == 3: + frame_np = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR) + + out.write(frame_np) + + out.release() + + except ImportError: + raise ImportError( + "Neither imageio nor opencv-python are available for video saving" + ) + + def _save_video_generic(self, video_obj, path, extension, fps): + """ + 通用视频保存方法 + """ + try: + # 首先尝试使用 _save_video_as_tensor，它已经包含了多种提取方法 + try: + self._save_video_as_tensor(video_obj, path, extension, fps) + return + except Exception as tensor_error: + print( + f"Tensor extraction failed, trying other methods: {str(tensor_error)[:100]}" + ) + + # 尝试各种可能的方法来获取视频数据 + if hasattr(video_obj, "_asdict"): + # 如果是 namedtuple 或类似结构 + try: + video_dict = video_obj._asdict() + if "frames" in video_dict: + video_fps = self._get_video_fps(video_obj) + self._save_tensor_video( + video_dict["frames"], path, extension, video_fps + ) + return + if "video" in video_dict: + video_fps = self._get_video_fps(video_obj) + self._save_tensor_video( + video_dict["video"], path, extension, video_fps + ) + return + except Exception: + pass + + if hasattr(video_obj, "__dict__"): + # 尝试从对象的属性中获取帧数据 + try: + video_dict = video_obj.__dict__ + for key, value in video_dict.items(): + if isinstance(value, torch.Tensor) and value.dim() >= 4: + video_fps = self._get_video_fps(video_obj) + self._save_tensor_video(value, path, extension, video_fps) + return + except Exception: + pass + + # 最后尝试直接调用对象的保存方法（不使用音频） + if hasattr(video_obj, "save"): + try: + video_obj.save(path) + return + except Exception: + pass + + raise Exception( + f"Could not find a way to save the video object. Type: {type(video_obj)}, Attributes: {[attr for attr in dir(video_obj) if not attr.startswith('_')][:10]}" + ) + + except Exception as e: + print(f"Error in generic video save: {str(e)}") + raise e + + def _get_first_or_default(self, value, default): + if isinstance(value, list): + for v in value: + if isinstance(v, type(default)): + return v + return default + return value if value is not None else default + + def _normalize_input(self, input_data, count): + if input_data is None: + return [None] * count + if not isinstance(input_data, list): + return [input_data] * count + return input_data + [input_data[-1]] * (count - len(input_data)) + + def _get_output_path(self, base_dir, user_path): + if not user_path or str(user_path).lower() in ["none", "."]: + return base_dir + return ( + os.path.join(base_dir, user_path) + if not os.path.isabs(str(user_path)) + else user_path + ) + + def _generate_filename( + self, + prefix, + suffix, + padding, + counter_start, + counter_end, + delimiter, + final_output_path, + extension, + ): + existing_files = [ + f + for f in os.listdir(final_output_path) + if f.endswith(f".{extension}") or f.endswith(".txt") + ] + + if suffix.strip(): + pattern_start = re.compile( + rf"^(\d+){re.escape(delimiter)}{re.escape(prefix)}{re.escape(delimiter)}{re.escape(suffix)}$" + ) + pattern_end = re.compile( + rf"^{re.escape(prefix)}{re.escape(delimiter)}{re.escape(suffix)}{re.escape(delimiter)}(\d+)$" + ) + else: + pattern_start = re.compile( + rf"^(\d+){re.escape(delimiter)}{re.escape(prefix)}$" + ) + pattern_end = re.compile( + rf"^{re.escape(prefix)}{re.escape(delimiter)}(\d+)$" + ) + + numbers = [] + for file in existing_files: + file_base = os.path.splitext(file)[0] + if counter_start and not counter_end: + match = pattern_start.match(file_base) + elif counter_end and not counter_start: + match = pattern_end.match(file_base) + elif counter_start and counter_end: + match_start = pattern_start.match(file_base) + match_end = pattern_end.match(file_base) + match = match_start or match_end + if match_start and match_end: + numbers.extend([int(match_start.group(1)), int(match_end.group(1))]) + continue + else: + match = None + if match: + numbers.append(int(match.group(1))) + + counter = max(numbers) + 1 if numbers else 1 + parts = [] + if counter_start: + parts.append(f"{counter:0{padding}d}") + parts.append(str(prefix)) + if suffix.strip(): + parts.append(str(suffix)) + if counter_end: + parts.append(f"{counter:0{padding}d}") + + return delimiter.join(parts) + + def _parse_video_container(self, value): + if isinstance(value, VideoContainer): + return value + if value is None: + return VideoContainer.MP4 + try: + return VideoContainer(str(value).lower()) + except ValueError as exc: + raise ValueError(f"Unsupported format: {value}") from exc + + def _parse_video_codec(self, value): + if isinstance(value, VideoCodec): + return value + if value is None: + return VideoCodec.H264 + try: + return VideoCodec(str(value).lower()) + except ValueError as exc: + raise ValueError(f"Unsupported codec: {value}") from exc + + def _update_progress(self, node_id, current, total): + if node_id: + PromptServer.instance.send_sync( + "progress", {"node": node_id, "value": current, "max": total} + ) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..74905ae --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "batch-process" +description = "A collection of nodes for batch processing texts, images, videos, and LoRAs in ComfyUI" +version = "1.2.0" +license = {file = "LICENSE"} + +[project.urls] +Repository = "https://github.com/Zar4X/ComfyUI-Batch-Process" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "zar4x" +DisplayName = "ComfyUI-Batch-Process" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d968b6e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +Pillow>=10.0.0 +torch>=2.0.0 + diff --git a/tests/asset_filter_js.test.mjs b/tests/asset_filter_js.test.mjs new file mode 100644 index 0000000..cc67ad8 --- /dev/null +++ b/tests/asset_filter_js.test.mjs @@ -0,0 +1,383 @@ +// Minimal litegraph stand-in used to exercise web/asset_filter.js without a +// browser. Run with: node tests/asset_filter_js.test.mjs +import assert from "node:assert"; +import { readFileSync } from "node:fs"; +import { fileURLToPath } from "node:url"; +import path from "node:path"; + +const here = path.dirname(fileURLToPath(import.meta.url)); +const source = readFileSync(path.join(here, "..", "web", "asset_filter.js"), "utf8"); + +// ─── Fake litegraph graph/node ─── +class FakeGraph { + constructor() { + this.links = new Map(); + this.nodes = new Map(); + this.nextLinkId = 1; + } + getNodeById(id) { + return this.nodes.get(id) ?? null; + } + add(node) { + node.graph = this; + this.nodes.set(node.id, node); + return node; + } +} + +class FakeNode { + constructor(id, graph) { + this.id = id; + this.inputs = []; + this.outputs = [{ name: "out", links: [] }]; + this.graph = graph; + graph?.add(this); + } + addInput(name, type) { + this.inputs.push({ name, type, link: null }); + this.onConnectionsChange?.(1, this.inputs.length - 1, false, null, null); + } + removeInput(index) { + const input = this.inputs[index]; + if (input?.link != null) this.graph.links.delete(input.link); + this.inputs.splice(index, 1); + this.onConnectionsChange?.(1, index, false, null, null); + } + // origin.connect(originSlot, targetNode, targetSlot) + connect(originSlot, target, targetSlot) { + const input = target.inputs[targetSlot]; + if (!input) return null; + if (input.link != null) this.graph.links.delete(input.link); + const id = this.graph.nextLinkId++; + this.graph.links.set(id, { + id, + origin_id: this.id, + origin_slot: originSlot, + target_id: target.id, + target_slot: targetSlot, + }); + input.link = id; + target.onConnectionsChange?.(1, targetSlot, true, this.graph.links.get(id), input); + return id; + } + setDirtyCanvas() {} +} + +// ─── Load the extension with a stubbed `app` ─── +let registered = null; +const app = { + registerExtension(ext) { + registered = ext; + }, + graph: { setDirtyCanvas() {} }, +}; + +const rafQueue = []; +globalThis.requestAnimationFrame = (fn) => { + rafQueue.push(fn); + return rafQueue.length; +}; +function flush() { + // Drain queued callbacks, including ones scheduled while draining. + let guard = 0; + while (rafQueue.length) { + if (++guard > 1000) throw new Error("requestAnimationFrame loop did not settle"); + rafQueue.shift()(); + } +} + +const moduleSource = source.replace( + /^import\s+\{\s*app\s*\}\s+from\s+["'][^"']+["'];?\s*$/m, + "" +); +const factory = new Function("app", `${moduleSource}\nreturn null;`); +factory(app); +assert.ok(registered, "extension should register itself"); + +// Apply the extension's prototype patches to our fake node class. The real +// backend declares 9 image + 3 video outputs, 0-based; the input caps are +// derived from these names. +const MAX_IMAGE = 9; +const MAX_VIDEO = 3; +const OUTPUT_NAMES = [ + ...Array.from({ length: MAX_IMAGE }, (_, i) => `image_${i}`), + ...Array.from({ length: MAX_VIDEO }, (_, i) => `video_${i}`), + "image_count", + "video_count", +]; +await registered.beforeRegisterNodeDef(FakeNode, { + name: "AssetFilter", + output_name: OUTPUT_NAMES, +}); + +function names(node) { + return node.inputs.map((i) => i.name); +} + +function makeFilterNode(graph) { + const node = new FakeNode(`f${graph.nextLinkId}_${graph.nodes.size}`, graph); + node.onAdded(); + flush(); + return node; +} + +// ─── Tests ─── +const tests = []; +function test(name, fn) { + tests.push([name, fn]); +} + +test("seeds one free image and one free video slot", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + assert.deepStrictEqual(names(node), ["image_0", "video_0"]); +}); + +test("connecting the free image slot appends exactly one new image slot", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + const src = new FakeNode("src", graph); + src.connect(0, node, 0); + flush(); + assert.deepStrictEqual(names(node), ["image_0", "image_1", "video_0"]); +}); + +test("repeated stabilize passes do not grow the input list", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + const src = new FakeNode("src", graph); + src.connect(0, node, 0); + flush(); + const before = names(node); + for (let i = 0; i < 20; i++) { + node.onConnectionsChange(1, 0, true, null, null); + flush(); + } + assert.deepStrictEqual(names(node), before, "layout must be stable"); +}); + +test("indices stay compact and never drift upward", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + const a = new FakeNode("a", graph); + const b = new FakeNode("b", graph); + a.connect(0, node, 0); + flush(); + b.connect(0, node, names(node).indexOf("image_1")); + flush(); + assert.deepStrictEqual(names(node), [ + "image_0", + "image_1", + "image_2", + "video_0", + ]); + // Disconnect the first image; the other is renumbered back to image_0. + const link = node.inputs[0].link; + graph.links.delete(link); + node.inputs[0].link = null; + node.onConnectionsChange(1, 0, false, null, null); + flush(); + assert.deepStrictEqual(names(node), ["image_0", "image_1", "video_0"]); +}); + +test("image slots stay grouped above video slots", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + const v = new FakeNode("v", graph); + v.connect(0, node, names(node).indexOf("video_0")); + flush(); + const i = new FakeNode("i", graph); + i.connect(0, node, names(node).indexOf("image_0")); + flush(); + assert.deepStrictEqual(names(node), [ + "image_0", + "image_1", + "video_0", + "video_1", + ]); + const order = names(node); + const lastImage = order.lastIndexOf("image_1"); + const firstVideo = order.indexOf("video_0"); + assert.ok(lastImage < firstVideo, "all image slots precede video slots"); +}); + +test("existing links survive a rebuild and point at the right origins", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + const img = new FakeNode("img", graph); + const vid = new FakeNode("vid", graph); + img.connect(0, node, names(node).indexOf("image_0")); + flush(); + vid.connect(0, node, names(node).indexOf("video_0")); + flush(); + const byName = Object.fromEntries( + node.inputs.map((input) => [input.name, input]) + ); + const imgLink = graph.links.get(byName["image_0"].link); + const vidLink = graph.links.get(byName["video_0"].link); + assert.strictEqual(imgLink?.origin_id, "img"); + assert.strictEqual(vidLink?.origin_id, "vid"); +}); + +test("multiple videos keep their own origins after renumbering", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + const v1 = new FakeNode("v1", graph); + const v2 = new FakeNode("v2", graph); + v1.connect(0, node, names(node).indexOf("video_0")); + flush(); + v2.connect(0, node, names(node).indexOf("video_1")); + flush(); + const byName = Object.fromEntries( + node.inputs.map((input) => [input.name, input]) + ); + assert.strictEqual(graph.links.get(byName["video_0"].link)?.origin_id, "v1"); + assert.strictEqual(graph.links.get(byName["video_1"].link)?.origin_id, "v2"); + assert.deepStrictEqual(names(node), [ + "image_0", + "video_0", + "video_1", + "video_2", + ]); +}); + +test("loading a legacy 1-based workflow normalizes to 0-based", () => { + const graph = new FakeGraph(); + const node = new FakeNode("loaded", graph); + // Simulate a workflow saved before the rename: drifted 1-based names and + // interleaving. Everything is renumbered from image_0 / video_0. + node.addInput("image_6116", "IMAGE"); + node.addInput("video_3", "*"); + node.addInput("image_2", "IMAGE"); + const src = new FakeNode("s", graph); + src.connect(0, node, 0); + src.connect(0, node, 1); + node.configure({}); + flush(); + assert.deepStrictEqual(names(node), [ + "image_0", + "image_1", + "video_0", + "video_1", + ]); +}); + +// Fill a zone to its cap by connecting the trailing free slot repeatedly. +function fillZone(graph, node, group, count) { + for (let i = 0; i < count; i++) { + const free = names(node).indexOf(`${group}_${i}`); + if (free < 0) break; + new FakeNode(`${group}_src${i}`, graph).connect(0, node, free); + flush(); + } +} + +test("image zone stops growing at the output count", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + fillZone(graph, node, "image", MAX_IMAGE); + const images = names(node).filter((n) => n.startsWith("image_")); + assert.strictEqual(images.length, MAX_IMAGE, "no slot past the last output"); + assert.strictEqual(images.at(-1), `image_${MAX_IMAGE - 1}`); + // Every slot is connected: a full zone keeps no trailing free slot. + const unconnected = node.inputs.filter( + (i) => i.name.startsWith("image_") && i.link == null + ); + assert.strictEqual(unconnected.length, 0); +}); + +test("video zone stops growing at its own smaller output count", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + fillZone(graph, node, "video", MAX_VIDEO); + const videos = names(node).filter((n) => n.startsWith("video_")); + assert.strictEqual(videos.length, MAX_VIDEO); + assert.strictEqual(videos.at(-1), `video_${MAX_VIDEO - 1}`); +}); + +test("a full zone stays stable under repeated stabilize passes", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + fillZone(graph, node, "image", MAX_IMAGE); + fillZone(graph, node, "video", MAX_VIDEO); + const before = names(node); + assert.strictEqual(before.length, MAX_IMAGE + MAX_VIDEO); + for (let i = 0; i < 20; i++) { + node.onConnectionsChange(1, 0, true, null, null); + flush(); + } + assert.deepStrictEqual(names(node), before); +}); + +test("freeing a slot in a full zone leaves the gap in place", () => { + const graph = new FakeGraph(); + const node = makeFilterNode(graph); + fillZone(graph, node, "image", MAX_IMAGE); + // Disconnect image_0. The zone is already at cap, so it must not renumber: + // outputs mirror inputs positionally, and shifting image_1 down to image_0 + // would silently move every remaining asset to a different output. + const link = node.inputs[0].link; + graph.links.delete(link); + node.inputs[0].link = null; + node.onConnectionsChange(1, 0, false, null, null); + flush(); + const images = node.inputs.filter((i) => i.name.startsWith("image_")); + assert.strictEqual(images.length, MAX_IMAGE, "zone stays at cap"); + const free = images.filter((i) => i.link == null); + assert.strictEqual(free.length, 1, "exactly one free slot"); + assert.strictEqual(free[0].name, "image_0", "the freed slot itself is free"); + // Every other slot keeps the origin it had before. + for (let i = 1; i < MAX_IMAGE; i++) { + const input = images[i]; + assert.strictEqual( + graph.links.get(input.link)?.origin_id, + `image_src${i}`, + `image_${i} kept its origin` + ); + } +}); + +test("an over-cap layout from a stale workflow is trimmed to the caps", () => { + const graph = new FakeGraph(); + const node = new FakeNode("stale", graph); + // A workflow saved before the caps existed: 12 image inputs, all connected. + const src = new FakeNode("s", graph); + for (let i = 1; i <= 12; i++) node.addInput(`image_${i}`, "IMAGE"); + for (let i = 0; i < 12; i++) src.connect(0, node, i); + node.configure({}); + flush(); + const images = names(node).filter((n) => n.startsWith("image_")); + assert.strictEqual(images.length, MAX_IMAGE); + assert.deepStrictEqual( + images, + Array.from({ length: MAX_IMAGE }, (_, i) => `image_${i}`) + ); +}); + +test("caps fall back to 9/3 when output names are unavailable", async () => { + // A defensive path: if nodeData carries no output names the extension must + // still cap rather than grow without bound. + class Fallback extends FakeNode {} + await registered.beforeRegisterNodeDef(Fallback, { name: "AssetFilter" }); + const graph = new FakeGraph(); + const node = new Fallback("fb", graph); + node.onAdded(); + flush(); + fillZone(graph, node, "video", 10); + const videos = names(node).filter((n) => n.startsWith("video_")); + assert.strictEqual(videos.length, 3); +}); + +let failed = 0; +for (const [name, fn] of tests) { + rafQueue.length = 0; + try { + await fn(); + console.log(`ok ${name}`); + } catch (error) { + failed++; + console.log(`FAIL ${name}\n ${error.message}`); + } +} +console.log(`\n${tests.length - failed}/${tests.length} passed`); +process.exit(failed ? 1 : 0); diff --git a/tests/test_asset_filter.py b/tests/test_asset_filter.py new file mode 100644 index 0000000..fe0bcff --- /dev/null +++ b/tests/test_asset_filter.py @@ -0,0 +1,377 @@ +import importlib.util +import os +import tempfile +import time +import unittest +from pathlib import Path + +import torch + + +MODULE_PATH = Path(__file__).resolve().parents[1] / "nodes" / "any_batch_process.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("any_batch_process_test", MODULE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def make_asset(path, age_hours): + Path(path).write_bytes(b"x") + if age_hours: + mtime = time.time() - age_hours * 3600.0 + os.utime(path, (mtime, mtime)) + + +class AssetFilterTests(unittest.TestCase): + def setUp(self): + self.module = load_module() + # Force the non-Comfy code path so tests never touch the real user + # directory state file. + self.module.folder_paths = None + self.module._memory_state.clear() + self.node = self.module.AssetFilter() + self.tmp = tempfile.mkdtemp(prefix="asset_filter_test_") + self.fresh_image = os.path.join(self.tmp, "fresh.png") + self.stale_image = os.path.join(self.tmp, "stale.png") + self.new_image = os.path.join(self.tmp, "new.png") + self.fresh_video = os.path.join(self.tmp, "fresh.mp4") + self.stale_video = os.path.join(self.tmp, "stale.mp4") + make_asset(self.fresh_image, 0) + make_asset(self.stale_image, 48) + make_asset(self.new_image, 0) + make_asset(self.fresh_video, 0) + make_asset(self.stale_video, 48) + self.image_a = torch.rand(1, 64, 64, 3) + self.image_b = torch.rand(1, 32, 32, 3) + + def filter(self, inputs, prompt, node_id="20", workflow_id="wf-test", **payloads): + prompt[node_id] = { + "class_type": "AssetFilter", + "inputs": {name: [source, 0] for name, source in inputs.items()}, + } + return self.node.filter_assets( + freshness_hours=24.0, + prompt=prompt, + node_id=node_id, + extra_pnginfo={"workflow": {"id": workflow_id}}, + **payloads, + ) + + def split(self, result): + """Unpack the flat output tuple into (images, videos, counts).""" + n_img = self.module._MAX_IMAGE_OUTPUTS + n_vid = self.module._MAX_VIDEO_OUTPUTS + self.assertEqual(len(result), n_img + n_vid + 2) + return ( + list(result[:n_img]), + list(result[n_img:n_img + n_vid]), + result[-2], + result[-1], + ) + + def kept(self, slots): + """Survivors in slot order, ignoring the None gaps between them.""" + return [value for value in slots if value is not None] + + def at(self, slots, index): + """The value on output _ (0-based, as shown in the UI).""" + return slots[index] + + def image_nodes(self, mapping): + return { + source: {"class_type": "LoadImage", "inputs": {"image": path}} + for source, path in mapping.items() + } + + def test_bootstrap_uses_mtime_and_positional_outputs(self): + prompt = self.image_nodes({"10": self.fresh_image, "11": self.stale_image}) + images, videos, image_count, video_count = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + self.assertEqual(image_count, 1) + self.assertEqual(video_count, 0) + kept = self.kept(images) + self.assertEqual(len(kept), 1) + self.assertTrue(torch.equal(kept[0], self.image_a)) + # No videos connected: every video output stays None so a reference + # consumer skips those slots instead of being pruned. + self.assertEqual(self.kept(videos), []) + + def test_bootstrap_all_stale_falls_back_to_pass_through(self): + prompt = self.image_nodes({"10": self.stale_image, "11": self.stale_image}) + images, _, image_count, _ = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + self.assertEqual(image_count, 2) + self.assertEqual(len(self.kept(images)), 2) + + def test_unchanged_name_is_filtered_on_next_run(self): + prompt = self.image_nodes({"10": self.fresh_image, "11": self.fresh_image}) + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + # Second run: slot 0 gets a new filename, slot 1 is untouched even + # though its file mtime is fresh. + prompt = self.image_nodes({"10": self.new_image, "11": self.fresh_image}) + images, _, image_count, _ = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + self.assertEqual(image_count, 1) + self.assertTrue(torch.equal(self.at(images, 0), self.image_a)) + self.assertIsNone(self.at(images, 1)) + + def test_rerun_without_changes_reuses_previous_selection(self): + prompt = self.image_nodes({"10": self.fresh_image, "11": self.stale_image}) + first = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + second = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + self.assertEqual(first[2], 1) + self.assertEqual(second[2], 1) + self.assertTrue(torch.equal(second[0][0], self.image_a)) + + def test_runtime_assets_without_source_file_are_kept(self): + prompt = { + "10": {"class_type": "EmptyImage", "inputs": {"width": 64, "height": 64}}, + } + images, _, image_count, _ = self.split( + self.filter({"image_0": "10"}, prompt, image_0=self.image_a) + ) + self.assertEqual(image_count, 1) + self.assertEqual(len(self.kept(images)), 1) + + def test_traces_through_intermediate_nodes(self): + prompt = { + "10": {"class_type": "VHS_LoadVideo", "inputs": {"video": self.fresh_video}}, + "11": {"class_type": "ImageScale", "inputs": {"image": ["10", 0]}}, + "12": {"class_type": "VHS_LoadVideo", "inputs": {"video": self.stale_video}}, + } + video_payload = object() + stale_payload = object() + _, videos, _, video_count = self.split( + self.filter( + {"video_0": "11", "video_1": "12"}, + prompt, + video_0=video_payload, + video_1=stale_payload, + ) + ) + self.assertEqual(video_count, 1) + self.assertIs(videos[0], video_payload) + + def test_non_media_widgets_are_not_mistaken_for_sources(self): + prompt = { + "10": {"class_type": "CLIPTextEncode", "inputs": {"text": "stale.png"}}, + } + _, videos, _, video_count = self.split( + self.filter({"video_0": "10"}, prompt, video_0=self.image_a) + ) + self.assertEqual(video_count, 1) + self.assertEqual(len(self.kept(videos)), 1) + + def test_original_sizes_are_preserved(self): + prompt = self.image_nodes({"10": self.fresh_image, "11": self.fresh_image}) + images, _, image_count, _ = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + self.assertEqual(image_count, 2) + # Each survivor keeps its own resolution: separate outputs mean no + # torch.cat, so no black-border padding. + self.assertEqual(tuple(images[0].shape), (1, 64, 64, 3)) + self.assertEqual(tuple(images[1].shape), (1, 32, 32, 3)) + + def test_flexible_inputs_accept_image_and_video_names(self): + inputs = self.module.AssetFilter.INPUT_TYPES()["optional"] + self.assertIn("image_0", inputs) + self.assertIn("image_3", inputs) + self.assertIn("video_12", inputs) + self.assertNotIn("anything_else", inputs) + self.assertEqual(inputs["image_0"][0], "IMAGE") + self.assertEqual(inputs["image_3"][0], "IMAGE") + self.assertEqual(inputs["video_12"][0], "*") + + def test_output_signature_matches_consumer_slot_counts(self): + cls = self.module.AssetFilter + n_img = self.module._MAX_IMAGE_OUTPUTS + n_vid = self.module._MAX_VIDEO_OUTPUTS + self.assertEqual(len(cls.RETURN_TYPES), n_img + n_vid + 2) + self.assertEqual(len(cls.RETURN_NAMES), len(cls.RETURN_TYPES)) + # 0-based, matching MiniMax H3's ref_image_0..8 / ref_video_0..2. + self.assertEqual(cls.RETURN_NAMES[0], "image_0") + self.assertEqual(cls.RETURN_NAMES[n_img - 1], f"image_{n_img - 1}") + self.assertEqual(cls.RETURN_NAMES[n_img], "video_0") + self.assertEqual(cls.RETURN_NAMES[-2:], ("image_count", "video_count")) + self.assertEqual(cls.RETURN_TYPES[:n_img], tuple(["IMAGE"] * n_img)) + self.assertEqual(cls.RETURN_TYPES[-2:], ("INT", "INT")) + # Each output carries one asset, so the engine must not treat them as + # lists to expand into separate executions. + self.assertFalse(hasattr(cls, "OUTPUT_IS_LIST")) + + def test_survivors_beyond_the_last_output_are_dropped(self): + n_img = self.module._MAX_IMAGE_OUTPUTS + count = n_img + 2 + paths = {} + payloads = {} + inputs = {} + for i in range(count): + path = os.path.join(self.tmp, f"surplus_{i}.png") + make_asset(path, 0) + paths[str(i)] = path + payloads[f"image_{i}"] = torch.rand(1, 8, 8, 3) + inputs[f"image_{i}"] = str(i) + images, _, image_count, _ = self.split( + self.filter(inputs, self.image_nodes(paths), **payloads) + ) + self.assertEqual(image_count, count) + # The tuple keeps its declared width; slots past the last output have + # nowhere to go, so those inputs are dropped. + self.assertEqual(len(images), n_img) + self.assertEqual(len(self.kept(images)), n_img) + for i in range(n_img): + self.assertTrue(torch.equal(self.at(images, i), payloads[f"image_{i}"])) + + def test_survivor_stays_on_its_own_output_slot(self): + # image_0 keeps its filename (dropped), image_1 changes (kept). The + # survivor must stay on output image_1 — compacting it onto image_0 + # would make a downstream preview show a different image than the slot + # it is wired to. + prompt = self.image_nodes({"10": self.fresh_image, "11": self.stale_image}) + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + prompt = self.image_nodes({"10": self.fresh_image, "11": self.new_image}) + images, _, image_count, _ = self.split( + self.filter( + {"image_0": "10", "image_1": "11"}, + prompt, + image_0=self.image_a, + image_1=self.image_b, + ) + ) + self.assertEqual(image_count, 1) + self.assertIsNone(self.at(images, 0)) + self.assertTrue(torch.equal(self.at(images, 1), self.image_b)) + + def test_holes_are_preserved_across_the_whole_zone(self): + # Survivors on slots 1 and 3 must appear on outputs 1 and 3, leaving + # 0 and 2 empty. Reference consumers skip the gaps. + paths = {} + payloads = {} + inputs = {} + for i in range(4): + path = os.path.join(self.tmp, f"hole_{i}.png") + make_asset(path, 0 if i % 2 == 1 else 48) + paths[str(i)] = path + payloads[f"image_{i}"] = torch.rand(1, 8 * (i + 1), 8 * (i + 1), 3) + inputs[f"image_{i}"] = str(i) + images, _, image_count, _ = self.split( + self.filter(inputs, self.image_nodes(paths), **payloads) + ) + self.assertEqual(image_count, 2) + self.assertIsNone(self.at(images, 0)) + self.assertIsNone(self.at(images, 2)) + self.assertTrue(torch.equal(self.at(images, 1), payloads["image_1"])) + self.assertTrue(torch.equal(self.at(images, 3), payloads["image_3"])) + + def test_legacy_1_based_inputs_still_execute(self): + # A workflow saved before the rename to 0-based names feeds image_1 / + # image_2; the backend must still accept them (the frontend renumbers + # on load, but raw execution must not break either). + prompt = self.image_nodes({"10": self.fresh_image, "11": self.fresh_image}) + images, _, image_count, _ = self.split( + self.filter( + {"image_1": "10", "image_2": "11"}, + prompt, + image_1=self.image_a, + image_2=self.image_b, + ) + ) + self.assertEqual(image_count, 2) + self.assertTrue(torch.equal(self.at(images, 1), self.image_a)) + self.assertTrue(torch.equal(self.at(images, 2), self.image_b)) + + def test_state_is_scoped_per_workflow_and_node(self): + prompt = self.image_nodes({"10": self.fresh_image}) + self.filter({"image_0": "10"}, prompt, image_0=self.image_a) + state = self.module._load_state() + self.assertIn("wf-test:20", state) + # A different node in a different workflow must not collide. + prompt_b = self.image_nodes({"10": self.fresh_image}) + self.filter( + {"image_0": "10"}, + prompt_b, + node_id="21", + workflow_id="wf-other", + image_0=self.image_a, + ) + state = self.module._load_state() + self.assertIn("wf-test:20", state) + self.assertIn("wf-other:21", state) + + def test_state_key_reads_nested_workflow_id(self): + self.assertEqual( + self.module._state_key({"workflow": {"id": "abc"}}, "7"), "abc:7" + ) + # Tolerate a flat shape and a missing id without raising. + self.assertEqual(self.module._state_key({"id": "flat"}, "7"), "flat:7") + self.assertEqual(self.module._state_key(None, "7"), "None:7") + + def test_unused_outputs_are_none_not_blockers(self): + # Only a video is connected, so every image output has nothing to emit. + # None is required here: a reference consumer skips None slots, whereas + # an ExecutionBlocker on any one slot prunes the whole consumer node. + prompt = { + "10": {"class_type": "VHS_LoadVideo", "inputs": {"video": self.fresh_video}}, + } + images, videos, image_count, video_count = self.split( + self.filter({"video_0": "10"}, prompt, video_0=self.image_a) + ) + self.assertEqual(image_count, 0) + self.assertEqual(video_count, 1) + self.assertTrue(all(value is None for value in images)) + self.assertEqual(len(self.kept(videos)), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/web/asset_filter.js b/web/asset_filter.js new file mode 100644 index 0000000..c5c1828 --- /dev/null +++ b/web/asset_filter.js @@ -0,0 +1,220 @@ +import { app } from "../../../scripts/app.js"; + +// AssetFilter grows its input list on the fly: connecting the last free +// image/video slot appends a new one, and unused free slots are collapsed. +// Image inputs always stay grouped above video inputs; a newly added image +// slot appears directly under the existing image slots (same for videos). +// Input names must stay "image_N" / "video_N" so the backend can parse them. +// +// Each zone stops growing at the number of matching outputs the backend +// declares, since an input past the last output has nowhere to go. The caps are +// read from the node's own output names rather than hardcoded, so changing +// _MAX_IMAGE_OUTPUTS/_MAX_VIDEO_OUTPUTS in Python moves the frontend limit too. +// +// litegraph has no "insert input in the middle" API, so when the layout +// deviates from the canonical order the node is rebuilt: inputs are removed, +// re-added in canonical order, and links are restored through the official +// origin.connect() API. + +const IMAGE_TYPE = "IMAGE"; +const VIDEO_TYPE = "*"; +const NAME_RE = /^(image|video)_(\d+)$/; +const GROUPS = ["image", "video"]; +// Used only if the output names cannot be read (they always can in ComfyUI). +const FALLBACK_LIMITS = { image: 9, video: 3 }; +const limits = { ...FALLBACK_LIMITS }; + +// The backend declares one image_N/video_N output per usable input slot; the +// names are 0-based, so the cap is one past the highest output index. +function readLimits(nodeData) { + const names = nodeData?.output_name ?? nodeData?.output ?? []; + const found = { image: 0, video: 0 }; + for (const name of names) { + const parsed = typeof name === "string" ? parseName(name) : null; + if (parsed) { + found[parsed.group] = Math.max( + found[parsed.group], + parsed.index + 1 + ); + } + } + for (const group of GROUPS) { + if (found[group] > 0) limits[group] = found[group]; + } +} + +function limitFor(group) { + return limits[group] ?? FALLBACK_LIMITS[group]; +} + +function typeForGroup(group) { + return group === "image" ? IMAGE_TYPE : VIDEO_TYPE; +} + +function parseName(name) { + const match = NAME_RE.exec(name); + return match + ? { group: match[1], index: parseInt(match[2], 10) } + : null; +} + +function groupInputs(node, group) { + return (node.inputs || []).filter( + (input) => parseName(input.name)?.group === group + ); +} + +// Canonical layout: image zone, then video zone. Each zone holds its +// connected inputs (renumbered 0..N-1) followed by exactly one free slot, +// capped at that zone's output count. Names are always reassigned from the +// position, so indices stay compact and can never drift upward across rebuilds. +function targetLayout(node) { + const target = []; + for (const group of GROUPS) { + const limit = limitFor(group); + const connectedCount = Math.min( + groupInputs(node, group).filter((input) => input.link != null).length, + limit + ); + // One trailing free slot to grow into, unless the zone is full. + const total = Math.min(connectedCount + 1, limit); + for (let i = 0; i < total; i++) { + target.push(`${group}_${i}`); + } + } + return target; +} + +// Names the connected inputs are expected to carry after a rebuild, in the +// same order captureConnections() returns them. Connections beyond the zone's +// cap get no name and are dropped. +function targetNamesForConnections(node) { + const names = []; + for (const group of GROUPS) { + const limit = limitFor(group); + const connected = groupInputs(node, group).filter( + (input) => input.link != null + ); + connected.forEach((_, i) => { + names.push(i < limit ? `${group}_${i}` : null); + }); + } + return names; +} + +function layoutMatches(node, target) { + const inputs = node.inputs || []; + return ( + inputs.length === target.length && + target.every((name, i) => inputs[i].name === name) + ); +} + +// Connected inputs grouped image-first, each in ascending slot order. The +// order here must match targetNamesForConnections(). +function captureConnections(node) { + const connections = []; + for (const group of GROUPS) { + for (const input of groupInputs(node, group)) { + if (input.link == null) continue; + const links = node.graph?.links; + const link = links?.get ? links.get(input.link) : links?.[input.link]; + if (!link) continue; + const origin = node.graph?.getNodeById?.(link.origin_id); + if (origin) { + connections.push({ origin, slot: link.origin_slot }); + } + } + } + return connections; +} + +function rebuild(node, target) { + const connections = captureConnections(node); + const names = targetNamesForConnections(node); + while (node.inputs.length) { + node.removeInput(node.inputs.length - 1); + } + for (const name of target) { + node.addInput(name, typeForGroup(parseName(name).group)); + } + // Reconnect by canonical name; connections and names are index-aligned. + // A null name means the connection sat past its zone's cap, so it is + // deliberately left disconnected. + connections.forEach((connection, i) => { + if (names[i] == null) return; + const index = node.inputs.findIndex( + (input) => input.name === names[i] + ); + if (index >= 0) { + connection.origin.connect(connection.slot, node, index); + } + }); + node.setDirtyCanvas?.(true, true); +} + +function stabilize(node) { + if (!node.inputs || node.removed || node._assetFilterBusy) return; + const target = targetLayout(node); + if (layoutMatches(node, target)) return; + // removeInput()/connect() below fire onConnectionsChange synchronously; + // the guard keeps those callbacks from re-entering the rebuild. + node._assetFilterBusy = true; + try { + rebuild(node, target); + } finally { + node._assetFilterBusy = false; + } +} + +const scheduled = new WeakSet(); +function scheduleStabilize(node) { + if (scheduled.has(node)) return; + scheduled.add(node); + requestAnimationFrame(() => { + scheduled.delete(node); + stabilize(node); + }); +} + +app.registerExtension({ + name: "Zar4X.AssetFilter", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== "AssetFilter") return; + + readLimits(nodeData); + + const origOnAdded = nodeType.prototype.onAdded; + nodeType.prototype.onAdded = function () { + const result = origOnAdded?.apply(this, arguments); + // Fresh node: seed one free image and one free video slot. + for (const group of GROUPS) { + if (groupInputs(this, group).length === 0) { + this.addInput(`${group}_0`, typeForGroup(group)); + } + } + return result; + }; + + const origConfigure = nodeType.prototype.configure; + nodeType.prototype.configure = function (info) { + const result = origConfigure?.apply(this, arguments); + scheduleStabilize(this); + return result; + }; + + const origOnConnectionsChange = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function ( + type, + slotIndex, + isConnected, + linkInfo, + ioSlot + ) { + const result = origOnConnectionsChange?.apply(this, arguments); + // LiteGraph NODE_INPUT === 1: only react to input connections. + if (type === 1) scheduleStabilize(this); + return result; + }; + }, +});