From c7b623b2f49b847fbfaf35a8bca14bc0fd82f47d Mon Sep 17 00:00:00 2001 From: Alan Burns Date: Sat, 11 Oct 2025 17:16:29 -0500 Subject: [PATCH] Initial commit --- LICENSE | 21 ++ README.md | 40 ++ __init__.py | 102 +++++ auto_adjust.py | 284 ++++++++++++++ auto_color_match.py | 222 +++++++++++ crop_utils.py | 339 +++++++++++++++++ docs/AutoAdjustNode.md | 33 ++ docs/AutoColorConfigNode.md | 21 ++ docs/AutoCropBorders.md | 31 ++ docs/AutoWBColorMatch.md | 33 ++ docs/CropImageByMargins.md | 24 ++ docs/CropMaskByMargins.md | 24 ++ docs/FilenameAppendSuffix.md | 24 ++ docs/FitAspectHeadSafe.md | 32 ++ docs/FluxResolutionPrepare.md | 32 ++ docs/GpuSyncCacheClear.md | 18 + docs/ImageMegapixelSelector.md | 22 ++ docs/LoadImageCombined.md | 29 ++ docs/MQBBoxMin.md | 26 ++ docs/MultiPromptNode.md | 21 ++ docs/OtherConfigNode.md | 28 ++ docs/OutpaintConfigNode.md | 26 ++ docs/OutpaintPaddingComputeNode.md | 28 ++ docs/PortraitUtils_PairedImageLoader.md | 27 ++ docs/SeedVR2Prep.md | 21 ++ docs/StitchByMask.md | 31 ++ docs/WorkflowConfig.md | 50 +++ filename_append_suffix.py | 76 ++++ flux_resolution_prepare.py | 285 ++++++++++++++ gpu_cache_clear.py | 99 +++++ image_megapixel_selector.py | 185 +++++++++ load_image_combined.py | 286 ++++++++++++++ multiprompt.py | 47 +++ other_config.py | 95 +++++ outpaint_config.py | 236 ++++++++++++ paired_image_loader.py | 481 ++++++++++++++++++++++++ smart_crop.py | 433 +++++++++++++++++++++ stitch_image_by_mask.py | 213 +++++++++++ workflow_config.py | 329 ++++++++++++++++ 39 files changed, 4354 insertions(+) create mode 100644 LICENSE create mode 100644 README.md create mode 100644 __init__.py create mode 100644 auto_adjust.py create mode 100644 auto_color_match.py create mode 100644 crop_utils.py create mode 100644 docs/AutoAdjustNode.md create mode 100644 docs/AutoColorConfigNode.md create mode 100644 docs/AutoCropBorders.md create mode 100644 docs/AutoWBColorMatch.md create mode 100644 docs/CropImageByMargins.md create mode 100644 docs/CropMaskByMargins.md create mode 100644 docs/FilenameAppendSuffix.md create mode 100644 docs/FitAspectHeadSafe.md create mode 100644 docs/FluxResolutionPrepare.md create mode 100644 docs/GpuSyncCacheClear.md create mode 100644 docs/ImageMegapixelSelector.md create mode 100644 docs/LoadImageCombined.md create mode 100644 docs/MQBBoxMin.md create mode 100644 docs/MultiPromptNode.md create mode 100644 docs/OtherConfigNode.md create mode 100644 docs/OutpaintConfigNode.md create mode 100644 docs/OutpaintPaddingComputeNode.md create mode 100644 docs/PortraitUtils_PairedImageLoader.md create mode 100644 docs/SeedVR2Prep.md create mode 100644 docs/StitchByMask.md create mode 100644 docs/WorkflowConfig.md create mode 100644 filename_append_suffix.py create mode 100644 flux_resolution_prepare.py create mode 100644 gpu_cache_clear.py create mode 100644 image_megapixel_selector.py create mode 100644 load_image_combined.py create mode 100644 multiprompt.py create mode 100644 other_config.py create mode 100644 outpaint_config.py create mode 100644 paired_image_loader.py create mode 100644 smart_crop.py create mode 100644 stitch_image_by_mask.py create mode 100644 workflow_config.py diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..07d70ff --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 PortraitUtils contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..e1a248e --- /dev/null +++ b/README.md @@ -0,0 +1,40 @@ +# PortraitUtils Node Suite + +PortraitUtils is a focused collection of ComfyUI nodes designed for pipelines using real potographic images of people, where image sizes are not standardized as they are in text-to-image generation. The suite streamlines prep work such as cropping, resolution targeting, metadata management, and side-by-side QA so you can concentrate on creative iteration instead of utility wiring. These nodes should work on any type of image, but some of them are designed for and tested on humans, particularly the smart cropping nodes. + +## Node Index + +Each entry links to a short, task-oriented reference page (coming soon) for deeper usage notes and parameter breakdowns. + +### Image Adjustments & Preparation +- `AutoAdjustNode` – One-click levels, tone, and color balancing. [Docs](docs/AutoAdjustNode.md) +- `AutoWBColorMatch` – White-balances and colour matches against a reference image using several algorithms. [Docs](docs/AutoWBColorMatch.md) +- `AutoColorConfigNode` – Lightweight toggle bundle for sharing Auto Adjust presets. [Docs](docs/AutoColorConfigNode.md) +- `FluxResolutionPrepare` – Crops and resizes to the nearest Flux-friendly resolution, with optional pre-upscale. [Docs](docs/FluxResolutionPrepare.md) +- `ImageMegapixelSelector` – Chooses image sizes by target megapixels to match downstream model limits. [Docs](docs/ImageMegapixelSelector.md) + +### Cropping, Framing, and Stitching +- `AutoCropBorders` – Removes uniform borders using adaptive fuzzy detection. GPU acceleration is currently broken, so leave that off for now. [Docs](docs/AutoCropBorders.md) +- `CropImageByMargins` – Manual margin-based cropper for RGB images. [Docs](docs/CropImageByMargins.md) +- `CropMaskByMargins` – Margin cropper designed for MASK tensors. [Docs](docs/CropMaskByMargins.md) +- `FitAspectHeadSafe` – Rescales and pads while preserving face-safe regions for portrait framing. [Docs](docs/FitAspectHeadSafe.md) +- `MQBBoxMin` – Utility for merging bounding boxes while enforcing minimum sizes. [Docs](docs/MQBBoxMin.md) +- `StitchByMask` – Composites foreground and background images using mask-driven blending. [Docs](docs/StitchByMask.md) + +### Prompting & Workflow Configuration +- `MultiPromptNode` – Cycles through prompt variants with per-slot weights. [Docs](docs/MultiPromptNode.md) +- `WorkflowConfig` – Centralized key/value store for workflow-wide settings. [Docs](docs/WorkflowConfig.md) +- `OtherConfigNode` – Additional config toggles for orchestration scripts. [Docs](docs/OtherConfigNode.md) +- `OutpaintConfigNode` – Packs outpainting preferences (steps, padding, sampler choices). [Docs](docs/OutpaintConfigNode.md) +- `OutpaintPaddingComputeNode` – Calculates padding offsets to feed outpaint chains. [Docs](docs/OutpaintPaddingComputeNode.md) + +### I/O & Utility Helpers +- `LoadImageCombined` – Hybrid single/batch image loader with auto-advance support. [Docs](docs/LoadImageCombined.md) +- `PortraitUtils_PairedImageLoader` – Loads synchronized source/output pairs for A/B comparison workflows. [Docs](docs/PortraitUtils_PairedImageLoader.md) +- `FilenameAppendSuffix` – Adds suffixes to filenames while respecting multi-extension cases. [Docs](docs/FilenameAppendSuffix.md) +- `GpuSyncCacheClear` – Forces a CUDA sync and empties selected caches to avoid VRAM buildup. [Docs](docs/GpuSyncCacheClear.md) +- `SeedVR2Prep` – Convenience helper for SeedVR2/VRAM-heavy pipelines (cache priming & cleanup). [Docs](docs/SeedVR2Prep.md) + +## License + +This project is released under the [MIT License](LICENSE). diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..04f61fb --- /dev/null +++ b/__init__.py @@ -0,0 +1,102 @@ +""" +Custom node registry that aggregates the standalone node modules in this package. + +Intentionally excludes the legacy `ImageUtils.py` bundle; only the extracted +modules are imported and registered here. +""" + +from .auto_adjust import ( + NODE_CLASS_MAPPINGS as AUTO_ADJUST_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as AUTO_ADJUST_DISPLAY_MAP, +) +from .outpaint_config import ( + NODE_CLASS_MAPPINGS as OUTPAINT_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as OUTPAINT_DISPLAY_MAP, +) +from .multiprompt import ( + NODE_CLASS_MAPPINGS as MULTIPROMPT_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as MULTIPROMPT_DISPLAY_MAP, +) +from .stitch_image_by_mask import ( + NODE_CLASS_MAPPINGS as STITCH_BY_MASK_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as STITCH_BY_MASK_DISPLAY_MAP, +) +from .crop_utils import ( + NODE_CLASS_MAPPINGS as CROP_UTILS_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as CROP_UTILS_DISPLAY_MAP, +) +from .gpu_cache_clear import ( + NODE_CLASS_MAPPINGS as GPU_CACHE_CLEAR_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as GPU_CACHE_CLEAR_DISPLAY_MAP, +) +from .filename_append_suffix import ( + NODE_CLASS_MAPPINGS as FILENAME_APPEND_SUFFIX_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as FILENAME_APPEND_SUFFIX_DISPLAY_MAP, +) +from .auto_color_match import ( + NODE_CLASS_MAPPINGS as AUTO_COLOR_MATCH_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as AUTO_COLOR_MATCH_DISPLAY_MAP, +) +from .image_megapixel_selector import ( + NODE_CLASS_MAPPINGS as IMAGE_MPX_SELECTOR_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as IMAGE_MPX_SELECTOR_DISPLAY_MAP, +) +from .flux_resolution_prepare import ( + NODE_CLASS_MAPPINGS as FLUX_PREP_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as FLUX_PREP_DISPLAY_MAP, +) +from .workflow_config import ( + NODE_CLASS_MAPPINGS as WORKFLOW_CONFIG_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as WORKFLOW_CONFIG_DISPLAY_MAP, +) +from .smart_crop import ( + NODE_CLASS_MAPPINGS as SMART_CROP_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as SMART_CROP_DISPLAY_MAP, +) +from .paired_image_loader import ( + NODE_CLASS_MAPPINGS as PAIRED_LOADER_CLASS_MAP, + NODE_DISPLAY_NAME_MAPPINGS as PAIRED_LOADER_DISPLAY_MAP, +) +from .load_image_combined import LoadImageCombined + + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# Merge the mappings exposed by the extracted modules. +NODE_CLASS_MAPPINGS.update(AUTO_ADJUST_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(OUTPAINT_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(MULTIPROMPT_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(STITCH_BY_MASK_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(CROP_UTILS_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(GPU_CACHE_CLEAR_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(FILENAME_APPEND_SUFFIX_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(AUTO_COLOR_MATCH_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(IMAGE_MPX_SELECTOR_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(FLUX_PREP_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(WORKFLOW_CONFIG_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(SMART_CROP_CLASS_MAP) +NODE_CLASS_MAPPINGS.update(PAIRED_LOADER_CLASS_MAP) + +NODE_DISPLAY_NAME_MAPPINGS.update(AUTO_ADJUST_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(OUTPAINT_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(MULTIPROMPT_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(STITCH_BY_MASK_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(CROP_UTILS_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(GPU_CACHE_CLEAR_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(FILENAME_APPEND_SUFFIX_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(AUTO_COLOR_MATCH_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_MPX_SELECTOR_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(FLUX_PREP_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(WORKFLOW_CONFIG_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(SMART_CROP_DISPLAY_MAP) +NODE_DISPLAY_NAME_MAPPINGS.update(PAIRED_LOADER_DISPLAY_MAP) + +# Register the load-image node which does not expose pre-built mappings. +NODE_CLASS_MAPPINGS["LoadImageCombined"] = LoadImageCombined +NODE_DISPLAY_NAME_MAPPINGS["LoadImageCombined"] = "Load Image (Combined)" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", +] diff --git a/auto_adjust.py b/auto_adjust.py new file mode 100644 index 0000000..d04c967 --- /dev/null +++ b/auto_adjust.py @@ -0,0 +1,284 @@ +# auto_adjust.py (extracted from ImageUtils.py) +# Minimal standalone module containing AutoAdjustNode and AutoColorConfigNode +import torch + + +def _auto_color(rgb, use_color, snap_midtones): + if not use_color: + return rgb + + ycbcr = _rgb_to_ycbcr(rgb) + Y, Cb, Cr = ycbcr[..., 0:1], ycbcr[..., 1:2], ycbcr[..., 2:3] + + if snap_midtones: + B, H, W, _ = rgb.shape + mask = (Y > 0.2) & (Y < 0.8) + def median_masked(ch, msk): + flat = ch[msk.expand_as(ch)].view(ch.shape[0], -1) if msk.any() else ch.view(ch.shape[0], -1) + return torch.quantile(flat, 0.5, dim=1, keepdim=True).view(ch.shape[0], 1, 1, 1) + cb_mid = median_masked(Cb, mask) + cr_mid = median_masked(Cr, mask) + Cb = Cb - cb_mid + Cr = Cr - cr_mid + else: + cb_mean = torch.mean(Cb, dim=(1, 2), keepdim=True) + cr_mean = torch.mean(Cr, dim=(1, 2), keepdim=True) + Cb = Cb - cb_mean + Cr = Cr - cr_mean + + ycbcr = torch.cat([Y, Cb, Cr], dim=-1) + out = _ycbcr_to_rgb(ycbcr) + return _clamp01(out) + + +def _auto_levels(rgb, use_levels, shadow_pct, highlight_pct, gamma_norm, precision_mode): + if not use_levels: + return rgb + + Y = _luma(rgb) + if shadow_pct == 0.0 and highlight_pct == 0.0: + low, high = torch.min(Y, dim=(1, 2), keepdim=True).values, torch.max(Y, dim=(1, 2), keepdim=True).values + else: + if shadow_pct + highlight_pct > 0: + if shadow_pct > 0: + low = (_percentiles_exact(Y, shadow_pct / 100.0) if precision_mode == "Exact" else _percentiles_hist(Y, shadow_pct / 100.0)) + else: + low = torch.min(Y, dim=(1, 2), keepdim=True).values + if highlight_pct > 0: + high = (_percentiles_exact(Y, 1.0 - (highlight_pct / 100.0)) if precision_mode == "Exact" else _percentiles_hist(Y, 1.0 - (highlight_pct / 100.0))) + else: + high = torch.max(Y, dim=(1, 2), keepdim=True).values + else: + # exact path if clips are both zero (already handled) or if needed + low = _percentiles_exact(Y, shadow_pct / 100.0) + high = _percentiles_exact(Y, 1.0 - (highlight_pct / 100.0)) + + stretched = _linear_stretch_scalar(rgb, low, high) + + if gamma_norm: + Ys = _luma(stretched) + mid_mask = (Ys > 0.25) & (Ys < 0.75) + if mid_mask.any(): + flat = Ys[mid_mask] + median = torch.quantile(flat.view(1, -1), 0.5).item() + if 0.35 < median < 0.65: + # steer midtones gently toward 0.5 + gamma = 0.0 + if median != 0.0 and median != 1.0: + # gamma solving: 0.5 = median ** gamma => gamma = log(0.5)/log(median) + gamma = torch.log(torch.tensor(0.5)) / torch.log(torch.tensor(median)) + gamma = float(torch.clamp(gamma, 0.85, 1.15)) + if gamma != 0.0: + stretched = torch.clamp(stretched, 1e-6, 1.0) ** gamma + + return _clamp01(stretched) + + +def _auto_tone(rgb, use_tone, mode, shadow_pct, highlight_pct, precision_mode): + if not use_tone: + return rgb + + if mode == "Per-channel": + if shadow_pct == 0.0 and highlight_pct == 0.0: + low = torch.min(rgb, dim=(1, 2), keepdim=True).values + high = torch.max(rgb, dim=(1, 2), keepdim=True).values + else: + def pct(ch, p): + if p == 0.0: + return torch.min(ch, dim=(1, 2), keepdim=True).values + return (_percentiles_exact(ch, p / 100.0) if precision_mode == "Exact" else _percentiles_hist(ch, p / 100.0)) + low = torch.cat([pct(rgb[..., i:i+1], shadow_pct) for i in range(3)], dim=-1) + def pct_hi(ch, p): + if p == 0.0: + return torch.max(ch, dim=(1, 2), keepdim=True).values + return (_percentiles_exact(ch, 1.0 - (p / 100.0)) if precision_mode == "Exact" else _percentiles_hist(ch, 1.0 - (p / 100.0))) + high = torch.cat([pct_hi(rgb[..., i:i+1], highlight_pct) for i in range(3)], dim=-1) + stretched = _linear_stretch(rgb, low, high) + return _clamp01(stretched) + + # Monochromatic: stretch only luma + ycbcr = _rgb_to_ycbcr(rgb) + Y = ycbcr[..., 0:1] + if shadow_pct == 0.0 and highlight_pct == 0.0: + low = torch.min(Y, dim=(1, 2), keepdim=True).values + high = torch.max(Y, dim=(1, 2), keepdim=True).values + else: + low = (_percentiles_exact(Y, shadow_pct / 100.0) if precision_mode == "Exact" else _percentiles_hist(Y, shadow_pct / 100.0)) if shadow_pct > 0.0 else torch.min(Y, dim=(1, 2), keepdim=True).values + high = (_percentiles_exact(Y, 1.0 - (highlight_pct / 100.0)) if precision_mode == "Exact" else _percentiles_hist(Y, 1.0 - (highlight_pct / 100.0))) if highlight_pct > 0.0 else torch.max(Y, dim=(1, 2), keepdim=True).values + + Ys = _linear_stretch_scalar(Y, low, high) + ycbcr = torch.cat([Ys, ycbcr[..., 1:2], ycbcr[..., 2:3]], dim=-1) + out = _ycbcr_to_rgb(ycbcr) + return _clamp01(out) + + +def _clamp01(x): + return torch.clamp(x, 0.0, 1.0) + + +def _linear_stretch(x, low, high): + eps = 1e-6 + return torch.clamp((x - low) / torch.clamp(high - low, min=eps), 0.0, 1.0) + + +def _linear_stretch_scalar(x, low, high): + eps = 1e-6 + return torch.clamp((x - low) / torch.clamp(high - low, min=eps), 0.0, 1.0) + + +def _luma(rgb): + return (0.2126 * rgb[..., 0:1]) + (0.7152 * rgb[..., 1:2]) + (0.0722 * rgb[..., 2:3]) + + +def _percentiles_exact(x, q): + flat = _to_1d(x) + return torch.quantile(flat, q, dim=1, keepdim=True).view(x.shape[0], 1, 1, x.shape[-1]) + + +def _percentiles_hist(x, q): + B, H, W, C = x.shape + flat = x.view(B, -1, C) + # 256-bin histogram approximation per channel + bins = 256 + edges = torch.linspace(0, 1, steps=bins + 1, device=x.device, dtype=x.dtype) + centers = (edges[:-1] + edges[1:]) / 2.0 + # build cdf per batch/channel + # counts shape: [B, C, bins] + counts = [] + for c in range(C): + ch = flat[..., c] + # bucketize + idx = torch.clamp((ch * (bins - 1)).long(), 0, bins - 1) + one_hot = torch.nn.functional.one_hot(idx, num_classes=bins).to(x.dtype) + counts.append(one_hot.sum(dim=1)) # [B, bins] + counts = torch.stack(counts, dim=1) # [B, C, bins] + cdf = torch.cumsum(counts, dim=-1) + cdf = cdf / cdf[..., -1:].clamp(min=1.0) + # find first bin where cdf >= q + q_idx = torch.argmax((cdf >= q).to(torch.int64), dim=-1) # [B, C] + vals = centers[q_idx] # [B, C] + return vals.view(B, 1, 1, C) + + +def _rgb_to_ycbcr(rgb): + # Expect rgb in [0,1] + r, g, b = rgb[..., 0:1], rgb[..., 1:2], rgb[..., 2:3] + Y = 0.299 * r + 0.587 * g + 0.114 * b + Cb = 0.5643 * (b - Y) + Cr = 0.7132 * (r - Y) + return torch.cat([Y, Cb, Cr], dim=-1) + + +def _to_1d(x): + B, H, W, C = x.shape + return x.view(B, -1, C) + + +def _ycbcr_to_rgb(ycbcr): + Y, Cb, Cr = ycbcr[..., 0:1], ycbcr[..., 1:2], ycbcr[..., 2:3] + r = Y + 1.4020 * Cr + g = Y - 0.3441 * Cb - 0.7141 * Cr + b = Y + 1.7720 * Cb + return torch.cat([r, g, b], dim=-1) + + +class AutoAdjustNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "precision": (["Histogram (fast)", "Exact"], {"default": "Exact"}), + + "auto_levels": ("BOOLEAN", {"default": True}), + "levels_shadow_clip_pct": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5.0, "step": 0.01}), + "levels_highlight_clip_pct": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5.0, "step": 0.01}), + "levels_gamma_normalize": ("BOOLEAN", {"default": False}), + + "auto_tone": ("BOOLEAN", {"default": True}), + "tone_mode": (["Per-channel", "Monochromatic"], {"default": "Per-channel"}), + "tone_shadow_clip_pct": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5.0, "step": 0.01}), + "tone_highlight_clip_pct": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5.0, "step": 0.01}), + + "auto_color": ("BOOLEAN", {"default": True}), + "snap_neutral_midtones": ("BOOLEAN", {"default": False}), + + "flip_horizontal": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply" + CATEGORY = "image/adjustment" + + def apply( + self, + image, + precision, + + auto_levels, levels_shadow_clip_pct, levels_highlight_clip_pct, levels_gamma_normalize, + auto_tone, tone_mode, tone_shadow_clip_pct, tone_highlight_clip_pct, + auto_color, snap_neutral_midtones, + + flip_horizontal, + ): + # Expect image in [0,1], shape [B,H,W,C] with C=3 or 4 (alpha is discarded) + assert image.ndim == 4 and image.shape[-1] in (3, 4), "Expected IMAGE tensor [B,H,W,3 or 4]" + rgb = image[..., :3].to(dtype=torch.float32) + rgb = _clamp01(rgb) + + # Precision switch only affects percentile path inside helpers + # (Histogram (fast) uses histogram-based percentiles; Exact uses torch.quantile) + # In this extraction we use the same helpers; selection logic is inside them where applicable. + if precision not in ("Histogram (fast)", "Exact"): + precision = "Histogram (fast)" + + # Auto Levels (global scalar stretch via luma) + rgb = _auto_levels(rgb, auto_levels, levels_shadow_clip_pct, levels_highlight_clip_pct, levels_gamma_normalize, precision) + + # Auto Tone (per-channel or monochrome stretch) + rgb = _auto_tone(rgb, auto_tone, tone_mode, tone_shadow_clip_pct, tone_highlight_clip_pct, precision) + + # Auto Color (neutral balance in YCbCr) + rgb = _auto_color(rgb, auto_color, snap_neutral_midtones) + + if flip_horizontal: + rgb = torch.flip(rgb, dims=[2]) + + out = _clamp01(rgb) + return (out,) + +# ============================================================ +# AutoColor Config (as before) +# ============================================================ + +class AutoColorConfigNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "auto_levels": ("BOOLEAN", {"default": False}), + "auto_tone": ("BOOLEAN", {"default": False}), + "auto_color": ("BOOLEAN", {"default": False}), + "flip_horizontal": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("BOOLEAN", "BOOLEAN", "BOOLEAN", "BOOLEAN") + RETURN_NAMES = ("auto_levels", "auto_tone", "auto_color", "flip_horizontal") + FUNCTION = "apply" + CATEGORY = "config" + + def apply(self, auto_levels, auto_tone, auto_color, flip_horizontal): + return (auto_levels, auto_tone, auto_color, flip_horizontal) + + +NODE_CLASS_MAPPINGS = { + "AutoAdjustNode": AutoAdjustNode, + "AutoColorConfigNode": AutoColorConfigNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AutoAdjustNode": "Auto Adjust (Levels/Tone/Color)", + "AutoColorConfigNode": "AutoColor Config", +} diff --git a/auto_color_match.py b/auto_color_match.py new file mode 100644 index 0000000..c1de3aa --- /dev/null +++ b/auto_color_match.py @@ -0,0 +1,222 @@ +from __future__ import annotations +import torch +import torch.nn.functional as F + +# --------------------------- +# sRGB <-> Linear helpers +# --------------------------- +def srgb_to_linear(x): + a = 0.055 + return torch.where(x <= 0.04045, x / 12.92, ((x + a) / (1 + a)).pow(2.4)) + +def linear_to_srgb(x): + a = 0.055 + return torch.where(x <= 0.0031308, x * 12.92, (1 + a) * x.pow(1/2.4) - a) + +# --------------------------- +# RGB <-> XYZ <-> Lab (D65) +# --------------------------- +# Matrices for D65 sRGB +_M_RGB2XYZ = torch.tensor([[0.4124564, 0.3575761, 0.1804375], + [0.2126729, 0.7151522, 0.0721750], + [0.0193339, 0.1191920, 0.9503041]], dtype=torch.float32) +_M_XYZ2RGB = torch.tensor([[ 3.2404542, -1.5371385, -0.4985314], + [-0.9692660, 1.8760108, 0.0415560], + [ 0.0556434, -0.2040259, 1.0572252]], dtype=torch.float32) +# D65 white point +_Xn, _Yn, _Zn = 0.95047, 1.00000, 1.08883 + +def rgb_to_lab(rgb): # [B,H,W,3] in 0..1 sRGB + B,H,W,C = rgb.shape + M = _M_RGB2XYZ.to(rgb.device, rgb.dtype) + # sRGB -> linear -> XYZ + lin = srgb_to_linear(rgb) + xyz = torch.einsum('bhwc,cd->bhwd', lin, M) + # Normalize by white point + x = xyz[...,0] / _Xn + y = xyz[...,1] / _Yn + z = xyz[...,2] / _Zn + eps = 216/24389 + kappa = 24389/27 + + def f(t): + return torch.where(t > eps, t.pow(1/3), (kappa * t + 16) / 116) + + fx, fy, fz = f(x), f(y), f(z) + L = 116*fy - 16 + a = 500*(fx - fy) + b = 200*(fy - fz) + return torch.stack([L, a, b], dim=-1) # [B,H,W,3] + +def lab_to_rgb(lab): # [B,H,W,3] + L, a, b = lab[...,0], lab[...,1], lab[...,2] + fy = (L + 16)/116 + fx = fy + a/500 + fz = fy - b/200 + + eps = 216/24389 + kappa = 24389/27 + + def finv(ft): + t3 = ft**3 + return torch.where(t3 > eps, t3, (116*ft - 16)/kappa) + + xr, yr, zr = finv(fx), finv(fy), finv(fz) + X = xr * _Xn + Y = yr * _Yn + Z = zr * _Zn + xyz = torch.stack([X,Y,Z], dim=-1) + M = _M_XYZ2RGB.to(lab.device, lab.dtype) + lin = torch.einsum('bhwc,cd->bhwd', xyz, M) + srgb = linear_to_srgb(lin).clamp(0,1) + return srgb + +# --------------------------- +# Utility +# --------------------------- +def to_bhwc(x): + if x.dim()==3: x = x.unsqueeze(0) + assert x.dim()==4 and x.shape[-1] in (3,4), f"Expected [B,H,W,3/4], got {tuple(x.shape)}" + if x.shape[-1]==4: x = x[...,:3] + return x.float().clamp(0,1) + +def resize_bhwc(x, h, w, mode='bilinear'): + return F.interpolate(x.permute(0,3,1,2), size=(h,w), mode=mode, align_corners=False).permute(0,2,3,1) + +def luminance(img): + # Rec.709 luma + return 0.2126*img[...,0] + 0.7152*img[...,1] + 0.0722*img[...,2] + +def per_image_mean_std(x): + # x: [B,H,W,3] + B = x.shape[0] + flat = x.view(B, -1, 3) + mean = flat.mean(dim=1, keepdim=True) + std = flat.std (dim=1, unbiased=False, keepdim=True).clamp_min(1e-6) + return mean, std + +# --------------------------- +# Methods +# --------------------------- +def wb_grayworld(img): + # scale channels so mean becomes gray + B = img.shape[0] + mean = img.view(B,-1,3).mean(dim=1) # [B,3] + gray = mean.mean(dim=1, keepdim=True) # [B,1] + gains = (gray / mean).view(B,1,1,3) + out = (img * gains).clamp(0,1) + return out + +def wb_highlight(img, percentile=95.0): + # use brightest-percent luminance pixels as "white patch" + B,H,W,_ = img.shape + Y = luminance(img).view(B, -1) + thr = torch.quantile(Y, percentile/100.0, dim=1, keepdim=True) # [B,1] + mask = (Y >= thr).view(B,H,W,1) + # avoid empty mask + eps = 1e-6 + sel = img * mask + count = mask.sum(dim=(1,2,3), keepdim=True).clamp_min(1.0) + mean_sel = sel.sum(dim=(1,2,3), keepdim=True) / count # [B,1,1,3] + target_white = torch.ones_like(mean_sel) * 0.95 # bring selected whites near 95% to avoid clipping + gains = (target_white / mean_sel).clamp(0.5, 2.0) + return (img * gains).clamp(0,1) + +def reinhard_match(src, ref, l_only=False): + # Convert to Lab, match mean/std + src_lab = rgb_to_lab(src) + ref_lab = rgb_to_lab(ref) + src_mean, src_std = per_image_mean_std(src_lab) + ref_mean, ref_std = per_image_mean_std(ref_lab) + + if l_only: + # match L only + L = (src_lab[...,0:1] - src_mean[...,0:1]) / src_std[...,0:1] * ref_std[...,0:1] + ref_mean[...,0:1] + out_lab = torch.cat([L, src_lab[...,1:],], dim=-1) + else: + out_lab = (src_lab - src_mean) / src_std * ref_std + ref_mean + + out = lab_to_rgb(out_lab).clamp(0,1) + return out + +# --------------------------- +# Comfy node +# --------------------------- +class AutoWBColorMatch: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "reference": ("IMAGE",), + "method": ([ + "wb_grayworld", + "wb_highlight", + "reinhard_lab", + "lab_l_only", + "wb_highlight+reinhard", + ], {"default": "wb_highlight+reinhard"}), + "percentile": ("FLOAT", {"default": 95.0, "min": 80.0, "max": 99.9, "step": 0.1}), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "clip_gamut": ("BOOLEAN", {"default": True}), + "force_size": ("BOOLEAN", {"default": False}), + "target_width": ("INT", {"default": 1440, "min": 16, "max": 8192, "step": 1}), + "target_height": ("INT", {"default": 1080, "min": 16, "max": 8192, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "run" + CATEGORY = "Image/Color" + + def run(self, image, reference, method="wb_highlight+reinhard", + percentile=95.0, strength=1.0, clip_gamut=True, + force_size=False, target_width=1440, target_height=1080): + + src = to_bhwc(image) + ref = to_bhwc(reference) + + if force_size: + th, tw = target_height, target_width + src_small = resize_bhwc(src, th, tw) + ref_small = resize_bhwc(ref, th, tw) + else: + src_small, ref_small = src, ref + + # 1) white balance / base correction + if method in ("wb_grayworld",): + base = wb_grayworld(src_small) + elif method in ("wb_highlight", "wb_highlight+reinhard"): + base = wb_highlight(src_small, percentile=float(percentile)) + elif method in ("reinhard_lab","lab_l_only"): + base = src_small + else: + base = src_small + + # 2) color match + if method == "reinhard_lab": + matched = reinhard_match(base, ref_small, l_only=False) + elif method == "lab_l_only": + matched = reinhard_match(base, ref_small, l_only=True) + elif method == "wb_highlight+reinhard": + matched = reinhard_match(base, ref_small, l_only=False) + else: + matched = base + + # If we resized for stats, reapply the transform back to original resolution + if force_size and (src.shape[1]!=matched.shape[1] or src.shape[2]!=matched.shape[2]): + matched = resize_bhwc(matched, src.shape[1], src.shape[2]) + + # 3) blend strength + out = (1 - strength) * src + strength * matched + if clip_gamut: + out = out.clamp(0,1) + + return (out,) + +NODE_CLASS_MAPPINGS = { + "AutoWBColorMatch": AutoWBColorMatch, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "AutoWBColorMatch": "Auto White-Balance + Color Match", +} diff --git a/crop_utils.py b/crop_utils.py new file mode 100644 index 0000000..15a6763 --- /dev/null +++ b/crop_utils.py @@ -0,0 +1,339 @@ +import torch +import torch.nn.functional as _iutils_F + +_iutils_torch = torch + + +def _iutils_ensure_bhwc_rgb(image): + t = image + if not isinstance(t, _iutils_torch.Tensor): + t = _iutils_torch.tensor(t) + t = t.to(_iutils_torch.float32) + if t.dim() == 3: + t = t.unsqueeze(0) + # Expect [B,H,W,C] + if t.shape[-1] == 3: + return t.clamp(0.0, 1.0) + elif t.shape[-1] > 3: + return t[..., :3].clamp(0.0, 1.0) + elif t.shape[-1] == 1: + return t.repeat(1, 1, 1, 3).clamp(0.0, 1.0) + else: + return t.clamp(0.0, 1.0) + + +def _iutils_dilate(mask): + k = _iutils_torch.ones((1, 1, 3, 3), device=mask.device) + return (_iutils_F.conv2d(mask, k, padding=1) > 0).float() + + +def _iutils_rgb2y(img_bhwc): + r, g, b = img_bhwc[..., 0], img_bhwc[..., 1], img_bhwc[..., 2] + return (0.2126 * r + 0.7152 * g + 0.0722 * b).unsqueeze(-1) + + +def _iutils_region_grow(img, seed_mask, border_rgb, thr_ch, max_iter=4096): + B, H, W, Cp = img.shape + dist = (img - border_rgb).abs().sum(dim=-1, keepdim=True).permute(0, 3, 1, 2) + thr_l1 = thr_ch.sum(dim=-1, keepdim=True).permute(0, 3, 1, 2) + eligible = (dist <= thr_l1).float() + cur = seed_mask.clone() + for _ in range(max_iter): + grown = _iutils_dilate(cur) * eligible + if _iutils_torch.equal(grown, cur): + break + cur = grown + return cur + + +def _iutils_bbox_from_mask(nonborder, pad): + B, _, H, W = nonborder.shape + boxes = [] + for b in range(B): + m = nonborder[b, 0] + ys = _iutils_torch.where(m.any(dim=1))[0] + xs = _iutils_torch.where(m.any(dim=0))[0] + if len(ys) == 0 or len(xs) == 0: + l, t, w, h = 0, 0, W, H + else: + top, bottom = int(ys[0].item()), int(ys[-1].item()) + left, right = int(xs[0].item()), int(xs[-1].item()) + l = max(0, left - pad) + t = max(0, top - pad) + r = min(W, right + 1 + pad) + btm = min(H, bottom + 1 + pad) + w = max(1, r - l) + h = max(1, btm - t) + boxes.append((l, t, w, h)) + return boxes + + +class AutoCropBorders: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "fuzz_mode": (["percent", "adaptive"], {"default": "adaptive"}), + "fuzz_percent": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 50.0, "step": 0.1}), + "adaptive_k": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), + "edge_margin_px": ("INT", {"default": 4, "min": 1, "max": 64, "step": 1}), + "pad_px": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}), + "use_luma_only": ("BOOLEAN", {"default": False}), + "max_growth_iter": ("INT", {"default": 4096, "min": 64, "max": 16384, "step": 64}), + "return_border_mask": ("BOOLEAN", {"default": False}), + "use_gpu": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "INT", "MASK") + RETURN_NAMES = ("image", "left", "top", "width", "height", "border_mask") + FUNCTION = "run" + CATEGORY = "image/transform" + + def run( + self, + image, + fuzz_mode="adaptive", + fuzz_percent=5.0, + adaptive_k=2.0, + edge_margin_px=4, + pad_px=0, + use_luma_only=False, + max_growth_iter=4096, + return_border_mask=False, + use_gpu=True, + ): + + img = _iutils_ensure_bhwc_rgb(image) + target_dev = ( + _iutils_torch.device("cuda") + if (use_gpu and _iutils_torch.cuda.is_available()) + else img.device + ) + if img.device != target_dev: + img = img.to(target_dev, non_blocking=True) + + B, H, W, C = img.shape + work = _iutils_rgb2y(img) if use_luma_only else img + Cp = work.shape[-1] + + m = int(edge_margin_px) + seed = _iutils_torch.zeros( + (B, 1, H, W), device=img.device, dtype=_iutils_torch.float32 + ) + seed[:, :, :m, :] = 1.0 + seed[:, :, -m:, :] = 1.0 + seed[:, :, :, :m] = 1.0 + seed[:, :, :, -m:] = 1.0 + + top = work[:, :m, :, :].reshape(B, -1, Cp) + bottom = work[:, -m:, :, :].reshape(B, -1, Cp) + left = work[:, :, :m, :].reshape(B, -1, Cp) + right = work[:, :, -m:, :].reshape(B, -1, Cp) + samples = _iutils_torch.cat([top, bottom, left, right], dim=1) + + med = samples.median(dim=1).values + absdev = (samples - med.unsqueeze(1)).abs() + mad = absdev.median(dim=1).values + 1e-6 + + base = _iutils_torch.full_like(med, float(fuzz_percent) / 100.0) + if str(fuzz_mode) == "adaptive": + base = base + float(adaptive_k) * mad + thr_ch = base.clamp(0.0, 1.0) + + border_rgb = med.view(B, 1, 1, Cp) + thr_ch_4d = thr_ch.view(B, 1, 1, Cp) + + grown = _iutils_region_grow(work, seed, border_rgb, thr_ch_4d, int(max_growth_iter)) + nonborder = (1.0 - grown).clamp(0, 1) + boxes = _iutils_bbox_from_mask(nonborder, int(pad_px)) + + # Crop + crops = [] + for b, (l, t, wc, hc) in enumerate(boxes): + crops.append(img[b : b + 1, t : t + hc, l : l + wc, :]) + max_h = max(c.shape[1] for c in crops) + max_w = max(c.shape[2] for c in crops) + out = [] + for c in crops: + pad_h = max_h - c.shape[1] + pad_w = max_w - c.shape[2] + out.append(_iutils_F.pad(c, (0, 0, 0, pad_w, 0, pad_h))) + out_img = _iutils_torch.cat(out, dim=0).clamp(0, 1) + + border_mask = grown.permute(0, 2, 3, 1) + if (border_mask.shape[1] != out_img.shape[1]) or (border_mask.shape[2] != out_img.shape[2]): + border_mask = ( + _iutils_F.interpolate( + border_mask.permute(0, 3, 1, 2), + size=(out_img.shape[1], out_img.shape[2]), + mode="nearest", + ) + .permute(0, 2, 3, 1) + ) + + # Enforce even dimensions (safe) + H_out, W_out = out_img.shape[1], out_img.shape[2] + + # Width + if (W_out % 2) == 1: + if W_out > 1: + out_img = out_img[:, :, : W_out - 1, :] + border_mask = border_mask[:, :, : W_out - 1, :] + W_out = W_out - 1 + else: + # pad one column by repeating last column + pad_col = out_img[:, :, -1:, :] + out_img = _iutils_torch.cat([out_img, pad_col], dim=2) + pad_col_m = border_mask[:, :, -1:, :] + border_mask = _iutils_torch.cat([border_mask, pad_col_m], dim=2) + W_out = 2 + + # Height + if (H_out % 2) == 1: + if H_out > 1: + out_img = out_img[:, : H_out - 1, :, :] + border_mask = border_mask[:, : H_out - 1, :, :] + H_out = H_out - 1 + else: + # pad one row by repeating last row + pad_row = out_img[:, -1:, :, :] + out_img = _iutils_torch.cat([out_img, pad_row], dim=1) + pad_row_m = border_mask[:, -1:, :, :] + border_mask = _iutils_torch.cat([border_mask, pad_row_m], dim=1) + H_out = 2 + + # Reported crop box (single image): clamp to >=2 + if B == 1: + l, t, wc, hc = boxes[0] + wc = wc if (wc % 2 == 0) else (wc - 1 if wc > 1 else 2) + hc = hc if (hc % 2 == 0) else (hc - 1 if hc > 1 else 2) + else: + l = t = wc = hc = 0 + + return ( + out_img, + int(l), + int(t), + int(wc), + int(hc), + border_mask if return_border_mask else _iutils_torch.zeros_like(border_mask), + ) + + +def _to_bhwc(x: torch.Tensor) -> torch.Tensor: + # Accept [H,W,C] or [B,H,W,C] + if x.dim() == 3: + x = x.unsqueeze(0) + if x.dim() != 4: + raise ValueError(f"Expected [B,H,W,C] or [H,W,C], got {tuple(x.shape)}") + return x + + +def _crop_bhwc(img: torch.Tensor, left: int, top: int, right: int, bottom: int, multiple: int) -> torch.Tensor: + B, H, W, C = img.shape + + # Clamp margins to safe range + left = max(0, min(left, W - 1)) + right = max(0, min(right, W - 1)) + top = max(0, min(top, H - 1)) + bottom = max(0, min(bottom, H - 1)) + + # Compute crop box + x0 = left + y0 = top + x1 = max(x0 + 1, W - right) # ensure at least 1 px + y1 = max(y0 + 1, H - bottom) + + # Optionally snap to a multiple (e.g., 64) + if multiple > 1: + # Round down width/height to a multiple, keeping the top-left fixed + new_w = ((x1 - x0) // multiple) * multiple + new_h = ((y1 - y0) // multiple) * multiple + new_w = max(1, new_w) + new_h = max(1, new_h) + x1 = min(W, x0 + new_w) + y1 = min(H, y0 + new_h) + + return img[:, y0:y1, x0:x1, :] + + +class CropImageByMargins: + """Crop an IMAGE tensor by pixel margins (left/top/right/bottom).""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "left_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "top_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "right_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "bottom_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "snap_multiple": ("INT", {"default": 1, "min": 1, "max": 512, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "crop" + CATEGORY = "Image/Transform" + + def crop(self, image, left_px, top_px, right_px, bottom_px, snap_multiple=1): + x = _to_bhwc(image) + y = _crop_bhwc(x, left_px, top_px, right_px, bottom_px, snap_multiple) + return (y.clamp(0.0, 1.0),) + + +class CropMaskByMargins: + """Same as above but for MASK input and output (single-channel).""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "left_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "top_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "right_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "bottom_px": ("INT", {"default": 0, "min": 0, "max": 16384}), + "snap_multiple": ("INT", {"default": 1, "min": 1, "max": 512, "step": 1}), + } + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "crop" + CATEGORY = "Mask/Transform" + + def crop(self, mask, left_px, top_px, right_px, bottom_px, snap_multiple=1): + # Normalize mask to [B,H,W,1] + m = mask + if m.dim() == 2: # [H,W] + m = m.unsqueeze(0).unsqueeze(-1) + elif m.dim() == 3: + if m.shape[0] == 1: # [1,H,W] -> [1,H,W,1] + m = m.permute(1, 2, 0).unsqueeze(0) + elif m.shape[-1] == 1: # [H,W,1] -> [1,H,W,1] + m = m.unsqueeze(0) + else: + raise ValueError(f"Unexpected MASK shape {tuple(m.shape)}") + elif m.dim() == 4 and m.shape[-1] == 1: + pass + else: + raise ValueError(f"Unexpected MASK shape {tuple(m.shape)}") + + y = _crop_bhwc(m, left_px, top_px, right_px, bottom_px, snap_multiple) + return (y.clamp(0.0, 1.0),) + + +NODE_CLASS_MAPPINGS = { + "AutoCropBorders": AutoCropBorders, + "CropImageByMargins": CropImageByMargins, + "CropMaskByMargins": CropMaskByMargins, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AutoCropBorders": "Auto-Crop Borders (fuzzy)", + "CropImageByMargins": "Crop by Margins (Image)", + "CropMaskByMargins": "Crop by Margins (Mask)", +} diff --git a/docs/AutoAdjustNode.md b/docs/AutoAdjustNode.md new file mode 100644 index 0000000..6e72b79 --- /dev/null +++ b/docs/AutoAdjustNode.md @@ -0,0 +1,33 @@ +# AutoAdjustNode + +## Overview +`AutoAdjustNode` performs a chained series of global tone corrections on an RGB image. It can optionally apply automatic levels, per-channel or monochrome tone stretching, neutral colour balancing in YCbCr, and a final horizontal flip. The implementation assumes floating-point tensors in `[0, 1]` and always returns RGB output. + +## Inputs +- `image` (`IMAGE`): Source image tensor shaped `[B, H, W, C]` (C must be 3 or 4; the alpha channel is discarded). +- `precision` (`STRING`): Selects percentile estimation mode for level/tone calculations. `"Histogram (fast)"` approximates quantiles via 256-bin histograms, while `"Exact"` uses `torch.quantile`. +- `auto_levels` (`BOOLEAN`, default `True`): Enables global min/max remapping based on luma percentiles. +- `levels_shadow_clip_pct` (`FLOAT`, default `0.1`): Percentage of darkest pixels to clip when computing the low anchor for levels. +- `levels_highlight_clip_pct` (`FLOAT`, default `0.1`): Percentage of brightest pixels to clip when computing the high anchor. +- `levels_gamma_normalize` (`BOOLEAN`, default `False`): When enabled, nudges midtones toward 50% grey if the median luma falls between 35% and 65%. +- `auto_tone` (`BOOLEAN`, default `True`): Toggles the tone mapping stage. +- `tone_mode` (`STRING`, default `"Per-channel"`): `"Per-channel"` stretches each channel independently; `"Monochromatic"` applies a single stretch derived from luma. +- `tone_shadow_clip_pct` (`FLOAT`, default `0.1`): Shadow percentile used by the tone stage. +- `tone_highlight_clip_pct` (`FLOAT`, default `0.1`): Highlight percentile used by the tone stage. +- `auto_color` (`BOOLEAN`, default `True`): Enables neutral balance in YCbCr space (centres Cb/Cr components). +- `snap_neutral_midtones` (`BOOLEAN`, default `False`): When `True`, only midtone pixels (luma in 0.2–0.8) contribute to the Cb/Cr statistics, reducing highlight/shadow bias. +- `flip_horizontal` (`BOOLEAN`, default `False`): Horizontally mirrors the result. + +## Outputs +- `IMAGE`: The processed RGB tensor with all three stages (levels, tone, colour) applied in order, plus any optional flip. + +## Processing Notes +- All stages operate on float32 tensors in `[0, 1]`; the node casts incoming data accordingly and clamps the final output. +- Levels and tone stages guard against degenerate ranges (when high equals low) and fall back to minimal stretching. +- The colour balance stage works in YCbCr and uses either masked medians or simple means depending on `snap_neutral_midtones`. +- Alpha channels are intentionally dropped so downstream nodes always receive RGB tensors. + +## Tips +- Theoretically, use `"Histogram (fast)"` when performance matters or when feeding large batches; switch to `"Exact"` for reproducibility-sensitive work. In actual practice, `"Exact"` is usually faster, so there's no reason not to use it. +- Keep `levels_shadow_clip_pct`/`levels_highlight_clip_pct` at zero when you need strict dynamic range preservation (e.g., HDR preparation). +- Combine this node with `AutoColorConfigNode` to share toggle states across multiple branches of a workflow.*** End Patch diff --git a/docs/AutoColorConfigNode.md b/docs/AutoColorConfigNode.md new file mode 100644 index 0000000..8ab7e1b --- /dev/null +++ b/docs/AutoColorConfigNode.md @@ -0,0 +1,21 @@ +# AutoColorConfigNode + +## Overview +`AutoColorConfigNode` bundles the toggle values used by `AutoAdjustNode` into a reusable configuration socket. It is designed to keep multiple processing branches in sync without duplicating widgets, and it emits plain boolean values that can be re-injected into other nodes. + +## Inputs +- `auto_levels` (`BOOLEAN`, default `False`): Desired state for the levels stage. +- `auto_tone` (`BOOLEAN`, default `False`): Desired state for the tone stage. +- `auto_color` (`BOOLEAN`, default `False`): Desired state for the colour balance stage. +- `flip_horizontal` (`BOOLEAN`, default `False`): Indicates whether downstream nodes should flip their outputs horizontally. + +## Outputs +- `auto_levels` (`BOOLEAN`) +- `auto_tone` (`BOOLEAN`) +- `auto_color` (`BOOLEAN`) +- `flip_horizontal` (`BOOLEAN`) + +## Usage Notes +- Route the outputs directly into the corresponding sockets on one or more `AutoAdjustNode` instances. +- Because all outputs are plain booleans, the node is safe to use with ComfyUI’s conditional routing or scripting nodes. +- The node stores no internal state; toggling any widget simply passes the value straight through.*** End Patch diff --git a/docs/AutoCropBorders.md b/docs/AutoCropBorders.md new file mode 100644 index 0000000..7843930 --- /dev/null +++ b/docs/AutoCropBorders.md @@ -0,0 +1,31 @@ +# AutoCropBorders + +## Overview +`AutoCropBorders` trims uniform borders from an image by sampling edge pixels, estimating their colour statistics, and region-growing inward until the content changes. It optionally returns the detected border mask and reports the crop rectangle, making it useful for automated subject isolation or pre-processing scanned artwork. + +## Inputs +- `image` (`IMAGE`): Tensor `[B, H, W, C]`. Values are clamped to `[0, 1]` and converted to float32. +- `fuzz_mode` (`STRING`, default `"adaptive"`): `"percent"` uses a fixed threshold derived from `fuzz_percent`; `"adaptive"` widens the tolerance based on border variance (`adaptive_k` multiplier). +- `fuzz_percent` (`FLOAT`, default `5.0`): Maximum per-channel deviation (in 0–100%) allowed when classifying border pixels in `"percent"` mode. +- `adaptive_k` (`FLOAT`, default `2.0`): Multiplier applied to the median absolute deviation to set per-channel thresholds in `"adaptive"` mode. +- `edge_margin_px` (`INT`, default `4`): Width of the seed strip sampled along each edge to build the border colour model. +- `pad_px` (`INT`, default `0`): Expands the resulting crop box outward by this many pixels (after detection). +- `use_luma_only` (`BOOLEAN`, default `False`): If `True`, detection runs on luma; otherwise it operates in RGB space. +- `max_growth_iter` (`INT`, default `4096`): Upper bound on the number of dilation iterations when growing the border region. +- `return_border_mask` (`BOOLEAN`, default `False`): When enabled, outputs the binary border mask aligned with the cropped image. +- `use_gpu` (`BOOLEAN`, default `True`): Moves temporary tensors to CUDA when available for faster region-growing. + +## Outputs +- `image` (`IMAGE`): Cropped RGB tensor with even dimensions enforced. +- `left` (`INT`), `top` (`INT`), `width` (`INT`), `height` (`INT`): Crop rectangle relative to the original image (`width`/`height` coerced to even when possible). +- `border_mask` (`MASK`): Binary mask of removed borders, resized to the cropped output. Zero tensor when `return_border_mask` is `False`. + +## Processing Notes +- The algorithm samples all four edges, computes medians and median absolute deviations, then grows a binary mask starting from the border strips. +- Detected borders are dilated and optionally padded, after which the surviving region is cropped and padded to even dimensions for latent-friendly downstream nodes. +- When `return_border_mask` is `False`, the mask socket is still supplied but contains zeros to simplify downstream wiring. + +## Tips +- Raise `edge_margin_px` for thicker frames to capture more representative samples. +- Switch to `"percent"` mode with a low `fuzz_percent` when dealing with synthetic borders that have exact RGB values. +- Disable `use_gpu` on systems with limited VRAM; the logic gracefully continues on CPU.*** End Patch diff --git a/docs/AutoWBColorMatch.md b/docs/AutoWBColorMatch.md new file mode 100644 index 0000000..1ea5a97 --- /dev/null +++ b/docs/AutoWBColorMatch.md @@ -0,0 +1,33 @@ +# AutoWBColorMatch + +## Overview +`AutoWBColorMatch` aligns an image’s white balance and global colour characteristics to match a reference. Multiple strategies are available: gray-world and highlight-based white balance, full or luminance-only Reinhard colour transfer in Lab space, or a hybrid pipeline that combines highlight balancing with Reinhard matching. + +## Inputs +- `image` (`IMAGE`): Source image tensor `[B, H, W, C]`. Alpha channels are discarded. +- `reference` (`IMAGE`): Target look that the source should emulate. +- `method` (`STRING`, default `"wb_highlight+reinhard"`): + - `wb_grayworld` – scales channels so the global mean becomes neutral grey. + - `wb_highlight` – computes gains from the brightest pixels. + - `reinhard_lab` – matches full Lab mean/std to the reference. + - `lab_l_only` – matches only the L component in Lab space. + - `wb_highlight+reinhard` – performs highlight white balance first, then full Reinhard matching. +- `percentile` (`FLOAT`, default `95.0`): Percentile used by the highlight detector when `wb_highlight` is involved (range `80.0–99.9`). +- `strength` (`FLOAT`, default `1.0`): Blends between the original (`0.0`) and fully matched result (`1.0`). +- `clip_gamut` (`BOOLEAN`, default `True`): Clamps the final output to `[0, 1]` after blending. +- `force_size` (`BOOLEAN`, default `False`): When `True`, resizes both images to the supplied `target_width`/`target_height` before computing statistics, then resizes the matched output back to the original resolution. +- `target_width`, `target_height` (`INT`, defaults `1440×1080`): Working resolution used when `force_size` is enabled. + +## Outputs +- `IMAGE`: RGB tensor after white balance and colour matching, blended according to `strength`. + +## Processing Notes +- Colour conversions follow the sRGB ↔ linear ↔ XYZ ↔ Lab pipeline with D65 white point constants. +- Highlight-based white balance uses a quantile mask to estimate channel gains, clamping them between 0.5 and 2.0 to avoid extreme corrections. +- Reinhard matching calculates per-image mean and standard deviation in Lab space; `lab_l_only` preserves chroma while adjusting luminance. +- When `force_size` is active, corrections occur at the specified working resolution for improved speed and comparability; the result is then resampled back to match the original dimensions. + +## Tips +- Use `wb_highlight+reinhard` to neutralise both colour casts and overall contrast in a single step. +- Dial down `strength` to keep some of the original styling when the reference is only a loose guide. +- Disable `clip_gamut` if you plan to perform additional tone mapping afterward and want to preserve unclamped intermediate values.*** End Patch diff --git a/docs/CropImageByMargins.md b/docs/CropImageByMargins.md new file mode 100644 index 0000000..2e24c38 --- /dev/null +++ b/docs/CropImageByMargins.md @@ -0,0 +1,24 @@ +# CropImageByMargins + +## Overview +`CropImageByMargins` removes a specified number of pixels from each edge of an RGB image. It supports optional snapping of the resulting dimensions to a chosen multiple, which is useful when downstream models require even or block-aligned shapes. + +## Inputs +- `image` (`IMAGE`): Tensor `[B, H, W, C]`. The node internally converts data to float32 and clamps to `[0, 1]`. +- `left_px` (`INT`, default `0`): Pixels to remove from the left edge. +- `top_px` (`INT`, default `0`): Pixels to remove from the top edge. +- `right_px` (`INT`, default `0`): Pixels to remove from the right edge. +- `bottom_px` (`INT`, default `0`): Pixels to remove from the bottom edge. +- `snap_multiple` (`INT`, default `1`): Optional modulus applied to the resulting width and height. The crop box is shrunk (never expanded) so that width/height are multiples of this value. Set to `1` to disable snapping. + +## Outputs +- `IMAGE`: Cropped tensor with the same batch size as the input. + +## Processing Notes +- Margins are clamped so the crop never goes outside the image bounds; at least one pixel is preserved in both axes. +- Snapping is performed after margins are applied, preserving the upper-left origin. +- Alpha channels are implicitly discarded by the helper utilities shared across the suite; the output is RGB. + +## Tips +- Use a `snap_multiple` of `64` when preparing tiles for SDXL-style upscalers or patch-based inference. +- Pair this node with `CropMaskByMargins` to keep images and masks aligned.*** End Patch diff --git a/docs/CropMaskByMargins.md b/docs/CropMaskByMargins.md new file mode 100644 index 0000000..5553653 --- /dev/null +++ b/docs/CropMaskByMargins.md @@ -0,0 +1,24 @@ +# CropMaskByMargins + +## Overview +`CropMaskByMargins` mirrors the behaviour of `CropImageByMargins` but operates on mask tensors. It normalises mask shapes into `[B, H, W, 1]`, applies pixel margins, and optionally snaps the result to a given multiple—ideal for keeping mask geometry aligned with cropped images. + +## Inputs +- `mask` (`MASK`): Accepts `[H, W]`, `[H, W, 1]`, `[1, H, W]`, or `[B, H, W, 1]`. Values are converted to float32 in `[0, 1]`. +- `left_px` (`INT`, default `0`): Pixels to remove from the left edge. +- `top_px` (`INT`, default `0`): Pixels to remove from the top edge. +- `right_px` (`INT`, default `0`): Pixels to remove from the right edge. +- `bottom_px` (`INT`, default `0`): Pixels to remove from the bottom edge. +- `snap_multiple` (`INT`, default `1`): Optional modulus applied to width/height after cropping (see the image variant for details). + +## Outputs +- `MASK`: Cropped mask tensor (same batch size as the normalised input). + +## Processing Notes +- Shape coercion covers common mask formats generated by ComfyUI widgets and scripts; unsupported layouts raise a descriptive `ValueError`. +- Margins are clamped so that the crop stays within bounds and retains at least a 1×1 area. +- Snapping uses the same helper as the image node, ensuring both outputs remain in lockstep when identical parameters are used. + +## Tips +- Always route the same margin values to both image and mask versions when preparing composites. +- Set `snap_multiple` to `8` or `16` before feeding masks into latent-space operations that expect even dimensions.*** End Patch diff --git a/docs/FilenameAppendSuffix.md b/docs/FilenameAppendSuffix.md new file mode 100644 index 0000000..32364e3 --- /dev/null +++ b/docs/FilenameAppendSuffix.md @@ -0,0 +1,24 @@ +# FilenameAppendSuffix + +## Overview +`FilenameAppendSuffix` builds a new filename by appending a suffix to an existing path. It supports removing multi-part extensions (e.g., `.tar.gz`), trimming whitespace, and optionally preserving the original directory component. + +## Inputs +- `filename` (`STRING`, default empty): Original filename or path. +- `suffix` (`STRING`, default `"supir"`): Text appended to the base name. Leave empty to omit the suffix altogether. +- `separator` (`STRING`, default `"-"`): Separator inserted between the base name and `suffix` when needed. +- `strip_all_extensions` (`BOOLEAN`, default `True`): When `True`, removes every trailing extension component (useful for multi-part archives). When `False`, only strips the final extension. +- `preserve_directory` (`BOOLEAN`, default `True`): When enabled, reattaches the original directory path to the generated filename. +- `trim_whitespace` (`BOOLEAN`, default `True`): Trims leading/trailing whitespace from the filename, suffix, and separator before processing. + +## Outputs +- `filename` (`STRING`): The rewritten filename/path. + +## Processing Notes +- If the base name already ends with the separator, the suffix is appended without duplicating it. +- When `suffix` is empty, the base name is returned as-is (subject to extension stripping). +- Directory handling uses `os.path.join`, so platform-specific separators are respected. + +## Tips +- Chain this node after `LoadImageCombined` or `PortraitUtils_PairedImageLoader` to build save paths that mirror the input file names. +- Disable `preserve_directory` when you want to relocate files into a new output folder without carrying over their original path.*** End Patch diff --git a/docs/FitAspectHeadSafe.md b/docs/FitAspectHeadSafe.md new file mode 100644 index 0000000..6d5578c --- /dev/null +++ b/docs/FitAspectHeadSafe.md @@ -0,0 +1,32 @@ +# FitAspectHeadSafe + +## Overview +`FitAspectHeadSafe` chooses an aspect ratio from a user-supplied list and computes a crop rectangle that keeps a detected subject centred with controllable headroom, footroom, and side margins. It is designed to work with the bounding boxes emitted by `MQBBoxMin`, maintaining portrait-friendly framing while matching downstream aspect targets. + +## Inputs +- `image` (`IMAGE`): Reference image for dimension lookup only. +- `x`, `y`, `w`, `h` (`INT`): Subject bounding box (e.g., from `MQBBoxMin`). Coordinates are expressed in pixels relative to the original image. +- `aspects_csv` (`STRING`, default `"2:3,3:4,1:1,9:16,16:9,5:8,8:5"`): Comma-separated aspect ratios (width:height). Both portrait and landscape ratios may be provided. +- `match_to` (`STRING`, default `"mq_box"`): Determines the target ratio. `"mq_box"` compares the candidate list to the subject box, `"image"` matches the full image ratio. +- `headroom_ratio` (`FLOAT`, default `0.12`): Fraction of the candidate height reserved above the subject when expanding the crop. +- `footroom_ratio` (`FLOAT`, default `0.06`): Fraction of height reserved below the subject. +- `side_margin_ratio` (`FLOAT`, default `0.08`): Lateral padding as a fraction of the candidate width. +- `bottom_priority` (`FLOAT`, default `0.75`): Bias when reconciling head/foot violations. Values closer to `1` favour the bottom margin (useful for keeping feet in frame). +- `horiz_gravity` (`STRING`, default `"center"`): Horizontal anchoring for the final crop. `"left"` and `"right"` pin the rectangle to the image edges; `"center"` balances around the subject. + +## Outputs +- `w` (`INT`), `h` (`INT`): Width and height of the recommended crop. +- `x` (`INT`), `y` (`INT`): Top-left coordinate of the crop within the original image. +- `aspect_used` (`STRING`): The ratio selected from `aspects_csv`. +- `debug` (`STRING`): Trace information describing the decision process (chosen ratio, applied margins, violation metrics). + +## Processing Notes +- The node parses the aspect list, finds the closest ratio to either the image or subject box, then expands the subject rectangle with head/foot/side margins. +- `_cover_min_rect` ensures the expanded rectangle is the smallest that matches the chosen ratio while fully covering the desired region. +- Vertical placement is refined by evaluating candidate positions and selecting the one with minimum weighted violation, using `bottom_priority` to bias the result. +- The final rectangle is clamped to the image boundaries and rounded to integers. + +## Tips +- Pair with `MQBBoxMin` to generate the `x`, `y`, `w`, `h` inputs automatically from a mask. +- Supply both portrait and landscape ratios in `aspects_csv` when you intend to reuse the node for multi-orientation workflows. +- Inspect the `debug` string when tuning headroom/footroom; it lists the applied margin sizes and any constraint violations.*** End Patch diff --git a/docs/FluxResolutionPrepare.md b/docs/FluxResolutionPrepare.md new file mode 100644 index 0000000..7d5b53a --- /dev/null +++ b/docs/FluxResolutionPrepare.md @@ -0,0 +1,32 @@ +# FluxResolutionPrepare + +## Overview +`FluxResolutionPrepare` crops and resizes an image to the closest legal resolution supported by Flux-family models. It optionally pre-upscales low-resolution inputs to a minimum megapixel budget, respects an optional user-supplied crop window, and reports the chosen aspect ratio along with crop efficiency metrics. + +## Inputs +- `image` (`IMAGE`): Source tensor `[1, H, W, C]`. Batched inputs are not supported. +- `min_megapixels` (`FLOAT`, default `0.95`): Minimum megapixel target used when `enable_pre_upscale` is `True`. The node upsamples the image if the current pixel count is below this threshold. +- `enable_pre_upscale` (`BOOLEAN`, default `True`): Toggles the pre-upscale pass. When disabled, the original resolution is used for cropping. +- `crop_width` (`INT`, optional, default `-1`): Width of a user-specified region of interest. Values `<= 0` disable the manual crop. +- `crop_height` (`INT`, optional, default `-1`): Height of the manual crop ROI. +- `crop_x` (`INT`, optional, default `0`): Left offset of the manual crop (pixels, before pre-upscale). +- `crop_y` (`INT`, optional, default `0`): Top offset of the manual crop (pixels, before pre-upscale). + +## Outputs +- `image` (`IMAGE`): Processed tensor sized to a Flux-compatible resolution. +- `ratio` (`STRING`): Human-readable aspect label (e.g., `3:4`, `9:21`). +- `target_width` (`INT`): Final width selected from the built-in Flux whitelist. +- `target_height` (`INT`): Final height. +- `area_loss_percent` (`FLOAT`): Percentage of the (pre-upscaled) working area discarded during cropping. +- `pre_scale_factor` (`FLOAT`): Scale multiplier applied by the pre-upscale stage (1.0 when disabled or unnecessary). + +## Processing Notes +- The node first pre-upscales by the minimum factor required to hit `min_megapixels`, using bicubic interpolation. +- If a manual crop is specified, it is interpreted in original-image coordinates and mapped to the pre-upscaled space. +- The best Flux resolution is chosen by minimizing scale deviation, area loss, and ratio error across a fixed candidate list (landscape/portrait pairs included). +- Final cropping is centred unless constrained by the boundaries of the manual crop; resizing uses bicubic interpolation. + +## Tips +- Feed the emitted `ratio` into UI labels/logs so you can audit whether the chosen aspect aligns with your prompt or downstream model. +- When passing manual crop dimensions, ensure they leave enough room for the target aspect ratio; otherwise, the node will fall back to the nearest fit that clamps to image bounds. +- Pair this node with `MQBBoxMin` and `FitAspectHeadSafe` to derive a subject-aware crop before committing to Flux sizing.*** End Patch diff --git a/docs/GpuSyncCacheClear.md b/docs/GpuSyncCacheClear.md new file mode 100644 index 0000000..5843f23 --- /dev/null +++ b/docs/GpuSyncCacheClear.md @@ -0,0 +1,18 @@ +# GpuSyncCacheClear + +## Overview +`GpuSyncCacheClear` is a utility node that forces a CUDA synchronisation, runs Python garbage collection, and empties the CUDA memory cache before passing the input image through untouched. It can help stabilise lengthy workflows by periodically freeing cached allocations. + +## Inputs +- `image` (`IMAGE`): Any image tensor. The node does not modify the data. + +## Outputs +- `image` (`IMAGE`): The same tensor received as input. + +## Processing Notes +- CUDA synchronisation and cache clearing are wrapped in `try/except` blocks so the node is safe to use on CPU-only systems or when CUDA is unavailable. +- Python’s garbage collector runs every time the node executes, which can slightly increase runtime but helps reclaim CPU memory between large steps. + +## Tips +- Insert this node between GPU-intensive stages (e.g., after diffusion sampling) to free VRAM before loading another model. +- Because the node simply echoes its input, it can sit inline without altering workflow semantics.*** End Patch diff --git a/docs/ImageMegapixelSelector.md b/docs/ImageMegapixelSelector.md new file mode 100644 index 0000000..9d8c2c8 --- /dev/null +++ b/docs/ImageMegapixelSelector.md @@ -0,0 +1,22 @@ +# ImageMegapixelSelector + +## Overview +`ImageMegapixelSelector` rescales an RGB image to the closest target in {1, 2, 3} megapixels (rounded with ties favouring larger). The resized output can enforce divisibility constraints on width and height, making it handy for models that require dimensions aligned to a specific multiple (e.g., 8 or 64). + +## Inputs +- `image` (`IMAGE`): Source tensor `[B, H, W, C]`. Batching is not yet supported (`B` must be 1). Any alpha channel is discarded prior to processing. +- `divisible_by` (`INT`, default `8`): Ensures both width and height in the result are multiples of this value. Set to `1` to disable snapping. + +## Outputs +- `image` (`IMAGE`): RGB tensor resized via Lanczos filtering to the chosen megapixel target, respecting the divisibility constraint. +- `target_megapixels` (`FLOAT`): The actual megapixel value achieved after rounding and modulus corrections. + +## Processing Notes +- The node computes the raw scale necessary to hit the nearest whole-megapixel target, then searches a neighbourhood of candidate widths that satisfy `divisible_by`. +- Heights are rounded by preserving aspect ratio where possible; several rounding strategies are tried to reduce distortion. +- Resampling uses PIL’s Lanczos filter on the CPU, preserving RGB ordering and clamping to `[0, 1]`. + +## Tips +- Combine this node with `FluxResolutionPrepare` when you want predictable intermediate resolutions prior to Flux cropping. +- Leave `divisible_by` at 8 for Stable Diffusion-style latents, or increase to 64 when feeding grids/tilers that operate on 64-pixel blocks. +- Because the node always outputs RGB, downstream mask compositors can assume three-channel data.*** End Patch diff --git a/docs/LoadImageCombined.md b/docs/LoadImageCombined.md new file mode 100644 index 0000000..37a80a1 --- /dev/null +++ b/docs/LoadImageCombined.md @@ -0,0 +1,29 @@ +# LoadImageCombined + +## Overview +`LoadImageCombined` unifies single-image selection and batch iteration within one node. In single mode it behaves like ComfyUI’s standard loader; in batch mode it iterates through a directory listing on each execution, keeping track of the last-delivered file to maintain sequence order. + +## Inputs +- `mode` (`STRING`, default `"Single"`): `"Single"` loads a user-selected file; `"Batch"` scans a directory and auto-advances each run. +- `input_dir` (`STRING`, default empty): Directory to scan in batch mode. Relative paths resolve against `ComfyUI/input`. +- `output_dir` (`STRING`, default empty): Convenience string passed through for downstream savers/loggers. +- `pattern` (`STRING`, default `"*"`): Glob pattern (e.g., `*.png`) applied in batch mode. +- `strip_trailing_numbers` (`BOOLEAN`, default `False`): When `True`, removes trailing `(1)`, `(2)`, etc., from filenames when computing the `filename_no_ext` output. +- `repeat_last` (`BOOLEAN`, default `False`): Keeps returning the last image instead of advancing on subsequent runs. +- `image` (`STRING`): File selector used in single mode. Only relevant when `mode == "Single"`. + +## Outputs +- `IMAGE`: RGB tensor loaded from disk (alpha channels are dropped). +- `filename_no_ext` (`STRING`): Base filename without extension (and optionally without trailing number suffix). +- `output_dir` (`STRING`): Passthrough of the input field. +- `width` (`INT`), `height` (`INT`): Dimensions of the loaded image. + +## Processing Notes +- Directory listings are case-insensitive and filtered to common image extensions (`png`, `jpg`, `jpeg`, `webp`, `tif`, `tiff`, `bmp`). +- Batch iteration state is held in-memory per directory/pattern combination and resets when the listing changes (new/removed files) or when `repeat_last` is enabled. +- All images are converted to RGB using Pillow, with Exif orientation applied. + +## Tips +- Chain this node with `PortraitUtils_PairedImageLoader` or comparison nodes to quickly inspect before/after renders. +- Use `repeat_last` when you need to re-run the same batch frame after adjusting downstream parameters. +- Combine with `FilenameAppendSuffix` to build descriptive output paths based on `filename_no_ext`.*** End Patch diff --git a/docs/MQBBoxMin.md b/docs/MQBBoxMin.md new file mode 100644 index 0000000..a7d528d --- /dev/null +++ b/docs/MQBBoxMin.md @@ -0,0 +1,26 @@ +# MQBBoxMin + +## Overview +`MQBBoxMin` (Median/Quantile Bounding Box) derives a robust subject bounding box from a mask. It auto-detects whether the mask represents foreground or background, applies quantile-based trimming to ignore outliers, and optionally inflates the result to maintain a minimum span. The node returns both the bounding box and a debug string describing the decision path. + +## Inputs +- `mask` (`MASK`): Input tensor in any common mask layout; converted to a 2D float array internally. +- `invert_mask` (`STRING`, default `"auto"`): Controls which side of the mask is treated as foreground. `"auto"` picks the side with less border mass, `"false"` keeps the mask as-is, `"true"` inverts it. +- `q_left`, `q_right`, `q_top`, `q_bottom` (`FLOAT`, defaults `0.005` / `0.995`): Quantiles used to crop the subject along each axis. Lower values trim only extreme borders; higher values create tighter boxes. +- `min_span_px` (`INT`, default `8`): Minimum width/height enforced after quantile cropping. +- `tight_crop` (`BOOLEAN`, default `True`): When `False`, the node simply returns the full image bounds. + +## Outputs +- `x`, `y`, `w`, `h` (`INT`): Bounding box coordinates in pixels relative to the mask. +- `debug` (`STRING`): Notes on the foreground selection, quantile bounds, and any aspect-ratio adjustments. + +## Processing Notes +- The mask is normalised to `[0, 1]`, with automatic conversion from 8-bit values when necessary. +- Quantile bounds are computed using cumulative sums along rows and columns, ensuring consistent behaviour even for noisy masks. +- The routine grows the quantile box by a padding proportional to image size and then searches for the minimal-area rectangle matching the image aspect ratio. Additional candidate aspect ratios (1:1, 2:3, 3:4, 9:16, etc.) are evaluated to minimise area. +- If a suitable aspect container cannot be found, the raw quantile rectangle is returned. + +## Tips +- Feed the output directly into `FitAspectHeadSafe` to generate camera-friendly crops. +- Raise `min_span_px` when working with thin subject masks to avoid degenerate boxes. +- For masks that include the background (white is background, black is subject), set `invert_mask` to `"true"` instead of relying on `"auto"`.*** End Patch diff --git a/docs/MultiPromptNode.md b/docs/MultiPromptNode.md new file mode 100644 index 0000000..cab5594 --- /dev/null +++ b/docs/MultiPromptNode.md @@ -0,0 +1,21 @@ +# MultiPromptNode + +## Overview +`MultiPromptNode` exposes four independent multiline string sockets that can hold prompt snippets, masks, or any free-form text. It does not alter or combine the strings; the node simply passes them through so they can be reused elsewhere in a workflow. + +## Inputs +- `subject_mask` (`STRING`, default empty): Free-form field typically used to describe subject-masking instructions. +- `inpaint_prompt` (`STRING`, default empty): Slot for inpainting-specific guidance. +- `outpaint_prompt` (`STRING`, default empty): Slot for outpainting instructions. +- `stitch_mask` (`STRING`, default empty): Slot for notes or textual encodings related to stitching. + +## Outputs +- `subject_mask` (`STRING`) +- `inpaint_prompt` (`STRING`) +- `outpaint_prompt` (`STRING`) +- `stitch_mask` (`STRING`) + +## Usage Notes +- All four sockets are optional; the node returns whatever was entered, including empty strings. +- Because outputs are plain strings, they can feed directly into conditioning nodes, prompt editors, or scripting helpers. +- Rename the node instance in your workflow to reflect the actual meaning of each socket if you repurpose them for other tasks.*** End Patch diff --git a/docs/OtherConfigNode.md b/docs/OtherConfigNode.md new file mode 100644 index 0000000..6cd0887 --- /dev/null +++ b/docs/OtherConfigNode.md @@ -0,0 +1,28 @@ +# OtherConfigNode + +## Overview +`OtherConfigNode` is a lightweight configuration hub that emits a handful of numeric and boolean values commonly reused across portrait workflows. Unlike `WorkflowConfig`, it does not interact with presets or metadata—its role is to provide a simple socket interface that keeps related parameters grouped. + +## Inputs +- `initial_upscale_megapixels` (`FLOAT`, default `1.0`): Megapixel target for the initial upscale stage. +- `final_upscale_megapixels` (`FLOAT`, default `4.0`): Megapixel target for the final upscale stage. +- `tight_crop` (`BOOLEAN`, default `False`) +- `blend_factor` (`FLOAT`, default `1.0`): Blend strength (0–1) for downstream compositing nodes. +- `stitch_bypass_mask` (`BOOLEAN`, default `False`) +- `input_dir` (`STRING`, default empty): Directory hint propagated to loaders or savers. + +## Outputs +- `initial_upscale (megapixels)` (`FLOAT`) +- `final_upscale (megapixels)` (`FLOAT`) +- `tight_crop` (`BOOLEAN`) +- `blend_factor` (`FLOAT`) +- `stitch_bypass_mask` (`BOOLEAN`) +- `input_dir` (`STRING`) + +## Processing Notes +- The node clamps upscale megapixels to a minimum of `0.1` and blend factors to the `[0, 1]` range. +- Whitespace is preserved in `input_dir`; trim it upstream if you require strict paths. + +## Tips +- Use this node as a compact fan-out hub when you don’t need the full preset system from `WorkflowConfig`. +- Route `blend_factor` into `StitchByMask.opacity` or similar nodes to keep mix strength consistent across passes.*** End Patch diff --git a/docs/OutpaintConfigNode.md b/docs/OutpaintConfigNode.md new file mode 100644 index 0000000..5e122ac --- /dev/null +++ b/docs/OutpaintConfigNode.md @@ -0,0 +1,26 @@ +# OutpaintConfigNode + +## Overview +`OutpaintConfigNode` collects user preferences for outpainting operations—gravity, measurement mode, and per-edge padding. Its outputs are raw values that can feed into padding calculators or custom scripts without additional parsing. + +## Inputs +- `mode` (`STRING`, default `"Percent"`): Determines whether outpainting requests are expressed as percentages of the current canvas or absolute pixels (`"Pixels"`). +- `gravity` (`STRING`, default `"center"`): Direction that should receive the majority of the expansion. Supports centre, single-direction, and corner presets. +- `horizontal_percent` (`FLOAT`, default `20.0`): Horizontal expansion percentage used when `mode` is `"Percent"`. +- `vertical_percent` (`FLOAT`, default `10.0`): Vertical expansion percentage used when `mode` is `"Percent"`. +- `left_px`, `right_px`, `top_px`, `bottom_px` (`INT`, defaults `0`): Explicit pixel padding for each side (used when `mode` is `"Pixels"`). + +## Outputs +- `mode` (`STRING`) +- `gravity` (`STRING`) +- `horizontal_percent` (`FLOAT`) +- `vertical_percent` (`FLOAT`) +- `left_px` (`INT`) +- `right_px` (`INT`) +- `top_px` (`INT`) +- `bottom_px` (`INT`) + +## Usage Notes +- All outputs are emitted regardless of mode, allowing downstream nodes to choose the correct interpretation. +- Percentages and pixel values are clamped to sensible minima (`>= 0`) before emission. +- Combine with `OutpaintPaddingComputeNode` to convert settings into absolute padding values for a specific image.*** End Patch diff --git a/docs/OutpaintPaddingComputeNode.md b/docs/OutpaintPaddingComputeNode.md new file mode 100644 index 0000000..5892e6f --- /dev/null +++ b/docs/OutpaintPaddingComputeNode.md @@ -0,0 +1,28 @@ +# OutpaintPaddingComputeNode + +## Overview +`OutpaintPaddingComputeNode` converts outpainting preferences (percentages or absolute pixels plus gravity) into concrete padding values for a given image. It ensures the final canvas dimensions remain even—useful for latent-based diffusion models—and supports both symmetric and directional growth. + +## Inputs +- `image` (`IMAGE`): Reference image whose width/height determine the base canvas. +- `mode` (`STRING`, default `"Percent"`): Matches the setting emitted by `OutpaintConfigNode`. `"Pixels"` bypasses percentage math. +- `gravity` (`STRING`, default `"center"`): Indicates where additional space should be allocated when working in percent mode. Supports centre, edges, and corners. +- `horizontal_percent` (`FLOAT`, default `20.0`): Horizontal expansion percentage. +- `vertical_percent` (`FLOAT`, default `10.0`): Vertical expansion percentage. +- `left_px`, `right_px`, `top_px`, `bottom_px` (`INT`, defaults `0`): Absolute pixel paddings used when `mode` is `"Pixels"`. + +## Outputs +- `left` (`INT`) +- `top` (`INT`) +- `right` (`INT`) +- `bottom` (`INT`) + +## Processing Notes +- The node first reads the image dimensions from the BHWC tensor and converts all numeric inputs to safe ranges (`>= 0`). +- In `"Pixels"` mode the provided per-edge values are used directly, followed by even-size enforcement. +- In `"Percent"` mode the node computes the total horizontal/vertical padding and splits it according to `gravity` (e.g., bottom-heavy when gravity is `"bottom"`). +- Even-dimension enforcement increments the right and/or bottom padding as needed so `(width + left + right)` and `(height + top + bottom)` are divisible by 2. + +## Tips +- Feed the outputs straight into padding-aware samplers or ControlNet preprocessing nodes. +- When chaining with `OutpaintConfigNode`, keep the gravity strings in sync to avoid unintended splits.*** End Patch diff --git a/docs/PortraitUtils_PairedImageLoader.md b/docs/PortraitUtils_PairedImageLoader.md new file mode 100644 index 0000000..c8858d3 --- /dev/null +++ b/docs/PortraitUtils_PairedImageLoader.md @@ -0,0 +1,27 @@ +# PortraitUtils_PairedImageLoader + +## Overview +`PortraitUtils_PairedImageLoader` pairs images from two directories (typically “source” and “output”) and emits the next matching duo on every run. It keeps its position per node instance, supports forward/backward stepping, and reports the active filename for UI display or logging. + +## Inputs +- `source_dir` (`STRING`): Directory containing the reference/source images. Relative paths resolve against the current working directory. +- `output_dir` (`STRING`): Directory containing the processed images to compare with the source set. +- `reverse` (`BOOLEAN`, default `False`): When enabled, stepping moves backwards through the matched list. +- `strip_trailing_numbers` (`BOOLEAN`, default `False`): Strips trailing ` (1)`, ` (2)`, etc., before matching filenames, useful when OS-level copy suffixes are present. +- `unique_id` (`UNIQUE_ID`, hidden): Injected by ComfyUI so the node can preserve its index across runs and avoid collisions when duplicated. + +## Outputs +- `source_image` (`IMAGE`) +- `output_image` (`IMAGE`) +- `filename` (`STRING`): The source filename associated with the current pair (base name only). + +## Matching & Caching Behaviour +- Filenames are matched case-insensitively by base name (optionally after stripping trailing numbers). Extensions may differ; the loader picks the first intersection and logs any extension mismatches. +- Directories are rescanned on each execution. The node maintains an index keyed by a signature derived from file names, sizes, and modification times so it can resume where it left off unless the listing changes. +- One-time warnings are printed when files exist only in one directory or when name collisions arise after stripping suffixes. +- Because `NOT_IDEMPOTENT = True`, ComfyUI will not cache outputs—the node always reloads the active pair, so edits on disk are visible immediately. + +## Tips +- Use the `reverse` toggle to walk backwards when you overshoot a comparison target. +- Connect the `filename` output to a `ShowText`-style widget (or metadata saver) to document which pair was reviewed. +- Keep directory sizes manageable; the loader reads both folders on every run to stay in sync.*** End Patch diff --git a/docs/SeedVR2Prep.md b/docs/SeedVR2Prep.md new file mode 100644 index 0000000..04e6b85 --- /dev/null +++ b/docs/SeedVR2Prep.md @@ -0,0 +1,21 @@ +# SeedVR2Prep + +## Overview +`SeedVR2Prep` prepares images for SeedVR2 and other VRAM-sensitive pipelines by enforcing even dimensions and optionally converting tensors to FP16 on the GPU. It ensures data is in `[B, H, W, 3]` format, repeating channels as necessary. + +## Inputs +- `image` (`IMAGE`): Input tensor. Non-tensor inputs are converted to `torch.Tensor` before processing. +- `ensure_even_dims` (`BOOLEAN`, default `True`): When enabled, trims or pads by repeating edge pixels so both width and height are even. +- `to_fp16_on_gpu` (`BOOLEAN`, default `True`): If CUDA is available, moves the tensor to GPU and converts it to float16. + +## Outputs +- `image` (`IMAGE`): Sanitised tensor with even dimensions and optional FP16 GPU residency. + +## Processing Notes +- Channel counts greater than 3 are truncated to RGB; single-channel inputs are repeated across RGB. +- Padding when width/height equals 1 duplicates the final column/row so dimensions never drop to zero. +- FP16 conversion only occurs when CUDA is available; otherwise the data remains on CPU in float32. + +## Tips +- Place this node right before VRAM-heavy samplers to cut down on memory usage and avoid odd-dimension errors. +- Disable `to_fp16_on_gpu` if a downstream node expects CPU tensors or full precision.*** End Patch diff --git a/docs/StitchByMask.md b/docs/StitchByMask.md new file mode 100644 index 0000000..6465684 --- /dev/null +++ b/docs/StitchByMask.md @@ -0,0 +1,31 @@ +# StitchByMask + +## Overview +`StitchByMask` composites two RGB images using a soft mask, with optional mask inversion, bypass blending, feathering, and enforced output dimensions. It is geared toward portrait editing workflows where foreground/background swaps or stitched passes need precise control. + +## Inputs +- `image_a` (`IMAGE`): Base layer tensor. +- `image_b` (`IMAGE`): Layer to composite over `image_a`. +- `mask` (`MASK`, default `None`): Blend mask. Required unless `bypass_mask` is `True`. Accepted in common mask layouts (`[B,H,W]`, `[B,H,W,1]`, `[H,W]`, `[H,W,1]`). +- `invert_mask` (`BOOLEAN`, default `False`): Flips the mask before blending. +- `bypass_mask` (`BOOLEAN`, default `False`): When enabled, ignores the supplied mask and uses a constant opacity value. +- `opacity` (`FLOAT`, default `1.0`): Global blend strength (0 preserves `image_a`, 1 favours `image_b`). Applied after optional inversion. +- `feather_radius` (`INT`, default `5`): Radius (in pixels) for disk dilation before Gaussian feathering. Set to `0` to disable feathering. +- `force_size` (`BOOLEAN`, default `False`): If `True`, both images (and the mask) are resized to `target_width` × `target_height` before blending. +- `target_width` (`INT`, default `1344`): Width used when `force_size` is `True`. +- `target_height` (`INT`, default `768`): Height used when `force_size` is `True`. + +## Outputs +- `stitched` (`IMAGE`): Composited RGB image, clamped to `[0, 1]`. +- `processed_mask` (`MASK`): The mask actually used for blending (after inversion, opacity scaling, feathering, and resizing). + +## Processing Notes +- Input tensors are converted to BHWC float32 arrays. Alpha channels, if present, are truncated to RGB. +- When `force_size` is enabled, image tensors are resized with bilinear sampling and masks with nearest-neighbour sampling to avoid edge blurring before feathering. +- Feathering expands the binary core of the mask via convolution with a disk kernel, then applies a Gaussian blur (`sigma ≈ radius/2`) and takes the element-wise maximum with the original mask to avoid revealing previously masked pixels. +- In bypass mode the node builds a uniform mask at the requested `opacity`, allowing linear interpolation between `image_a` and `image_b` without supplying a mask. + +## Tips +- Set `feather_radius` to `0` when blending masks that already contain feathered edges. +- Use bypass mode to create simple crossfades or to output the second pass at reduced opacity without worrying about masks. +- When enforcing a specific resolution (e.g., 1344×768 for Flux), ensure both inputs already match the target aspect to avoid distortion.*** End Patch diff --git a/docs/WorkflowConfig.md b/docs/WorkflowConfig.md new file mode 100644 index 0000000..9707d10 --- /dev/null +++ b/docs/WorkflowConfig.md @@ -0,0 +1,50 @@ +# WorkflowConfig + +## Overview +`WorkflowConfig` centralises high-level settings (upscale amounts, cropping policy, stitch opacity, seed control, notes) and exposes them both as individual sockets and as a JSON bundle. It supports loading/saving presets, honours server-side random seed generation, and attempts to keep workflow metadata in sync when special seed values are used. + +## Inputs +Required parameters: +- `template_name` (`STRING`): Built-in preset to load (currently `"Default"`). Templates provide fallback values for all outputs. +- `saved_preset` (`STRING`): Filename stem of a saved preset in the workspace (`` to disable). +- `load_saved_preset` (`BOOLEAN`, default `False`): When `True`, merges the chosen `saved_preset` over the template defaults. +- `initial_upscale` (`FLOAT`, default `1.0`): Upscale factor applied before the main pass. +- `final_upscale` (`FLOAT`, default `4.0`): Target upscale factor for the final output. +- `tight_crop` (`BOOLEAN`, default `False`): Indicates whether downstream nodes should prefer tight subject crops. +- `stitch_bypass_mask` (`BOOLEAN`, default `False`): Signals that stitched blends should ignore masks. +- `stitch_opacity` (`FLOAT`, default `1.0`): Preferred opacity for stitch operations (0–1). +- `crop_fuzz_percent` (`FLOAT`, default `5.0`): Suggested fuzziness allowance for auto-cropping nodes. +- `crop_pad_px` (`INT`, default `0`): Additional padding to apply around detected crops. +- `seed` (`INT`, default `0`): Primary seed value. Special values `-1`, `-2`, `-3` trigger server-side random generation (increment/decrement where possible). +- `job_notes` (`STRING`, default empty): Arbitrary notes stored alongside the configuration. + +Optional parameters: +- `save_preset_as` (`STRING`): Filename stem to use when persisting the current configuration. +- `write_preset` (`BOOLEAN`, default `False`): When `True`, writes `save_preset_as` to disk in the preset folder. + +Hidden inputs (managed by ComfyUI): +- `prompt` (`PROMPT`): Used to persist generated seeds back into the executing graph. +- `extra_pnginfo` (`EXTRA_PNGINFO`): Allows embedding the resolved seed into workflow metadata. +- `unique_id` (`UNIQUE_ID`): Required to update workflow widgets when special seeds are consumed. + +## Outputs +- `initial_upscale` (`FLOAT`) +- `final_upscale` (`FLOAT`) +- `tight_crop` (`BOOLEAN`) +- `stitch_bypass_mask` (`BOOLEAN`) +- `stitch_opacity` (`FLOAT`) +- `crop_fuzz_percent` (`FLOAT`) +- `crop_pad_px` (`INT`) +- `seed` (`INT`): The resolved seed after handling special values. +- `job_notes` (`STRING`) +- `config_bundle` (`STRING`): JSON string containing all of the above values. + +## Processing Notes +- Templates and saved presets are merged in order: template defaults → saved preset (if loaded) → current widget values. +- When `seed` is `-1`, `-2`, or `-3`, the node generates a random seed server-side. Because the client-side UI cannot receive the update automatically, `WorkflowConfig` attempts to write the new seed back into both the workflow metadata and the executing prompt using the hidden inputs. +- Presets are stored under `ComfyUI/user/default/workflows/workflow_config_presets/` relative to the install. Files are JSON-formatted and written only when `write_preset` is explicitly toggled. + +## Tips +- Drive other nodes (e.g., `AutoCropBorders`, `StitchByMask`, custom scripts) using the JSON bundle to keep all config values in sync. +- Use `job_notes` as a quick logging field—its contents travel alongside the bundle and can be embedded into output metadata downstream. +- Save presets for different portrait styles (tight headshots vs. wide environmental shots) and swap them via the dropdown without rearranging your workflow.*** End Patch diff --git a/filename_append_suffix.py b/filename_append_suffix.py new file mode 100644 index 0000000..f8fb044 --- /dev/null +++ b/filename_append_suffix.py @@ -0,0 +1,76 @@ +import os +import re + + +def _strip_extensions(path: str, strip_all: bool = True) -> tuple[str, str]: + """ + Return (directory, base_without_extension). + If strip_all is True, remove every trailing extension segment (e.g. .tar.gz). + """ + path = path.strip() + directory = os.path.dirname(path) + basename = os.path.basename(path) + if strip_all: + basename = re.sub(r"(?:\.[A-Za-z0-9]{1,5})+$", "", basename) + else: + basename = os.path.splitext(basename)[0] + return directory, basename + + +class FilenameAppendSuffix: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "filename": ("STRING", {"multiline": False, "default": ""}), + "suffix": ("STRING", {"multiline": False, "default": "supir"}), + "separator": ("STRING", {"multiline": False, "default": "-"}), + "strip_all_extensions": ("BOOLEAN", {"default": True}), + "preserve_directory": ("BOOLEAN", {"default": True}), + "trim_whitespace": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + FUNCTION = "build" + CATEGORY = "Utils/IO" + + def build( + self, + filename: str, + suffix: str = "supir", + separator: str = "-", + strip_all_extensions: bool = True, + preserve_directory: bool = True, + trim_whitespace: bool = True, + ): + fn = filename if isinstance(filename, str) else str(filename) + sx = suffix if isinstance(suffix, str) else str(suffix) + sep = separator if isinstance(separator, str) else "-" + + if trim_whitespace: + fn = fn.strip() + sx = sx.strip() + sep = sep.strip() or "-" + + directory, base = _strip_extensions(fn, strip_all=strip_all_extensions) + + if not sx: + out = base + else: + out = f"{base}{sx}" if base.endswith(sep) else f"{base}{sep}{sx}" + + if preserve_directory and directory: + out = os.path.join(directory, out) + + return (out,) + + +NODE_CLASS_MAPPINGS = { + "FilenameAppendSuffix": FilenameAppendSuffix, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FilenameAppendSuffix": "Filename: Append Suffix", +} diff --git a/flux_resolution_prepare.py b/flux_resolution_prepare.py new file mode 100644 index 0000000..e38c9ee --- /dev/null +++ b/flux_resolution_prepare.py @@ -0,0 +1,285 @@ +""" +Flux-friendly aspect cropper and resolution normalizer. + +This node crops an image to the closest allowed aspect ratio while discarding +as little content as possible, optionally pre-upscaling low resolution inputs, +then resizes the crop to the nearest Flux-compatible resolution. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch +import torch.nn.functional as F + +@dataclass(frozen=True) +class _TargetResolution: + width: int + height: int + label: str + ratio_label_override: Optional[str] = None + + @property + def ratio(self) -> float: + return self.width / self.height + + @property + def ratio_label(self) -> str: + if self.ratio_label_override: + return self.ratio_label_override + num, den = _normalize_ratio(self.width, self.height) + return f"{num}:{den}" + + @property + def area(self) -> int: + return self.width * self.height + + +def _normalize_ratio(num: int, den: int) -> Tuple[int, int]: + g = math.gcd(num, den) + return num // g, den // g + + +def _swap_ratio_label(label: Optional[str]) -> Optional[str]: + if not label or ":" not in label: + return label + left, right = label.split(":", 1) + return f"{right}:{left}" + + +def _build_target_resolutions() -> Tuple[_TargetResolution, ...]: + """Flux legal resolution list including landscape inverses.""" + base_specs = [ + (896, 1152, None), + (832, 1216, None), + (768, 1344, None), + (640, 1536, "9:21"), + (1024, 1024, None), + (1152, 1728, None), + (1216, 1664, None), + (1088, 1920, None), + (960, 2176, "9:21"), + (1408, 1408, None), + ] + + targets: List[_TargetResolution] = [] + seen = set() + for width, height, ratio_label in base_specs: + for w, h, lbl in ( + (width, height, ratio_label), + (height, width, _swap_ratio_label(ratio_label)), + ): + key = (w, h) + if key in seen: + continue + seen.add(key) + targets.append(_TargetResolution(w, h, f"{w}x{h}", lbl)) + return tuple(targets) + + +_TARGET_RESOLUTIONS: Tuple[_TargetResolution, ...] = _build_target_resolutions() + + +def _ensure_bhwc(image) -> torch.Tensor: + tensor = image + if not isinstance(tensor, torch.Tensor): + tensor = torch.tensor(tensor) + if tensor.dim() == 3: + tensor = tensor.unsqueeze(0) + if tensor.dim() != 4: + raise ValueError(f"Expected IMAGE tensor [B,H,W,C], got {tuple(tensor.shape)}") + if tensor.shape[-1] not in (3, 4): + raise ValueError("IMAGE tensor must have 3 (RGB) or 4 (RGBA) channels") + return tensor.to(torch.float32).clamp(0.0, 1.0) + + +def _resize_sample(sample: torch.Tensor, height: int, width: int) -> torch.Tensor: + if sample.shape[0] == height and sample.shape[1] == width: + return sample + # sample: [H, W, C] + tensor = sample.permute(2, 0, 1).unsqueeze(0) + tensor = F.interpolate( + tensor, + size=(height, width), + mode="bicubic", + align_corners=False, + ) + return tensor.squeeze(0).permute(1, 2, 0) + + +def _compute_crop_dims(width: int, height: int, ratio: float) -> Tuple[int, int]: + if ratio <= 0.0: + return width, height + + current_ratio = width / height + if current_ratio > ratio: + target_height = height + target_width = max(1, min(width, int(round(target_height * ratio)))) + else: + target_width = width + target_height = max(1, min(height, int(round(target_width / ratio)))) + + target_width = min(target_width, width) + target_height = min(target_height, height) + return target_width, target_height + + +def _center_crop(sample: torch.Tensor, width: int, height: int) -> Tuple[torch.Tensor, int, int]: + src_h, src_w = sample.shape[:2] + left = max(0, (src_w - width) // 2) + top = max(0, (src_h - height) // 2) + right = min(src_w, left + width) + bottom = min(src_h, top + height) + # Re-align if rounding caused overflow. + left = max(0, right - width) + top = max(0, bottom - height) + cropped = sample[top:bottom, left:right, :] + return cropped, left, top + + +def _select_target_combo(width: int, height: int): + if width <= 0 or height <= 0: + raise ValueError("Invalid crop dimensions encountered") + + original_area = width * height + original_ratio = width / height + best = None + best_key = None + + for target in _TARGET_RESOLUTIONS: + ratio = target.ratio + crop_w, crop_h = _compute_crop_dims(width, height, ratio) + if crop_w <= 0 or crop_h <= 0: + continue + + crop_area = crop_w * crop_h + area_loss = 1.0 - (crop_area / max(original_area, 1)) + area_loss = max(0.0, min(1.0, area_loss)) + + cropped_ratio = crop_w / crop_h + ratio_delta = abs(cropped_ratio - ratio) + origin_delta = abs(cropped_ratio - original_ratio) + + scale_w = target.width / max(crop_w, 1) + scale_h = target.height / max(crop_h, 1) + scale_dev = max(abs(scale_w - 1.0), abs(scale_h - 1.0)) + scale_mean = (abs(scale_w - 1.0) + abs(scale_h - 1.0)) * 0.5 + + key = (scale_dev, area_loss, scale_mean, ratio_delta, origin_delta, target.area) + if best_key is None or key < best_key: + best = (target, crop_w, crop_h, area_loss) + best_key = key + + if best is None: + raise RuntimeError("Failed to choose Flux target resolution") + return best + + +def _maybe_pre_upscale(sample: torch.Tensor, min_megapixels: float) -> Tuple[torch.Tensor, float]: + height, width = sample.shape[:2] + area = height * width + min_pixels = max(0.0, min_megapixels) * 1_000_000.0 + if min_pixels <= 0.0 or area >= min_pixels: + return sample, 1.0 + + scale = math.sqrt(min_pixels / max(area, 1.0)) + new_height = max(1, int(round(height * scale))) + new_width = max(1, int(round(width * scale))) + if new_height == height and new_width == width: + return sample, 1.0 + upscaled = _resize_sample(sample, new_height, new_width) + return upscaled, new_width / max(width, 1.0) + + +class FluxResolutionPrepare: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "min_megapixels": ("FLOAT", {"default": 0.95, "min": 0.1, "max": 64.0, "step": 0.05}), + "enable_pre_upscale": ("BOOLEAN", {"default": True}), + }, + "optional": { + "crop_width": ("INT", {"default": -1}), + "crop_height": ("INT", {"default": -1}), + "crop_x": ("INT", {"default": 0}), + "crop_y": ("INT", {"default": 0}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "FLOAT", "FLOAT") + RETURN_NAMES = ("image", "ratio", "target_width", "target_height", "area_loss_percent", "pre_scale_factor") + FUNCTION = "apply" + CATEGORY = "image/transform" + + def apply( + self, + image, + min_megapixels=0.95, + enable_pre_upscale=True, + crop_width=-1, + crop_height=-1, + crop_x=0, + crop_y=0, + ): + tensor = _ensure_bhwc(image) + if tensor.shape[0] != 1: + raise ValueError("FluxResolutionPrepare expects an unbatched IMAGE tensor (B=1)") + + sample = tensor[0] + orig_height, orig_width = sample.shape[:2] + pre_scale = 1.0 + if enable_pre_upscale: + sample, pre_scale = _maybe_pre_upscale(sample, min_megapixels) + + sample_height, sample_width = sample.shape[:2] + scale_x = sample_width / max(1, orig_width) + scale_y = sample_height / max(1, orig_height) + + orig_area = float(sample_height * sample_width) + + work = sample + if isinstance(crop_width, (int, float)) and isinstance(crop_height, (int, float)): + if crop_width > 0 and crop_height > 0: + sx = int(round(max(0.0, crop_x) * scale_x)) + sy = int(round(max(0.0, crop_y) * scale_y)) + sw = int(round(crop_width * scale_x)) + sh = int(round(crop_height * scale_y)) + if sw > 0 and sh > 0: + sx = max(0, min(sx, sample_width - 1)) + sy = max(0, min(sy, sample_height - 1)) + sw = max(1, min(sw, sample_width - sx)) + sh = max(1, min(sh, sample_height - sy)) + work = sample[sy : sy + sh, sx : sx + sw, :] + + work_height, work_width = work.shape[:2] + target, crop_w, crop_h, _ = _select_target_combo(work_width, work_height) + cropped, _, _ = _center_crop(work, crop_w, crop_h) + + resized = _resize_sample(cropped, target.height, target.width).clamp(0.0, 1.0) + output = resized.unsqueeze(0) + + final_area = float(crop_w * crop_h) + total_area_loss = max(0.0, 1.0 - (final_area / max(orig_area, 1.0))) + area_loss_percent = float(total_area_loss * 100.0) + return ( + output, + target.ratio_label, + int(target.width), + int(target.height), + area_loss_percent, + float(pre_scale), + ) + + +NODE_CLASS_MAPPINGS = { + "FluxResolutionPrepare": FluxResolutionPrepare, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxResolutionPrepare": "Flux Resolution Prepare", +} diff --git a/gpu_cache_clear.py b/gpu_cache_clear.py new file mode 100644 index 0000000..0c38bea --- /dev/null +++ b/gpu_cache_clear.py @@ -0,0 +1,99 @@ +import gc as _iutils_gc +import torch as _iutils_torch + + +def _iutils_ensure_bhwc_rgb(image): + t = image + if not isinstance(t, _iutils_torch.Tensor): + t = _iutils_torch.tensor(t) + t = t.to(_iutils_torch.float32) + if t.dim() == 3: + t = t.unsqueeze(0) + # Expect [B,H,W,C] + if t.shape[-1] == 3: + return t.clamp(0.0, 1.0) + elif t.shape[-1] > 3: + return t[..., :3].clamp(0.0, 1.0) + elif t.shape[-1] == 1: + return t.repeat(1, 1, 1, 3).clamp(0.0, 1.0) + else: + return t.clamp(0.0, 1.0) + + +class GpuSyncCacheClear: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"image": ("IMAGE",)}} + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply" + CATEGORY = "image/utils" + + def apply(self, image): + try: + if _iutils_torch.cuda.is_available(): + _iutils_torch.cuda.synchronize() + except Exception: + pass + _iutils_gc.collect() + try: + if _iutils_torch.cuda.is_available(): + _iutils_torch.cuda.empty_cache() + except Exception: + pass + return (image,) + + +class SeedVR2Prep: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "ensure_even_dims": ("BOOLEAN", {"default": True}), + "to_fp16_on_gpu": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply" + CATEGORY = "image/utils" + + def apply(self, image, ensure_even_dims=True, to_fp16_on_gpu=True): + x = _iutils_ensure_bhwc_rgb(image) + if ensure_even_dims: + H, W = x.shape[1], x.shape[2] + # Width + if (W % 2) == 1: + if W > 1: + x = x[:, :, : W - 1, :] + else: + # pad a column by repeating the last column + x = _iutils_torch.cat([x, x[:, :, -1:, :]], dim=2) + # Height + H, W = x.shape[1], x.shape[2] + if (H % 2) == 1: + if H > 1: + x = x[:, : H - 1, :, :] + else: + # pad a row by repeating last row + x = _iutils_torch.cat([x, x[:, -1:, :, :]], dim=1) + + if to_fp16_on_gpu and _iutils_torch.cuda.is_available(): + if x.device.type != "cuda": + x = x.to("cuda", non_blocking=True) + x = x.to(dtype=_iutils_torch.float16) + return (x,) + + +NODE_CLASS_MAPPINGS = { + "GpuSyncCacheClear": GpuSyncCacheClear, + "SeedVR2Prep": SeedVR2Prep, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GpuSyncCacheClear": "GPU Sync & Cache Clear", + "SeedVR2Prep": "SeedVR2 Prep (sanitize)", +} diff --git a/image_megapixel_selector.py b/image_megapixel_selector.py new file mode 100644 index 0000000..1364d14 --- /dev/null +++ b/image_megapixel_selector.py @@ -0,0 +1,185 @@ +import math +from typing import Tuple + +import numpy as np +import torch +from PIL import Image + + +class ImageMegapixelSelector: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "divisible_by": ( + "INT", + { + "default": 8, + "min": 1, + "max": 128, + "step": 1, + "tooltip": "Force output width/height to be divisible by this value.", + }, + ), + } + } + + RETURN_TYPES = ("IMAGE", "FLOAT") + RETURN_NAMES = ("image", "target_megapixels") + FUNCTION = "select" + CATEGORY = "image/analysis" + + @staticmethod + def _to_bhwc(image: torch.Tensor) -> torch.Tensor: + if not isinstance(image, torch.Tensor): + image = torch.tensor(image) + if image.dim() == 3: + image = image.unsqueeze(0) + if image.dim() != 4: + raise ValueError(f"Expected IMAGE tensor [B,H,W,C], got {tuple(image.shape)}") + if image.shape[-1] not in (3, 4): + raise ValueError("ImageMegapixelSelector expects RGB(A) images.") + if image.shape[-1] == 4: + image = image[..., :3] + return image.to(torch.float32).clamp(0.0, 1.0) + + @staticmethod + def _choose_megapixels(total_pixels: float) -> float: + megapixels = total_pixels / 1_000_000.0 + if megapixels <= 1.0: + return 1.0 + if megapixels >= 3.0: + return 3.0 + # Between 1 and 3: round to nearest legal breakpoint {1, 2, 3} with ties rounding up. + candidates = [1.0, 2.0, 3.0] + distances = [abs(megapixels - c) for c in candidates] + min_distance = min(distances) + # Pick the smallest candidate achieving the min distance, but prefer larger on ties. + best = max(c for c, d in zip(candidates, distances) if abs(d - min_distance) < 1e-9) + return best + + @staticmethod + def _round_to_multiple(value: int, modulus: int) -> int: + if modulus <= 1: + return max(1, value) + return max(modulus, int(round(value / modulus)) * modulus) + + @staticmethod + def _enforce_mod_dims( + raw_width: int, + raw_height: int, + aspect: float, + modulus: int, + ) -> Tuple[int, int]: + if modulus <= 1: + return max(1, raw_width), max(1, raw_height) + + raw_pixels = max(1, raw_width * raw_height) + raw_ratio = aspect + best = None + + width_candidates = set() + rounded_base = ImageMegapixelSelector._round_to_multiple(raw_width, modulus) + width_candidates.add(rounded_base) + for offset in range(-4, 5): + candidate = raw_width + offset * modulus + if candidate < modulus: + continue + width_candidates.add(ImageMegapixelSelector._round_to_multiple(candidate, modulus)) + + for w in sorted(width_candidates): + if w < modulus: + continue + ratio_based_height = w / raw_ratio + height_values = set() + for base in ( + ratio_based_height, + math.floor(ratio_based_height), + math.ceil(ratio_based_height), + raw_height, + ): + if base <= 0: + continue + height_values.add(ImageMegapixelSelector._round_to_multiple(int(round(base)), modulus)) + + for h in height_values: + if h < modulus: + continue + ratio = w / h + ratio_diff = abs(ratio - raw_ratio) + pixel_diff = abs((w * h) - raw_pixels) + dim_diff = abs(w - raw_width) + abs(h - raw_height) + key = (ratio_diff, pixel_diff, dim_diff) + if best is None or key < best[0]: + best = (key, w, h) + + if best is not None: + _, w_opt, h_opt = best + return max(1, w_opt), max(1, h_opt) + + # Fallback: simple rounding + return ( + ImageMegapixelSelector._round_to_multiple(raw_width, modulus), + ImageMegapixelSelector._round_to_multiple(raw_height, modulus), + ) + + @staticmethod + def _resize_lanczos(sample: torch.Tensor, width: int, height: int) -> torch.Tensor: + device = sample.device + dtype = sample.dtype + mode = "RGB" + + np_img = sample.detach().cpu().numpy() + np_img = np.clip(np_img, 0.0, 1.0) + np_img = (np_img * 255.0).round().astype(np.uint8) + + pil_image = Image.fromarray(np_img, mode=mode) + resized = pil_image.resize((int(width), int(height)), Image.LANCZOS) + arr = np.asarray(resized).astype(np.float32) / 255.0 + + if arr.ndim == 2: + arr = np.repeat(arr[..., None], 3, axis=-1) + if arr.shape[-1] > 3: + arr = arr[..., :3] + + tensor = torch.from_numpy(arr).to(device=device, dtype=dtype) + return tensor.clamp(0.0, 1.0) + + def select(self, image: torch.Tensor, divisible_by: int = 8): + tensor = self._to_bhwc(image) + if tensor.shape[0] != 1: + raise ValueError("ImageMegapixelSelector currently supports unbatched images (B=1).") + + sample = tensor[0] + height, width = sample.shape[:2] + total_pixels = float(height * width) + target_mp = self._choose_megapixels(total_pixels) + target_pixels = max(1.0, target_mp * 1_000_000.0) + + scale = math.sqrt(target_pixels / max(total_pixels, 1.0)) + raw_width = max(1, int(round(width * scale))) + raw_height = max(1, int(round(height * scale))) + modulus = max(1, int(divisible_by)) + + aspect = width / max(1.0, float(height)) + target_width, target_height = self._enforce_mod_dims(raw_width, raw_height, aspect, modulus) + + if target_width == width and target_height == height: + resized = sample + else: + resized = self._resize_lanczos(sample, target_width, target_height) + + actual_mp = float(target_width * target_height) / 1_000_000.0 + output = resized.unsqueeze(0).to(image.dtype if isinstance(image, torch.Tensor) else torch.float32) + + return (output, actual_mp) + + +NODE_CLASS_MAPPINGS = { + "ImageMegapixelSelector": ImageMegapixelSelector, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageMegapixelSelector": "Image Megapixel Selector", +} diff --git a/load_image_combined.py b/load_image_combined.py new file mode 100644 index 0000000..c476f09 --- /dev/null +++ b/load_image_combined.py @@ -0,0 +1,286 @@ +# load_image_combined.py +import os +import re +import glob +import hashlib +import numpy as np +from PIL import Image, ImageOps +import torch +import folder_paths + +# =========================== +# Shared utils / constants +# =========================== + +VALID_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".tif", ".tiff", ".bmp"} +_VALID_EXTS = VALID_EXTS # alias for any older references + + +def _coerce_str(value) -> str: + """Accept strings coming in as plain text, tuples, lists, or simple dict wrappers.""" + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, (tuple, list)) and value: + return _coerce_str(value[0]) + if isinstance(value, dict): + for key in ("string", "value", "text", "path"): + if key in value: + return _coerce_str(value[key]) + return str(value) + + +def _coerce_pattern(value) -> str: + s = _coerce_str(value).strip() + if not s or s.lower() in {"", "none"}: + return "*" + # Treat legacy numeric values (likely old batch_index) as request for default "*" + try: + float(s) + return "*" + except ValueError: + pass + return s + + +def _resolve_input_dir(input_dir: str) -> str: + """Resolve input directory against ComfyUI's input folder if relative.""" + if not input_dir: + return folder_paths.get_input_directory() + if os.path.isabs(input_dir): + return input_dir + return os.path.join(folder_paths.get_input_directory(), input_dir) + +def _basename_no_ext(filename: str, strip_numbers=False) -> str: + """Remove extension and optionally strip trailing (1)/(2) from filenames.""" + base = os.path.splitext(filename)[0] + if strip_numbers: + base = re.sub(r"\(\d+\)$", "", base).strip() + return base + +# ---------------------------------------------------------------- +# Batch listing index tracking (per-listing, in-process persistence) +# Keyed by (base, listing key), advances on each call unless +# repeat_last=True; not persisted to disk. +# ---------------------------------------------------------------- +_INDEX_STATE = {} # { base_dir: { listing_key: last_used_index } } + +def _listing_key(files_sorted, strip_numbers, pattern): + """ + Stable identifier for this specific listing. Combine directory path, + pattern, strip flag, and file names to disambiguate. + """ + if not files_sorted: + return "empty" + dir_path = os.path.dirname(files_sorted[0]) + parts = [os.path.basename(f).lower() for f in files_sorted] + seed = f"{os.path.abspath(dir_path).lower()}|pat={pattern}|strip={bool(strip_numbers)}|" + "|".join(parts) + return hashlib.sha256(seed.encode("utf-8")).hexdigest() + +def _peek_last_index(base_dir, key): + return int(_INDEX_STATE.get(base_dir, {}).get(key, -1)) + +def _choose_index_and_update(base_dir, key, n, repeat_last): + last = _peek_last_index(base_dir, key) + if last < 0: + next_idx = 0 + else: + next_idx = min(last, n - 1) if repeat_last else ((last + 1) % n) + d = _INDEX_STATE.get(base_dir) + if d is None: + d = {} + _INDEX_STATE[base_dir] = d + d[key] = next_idx + return next_idx + +# ============================================================ +# Node: Load Image (Combined) +# ============================================================ + +class LoadImageCombined: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + # Enumerate files directly from input dir, + # not folder_paths.get_filename_list("input") + input_dir = folder_paths.get_input_directory() + files = [ + f for f in os.listdir(input_dir) + if os.path.isfile(os.path.join(input_dir, f)) + ] + return { + "required": { + "mode": (["Single", "Batch"], {"default": "Single"}), + "input_dir": ("STRING", { + "default": "", + "multiline": False, + "placeholder": "Path or relative to ComfyUI/input" + }), + "output_dir": ("STRING", { + "default": "", + "multiline": False, + "placeholder": "Output Directory" + }), + "pattern": ("STRING", { + "default": "*", + "multiline": False, + "placeholder": "e.g., *.png" + }), + "strip_trailing_numbers": ("BOOLEAN", {"default": False}), + "repeat_last": ("BOOLEAN", {"default": False}), + "image": (sorted(files), {"image_upload": True}), + } + } + + CATEGORY = "image/io" + RETURN_TYPES = ("IMAGE", "STRING", "STRING", "INT", "INT") + RETURN_NAMES = ("IMAGE", "filename_no_ext", "output_dir", "width", "height") + FUNCTION = "load_image" + + def _load_pil(self, path): + im = Image.open(path) + im = ImageOps.exif_transpose(im) + # Always discard alpha (convert to RGB) to avoid downstream errors + if "A" in im.getbands(): + arr = np.array(im.convert("RGB")).astype(np.float32) / 255.0 + else: + arr = np.array(im.convert("RGB")).astype(np.float32) / 255.0 + return arr + + def _single_mode(self, image_choice, strip_numbers): + image_path = folder_paths.get_annotated_filepath(image_choice) + filename = os.path.basename(image_path) + filename_no_ext = _basename_no_ext(filename, strip_numbers) + arr = self._load_pil(image_path) + img_t = torch.from_numpy(arr)[None,] + h, w = arr.shape[0], arr.shape[1] + return img_t, filename_no_ext, int(w), int(h) + + def _gather_batch_files(self, input_dir, pattern): + if input_dir is None or str(input_dir).strip() == "": + raise ValueError("Batch mode requires 'input_dir'. Please specify a folder (absolute or relative to ComfyUI/input).") + base = _resolve_input_dir(input_dir) + if not os.path.isdir(base): + raise ValueError(f"Input directory not found: {base}") + + pat = pattern.strip() if pattern and pattern.strip() != "" else "*" + search_glob = os.path.join(base, pat) + candidates = glob.glob(search_glob) + files = [ + f for f in candidates + if os.path.isfile(f) and os.path.splitext(f)[1].lower() in _VALID_EXTS + ] + files_sorted = sorted(files, key=lambda s: s.lower()) + return base, pat, files_sorted + + def _batch_mode_auto_advance(self, input_dir, pattern, strip_numbers, repeat_last): + base, pat, files_sorted = self._gather_batch_files(input_dir, pattern) + if not files_sorted: + raise ValueError(f"No images found in '{base}' with pattern '{pat}'") + + key = _listing_key(files_sorted, strip_numbers, pat) + n = len(files_sorted) + use_idx = _choose_index_and_update(base, key, n, repeat_last) + + path = files_sorted[use_idx] + filename = os.path.basename(path) + filename_no_ext = _basename_no_ext(filename, strip_numbers) + arr = self._load_pil(path) + img_t = torch.from_numpy(arr)[None,] + h, w = arr.shape[0], arr.shape[1] + return img_t, filename_no_ext, int(w), int(h) + + def load_image(self, mode, input_dir, output_dir, pattern, strip_trailing_numbers, repeat_last, image): + input_dir = _coerce_str(input_dir).strip() + output_dir = _coerce_str(output_dir).strip() + pattern = _coerce_pattern(pattern) + if str(mode) == "Batch": + if not input_dir: + raise ValueError("Batch mode requires 'input_dir'. Please specify a folder (absolute or relative to ComfyUI/input).") + img_t, filename_no_ext, w, h = self._batch_mode_auto_advance( + input_dir, pattern, strip_trailing_numbers, repeat_last + ) + return img_t, filename_no_ext, str(output_dir or ""), w, h + else: + img_t, filename_no_ext, w, h = self._single_mode(image, strip_trailing_numbers) + return img_t, filename_no_ext, str(output_dir or ""), w, h + + @classmethod + def IS_CHANGED(s, mode, input_dir, output_dir, pattern, strip_trailing_numbers, repeat_last, image): + input_dir = _coerce_str(input_dir).strip() + output_dir = _coerce_str(output_dir).strip() + pattern = _coerce_pattern(pattern) + if str(mode) != "Batch": + image_path = folder_paths.get_annotated_filepath(image) + m = hashlib.sha256() + # Hash file bytes to detect actual content changes; also include strip flag + try: + with open(image_path, "rb") as f: + while True: + chunk = f.read(65536) + if not chunk: + break + m.update(chunk) + except Exception: + m.update(b"ERROR:cannot_read_single_image") + m.update(f"|strip={bool(strip_trailing_numbers)}|".encode("utf-8")) + m.update(f"|output_dir={str(output_dir)}|".encode("utf-8")) + return m.digest().hex() + + # Batch mode hashing: reflect directory listing, file sizes/mtimes, last-used index, strip flag, and pattern. + if not input_dir: + m = hashlib.sha256() + m.update(b"ERROR:no_input_dir_for_batch") + return m.digest().hex() + + base = _resolve_input_dir(input_dir) + pat = pattern if pattern else "*" + candidates = [ + f for f in glob.glob(os.path.join(base, pat)) + if os.path.isfile(f) and os.path.splitext(f)[1].lower() in _VALID_EXTS + ] + files_sorted = sorted(candidates, key=lambda s: s.lower()) + + m = hashlib.sha256() + for fp in files_sorted: + try: + st = os.stat(fp) + m.update(os.path.abspath(fp).lower().encode("utf-8")) + m.update(str(st.st_size).encode("utf-8")) + m.update(str(int(st.st_mtime)).encode("utf-8")) + except Exception: + m.update(os.path.abspath(fp).lower().encode("utf-8")) + m.update(b"\n") + + key = _listing_key(files_sorted, strip_trailing_numbers, pat) + last_used = _peek_last_index(base, key) + m.update(f"|last_used={last_used}|".encode("utf-8")) + m.update(f"|strip={bool(strip_trailing_numbers)}|".encode("utf-8")) + m.update(f"|pattern={pat}|".encode("utf-8")) + m.update(f"|output_dir={str(output_dir)}|".encode("utf-8")) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, mode, input_dir, output_dir, pattern, strip_trailing_numbers, repeat_last, image): + input_dir = _coerce_str(input_dir).strip() + pattern = _coerce_pattern(pattern) + if str(mode) == "Batch": + if not input_dir: + return "Batch mode requires 'input_dir'." + base = _resolve_input_dir(input_dir) + if not os.path.isdir(base): + return f"Input directory not found: {base}" + pat = pattern.strip() if pattern and pattern.strip() != "" else "*" + candidates = [ + f for f in glob.glob(os.path.join(base, pat)) + if os.path.isfile(f) and os.path.splitext(f)[1].lower() in _VALID_EXTS + ] + if not candidates: + return f"No images found in '{base}' with pattern '{pat}'" + return True + if not folder_paths.exists_annotated_filepath(image): + return f"Invalid image file: {image}" + return True diff --git a/multiprompt.py b/multiprompt.py new file mode 100644 index 0000000..ab49ef8 --- /dev/null +++ b/multiprompt.py @@ -0,0 +1,47 @@ +class MultiPromptNode: + """Simple shim that exposes four multiline prompt sockets.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "subject_mask": ( + "STRING", + {"default": "", "multiline": True, "placeholder": "subject mask"}, + ), + "inpaint_prompt": ( + "STRING", + {"default": "", "multiline": True, "placeholder": "inpaint prompt"}, + ), + "outpaint_prompt": ( + "STRING", + {"default": "", "multiline": True, "placeholder": "outpaint prompt"}, + ), + "stitch_mask": ( + "STRING", + {"default": "", "multiline": True, "placeholder": "stitch mask"}, + ), + } + } + + RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING") + RETURN_NAMES = ("subject_mask", "inpaint_prompt", "outpaint_prompt", "stitch_mask") + FUNCTION = "apply" + CATEGORY = "config" + + def apply(self, subject_mask, inpaint_prompt, outpaint_prompt, stitch_mask): + return ( + str(subject_mask), + str(inpaint_prompt), + str(outpaint_prompt), + str(stitch_mask), + ) + + +NODE_CLASS_MAPPINGS = { + "MultiPromptNode": MultiPromptNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MultiPromptNode": "Multi-Prompt", +} diff --git a/other_config.py b/other_config.py new file mode 100644 index 0000000..140c190 --- /dev/null +++ b/other_config.py @@ -0,0 +1,95 @@ +class OtherConfigNode: + """ + Outputs: + - initial_upscale (megapixels) + - final_upscale (megapixels) + - tight_crop (BOOLEAN, default False) + - blend_factor (FLOAT 0-1) + - input_dir (STRING) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "initial_upscale_megapixels": ( + "FLOAT", + { + "default": 1.0, + "min": 0.1, + "max": 100.0, + "step": 0.1, + "label": "initial_upscale (megapixels)", + }, + ), + "final_upscale_megapixels": ( + "FLOAT", + { + "default": 4.0, + "min": 0.1, + "max": 100.0, + "step": 0.1, + "label": "final_upscale (megapixels)", + }, + ), + "tight_crop": ("BOOLEAN", {"default": False}), + "blend_factor": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + "stitch_bypass_mask": ("BOOLEAN", {"default": False}), + "input_dir": ( + "STRING", + {"default": "", "multiline": False, "placeholder": "Input Directory"}, + ), + } + } + + RETURN_TYPES = ( + "FLOAT", + "FLOAT", + "BOOLEAN", + "FLOAT", + "BOOLEAN", + "STRING", + ) + RETURN_NAMES = ( + "initial_upscale (megapixels)", + "final_upscale (megapixels)", + "tight_crop", + "blend_factor", + "stitch_bypass_mask", + "input_dir", + ) + FUNCTION = "apply" + CATEGORY = "config" + + def apply( + self, + initial_upscale_megapixels, + final_upscale_megapixels, + tight_crop, + blend_factor, + stitch_bypass_mask, + input_dir, + ): + initial_upscale_megapixels = max(0.1, float(initial_upscale_megapixels)) + final_upscale_megapixels = max(0.1, float(final_upscale_megapixels)) + blend_factor = min(1.0, max(0.0, float(blend_factor))) + return ( + initial_upscale_megapixels, + final_upscale_megapixels, + bool(tight_crop), + blend_factor, + bool(stitch_bypass_mask), + str(input_dir), + ) + + +NODE_CLASS_MAPPINGS = { + "OtherConfigNode": OtherConfigNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "OtherConfigNode": "Other Config", +} diff --git a/outpaint_config.py b/outpaint_config.py new file mode 100644 index 0000000..6e33785 --- /dev/null +++ b/outpaint_config.py @@ -0,0 +1,236 @@ +import torch + + +class OutpaintConfigNode: + """ + Control panel for outpainting preferences (no image). + Outputs raw settings for Set/Get. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mode": (["Percent", "Pixels"], {"default": "Percent"}), + "gravity": ( + [ + "center", + "left", + "right", + "top", + "bottom", + "top left", + "top right", + "bottom right", + "bottom left", + ], + {"default": "center"}, + ), + "horizontal_percent": ( + "FLOAT", + {"default": 20.0, "min": 0.0, "max": 10000.0, "step": 0.1}, + ), + "vertical_percent": ( + "FLOAT", + {"default": 10.0, "min": 0.0, "max": 10000.0, "step": 0.1}, + ), + "left_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "right_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "top_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "bottom_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + } + } + + RETURN_TYPES = ( + "STRING", + "STRING", + "FLOAT", + "FLOAT", + "INT", + "INT", + "INT", + "INT", + ) + RETURN_NAMES = ( + "mode", + "gravity", + "horizontal_percent", + "vertical_percent", + "left_px", + "right_px", + "top_px", + "bottom_px", + ) + FUNCTION = "apply" + CATEGORY = "config" + + def apply( + self, + mode, + gravity, + horizontal_percent, + vertical_percent, + left_px, + right_px, + top_px, + bottom_px, + ): + + horizontal_percent = max(0.0, float(horizontal_percent)) + vertical_percent = max(0.0, float(vertical_percent)) + left_px = max(0, int(left_px)) + right_px = max(0, int(right_px)) + top_px = max(0, int(top_px)) + bottom_px = max(0, int(bottom_px)) + + return ( + mode, + gravity, + horizontal_percent, + vertical_percent, + left_px, + right_px, + top_px, + bottom_px, + ) + + +class OutpaintPaddingComputeNode: + """ + Compute absolute pixel paddings (left, top, right, bottom) for outpainting. + - Pixels mode: gravity ignored; enforce even final size. + - Percent mode: compute by % + gravity (including corners); enforce even final size. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mode": ("STRING", {"default": "Percent"}), + "gravity": ("STRING", {"default": "center"}), + "horizontal_percent": ( + "FLOAT", + {"default": 20.0, "min": 0.0, "max": 10000.0, "step": 0.1}, + ), + "vertical_percent": ( + "FLOAT", + {"default": 10.0, "min": 0.0, "max": 10000.0, "step": 0.1}, + ), + "left_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "right_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "top_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "bottom_px": ("INT", {"default": 0, "min": 0, "max": 1000000}), + } + } + + RETURN_TYPES = ("INT", "INT", "INT", "INT") + RETURN_NAMES = ("left", "top", "right", "bottom") + FUNCTION = "compute" + CATEGORY = "image/transform" + + def _get_wh(self, image_tensor): + if not torch.is_tensor(image_tensor): + image_tensor = torch.tensor(image_tensor) + if image_tensor.dim() != 4: + raise ValueError("Expected IMAGE tensor of shape [B,H,W,C].") + _, H, W, _ = image_tensor.shape + return int(W), int(H) + + def _enforce_even_final_size(self, W, H, left, top, right, bottom): + if (W + left + right) % 2 != 0: + right += 1 + if (H + top + bottom) % 2 != 0: + bottom += 1 + return left, top, right, bottom + + def compute( + self, + image, + mode, + gravity, + horizontal_percent, + vertical_percent, + left_px, + right_px, + top_px, + bottom_px, + ): + + W, H = self._get_wh(image) + + g = str(gravity).lower().strip() + valid = { + "center", + "left", + "right", + "top", + "bottom", + "top left", + "top right", + "bottom right", + "bottom left", + } + if g not in valid: + g = "center" + + hp = max(0.0, float(horizontal_percent)) + vp = max(0.0, float(vertical_percent)) + lp = max(0, int(left_px)) + rp = max(0, int(right_px)) + tp = max(0, int(top_px)) + bp = max(0, int(bottom_px)) + + if mode == "Pixels": + left, top, right, bottom = lp, tp, rp, bp + left, top, right, bottom = self._enforce_even_final_size( + W, H, left, top, right, bottom + ) + return (left, top, right, bottom) + + if hp == 0.0: + left = 0 + right = 0 + else: + horz_total = int(round(W * hp / 100.0)) + if g in ("left", "top left", "bottom left"): + left, right = horz_total, 0 + elif g in ("right", "top right", "bottom right"): + right, left = horz_total, 0 + else: + left = horz_total // 2 + right = horz_total - left + + if vp == 0.0: + top = 0 + bottom = 0 + else: + vert_total = int(round(H * vp / 100.0)) + if g in ("top", "top left", "top right"): + top, bottom = vert_total, 0 + elif g in ("bottom", "bottom left", "bottom right"): + bottom, top = vert_total, 0 + else: + top = vert_total // 2 + bottom = vert_total - top + + left, top, right, bottom = self._enforce_even_final_size( + W, H, left, top, right, bottom + ) + return ( + max(0, int(left)), + max(0, int(top)), + max(0, int(right)), + max(0, int(bottom)), + ) + + +NODE_CLASS_MAPPINGS = { + "OutpaintConfigNode": OutpaintConfigNode, + "OutpaintPaddingComputeNode": OutpaintPaddingComputeNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "OutpaintConfigNode": "Outpaint Config", + "OutpaintPaddingComputeNode": "Outpaint Padding Compute", +} diff --git a/paired_image_loader.py b/paired_image_loader.py new file mode 100644 index 0000000..d220922 --- /dev/null +++ b/paired_image_loader.py @@ -0,0 +1,481 @@ +from __future__ import annotations + +import hashlib +import os +import re +import uuid +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Iterable, List, Tuple + +import numpy as np +import torch +from PIL import Image, ImageOps + +import node_helpers + +_IMAGE_EXTENSIONS = { + ".png", + ".jpg", + ".jpeg", + ".bmp", + ".gif", + ".tif", + ".tiff", + ".webp", +} + +_TRAILING_NUMBER_PATTERN = re.compile(r"\s*\(\d+\)$") +_NATURAL_KEY_PATTERN = re.compile(r"(\d+)") + + +def _natural_sort_key(value: str) -> List[object]: + return [ + int(chunk) if chunk.isdigit() else chunk.lower() + for chunk in _NATURAL_KEY_PATTERN.split(value) + ] + + +def _normalize_base(stem: str, strip_numbers: bool) -> str: + base = stem + if strip_numbers: + base = _TRAILING_NUMBER_PATTERN.sub("", base) + return base.lower() + + +def _resolve_directory(path_str: str) -> Path: + path = Path(path_str).expanduser() + if not path.is_absolute(): + path = (Path.cwd() / path).resolve() + else: + path = path.resolve() + return path + + +@dataclass(frozen=True) +class _FileEntry: + path: Path + name: str + stem: str + normalized_key: str + ext: str + mtime_ns: int + size: int + + +@dataclass(frozen=True) +class _Pair: + key: str + display_name: str + source: _FileEntry + output: _FileEntry + + +@dataclass +class _NodeState: + pairs: List[_Pair] = field(default_factory=list) + index: int = -1 + signature: Tuple | None = None + warn_signature: Tuple | None = None + + +_STATE: Dict[str, _NodeState] = {} +_SIGNATURE_INDEX: Dict[Tuple, int] = {} + + +def _coerce_str(value) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, (list, tuple)) and value: + return _coerce_str(value[0]) + return str(value) + + +class PortraitUtils_PairedImageLoader: + CATEGORY = "PortraitUtils" + RETURN_TYPES = ("IMAGE", "IMAGE", "STRING") + FUNCTION = "load_next_pair" + RETURN_NAMES = ("source_image", "output_image", "filename") + NOT_IDEMPOTENT = True + + def __init__(self): + self._state_key = uuid.uuid4().hex + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "source_dir": ("STRING", {"default": "", "multiline": False}), + "output_dir": ("STRING", {"default": "", "multiline": False}), + "reverse": ("BOOLEAN", {"default": False}), + "strip_trailing_numbers": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Strip trailing \" (n)\" suffixes before matching.", + }, + ), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, + } + + def load_next_pair( + self, + source_dir, + output_dir, + reverse=False, + strip_trailing_numbers=False, + unique_id=None, + ): + if isinstance(source_dir, list): + source_dir = source_dir[0] + if isinstance(output_dir, list): + output_dir = output_dir[0] + if isinstance(reverse, list): + reverse = reverse[0] + if isinstance(strip_trailing_numbers, list): + strip_trailing_numbers = strip_trailing_numbers[0] + if isinstance(unique_id, list) and unique_id: + unique_id = unique_id[0] + + reverse = bool(reverse) + strip_trailing_numbers = bool(strip_trailing_numbers) + + current_key = getattr(self, "_state_key", None) + if unique_id is not None: + state_key = str(unique_id) + else: + if current_key is None: + current_key = uuid.uuid4().hex + self._state_key = current_key + state_key = current_key + + migrated_state = None + if current_key is not None and state_key != current_key: + migrated_state = _STATE.pop(current_key, None) + + self._state_key = state_key + + if not source_dir: + raise ValueError("PortraitUtils_PairedImageLoader: source_dir is required.") + if not output_dir: + raise ValueError("PortraitUtils_PairedImageLoader: output_dir is required.") + + src_path = _resolve_directory(source_dir) + out_path = _resolve_directory(output_dir) + + if not src_path.exists() or not src_path.is_dir(): + raise ValueError( + f"PortraitUtils_PairedImageLoader: source_dir not found: {src_path}" + ) + if not out_path.exists() or not out_path.is_dir(): + raise ValueError( + f"PortraitUtils_PairedImageLoader: output_dir not found: {out_path}" + ) + + state = _STATE.setdefault(state_key, migrated_state or _NodeState()) + + pairs, signature, warnings = self._scan_directories( + src_path, out_path, strip_trailing_numbers + ) + + state.pairs = pairs + if signature != state.signature: + state.signature = signature + if state.index >= len(pairs): + state.index = len(pairs) - 1 + + if signature != state.warn_signature and any(warnings.values()): + self._emit_warnings(warnings, src_path, out_path) + + state.warn_signature = signature + + if not pairs: + raise RuntimeError( + "PortraitUtils_PairedImageLoader: no matching images found " + f"between {src_path} and {out_path}." + ) + + if state.index < 0: + next_index = len(pairs) - 1 if reverse else 0 + else: + next_index = (state.index - 1) % len(pairs) if reverse else ( + (state.index + 1) % len(pairs) + ) + + state.index = next_index + pair = pairs[next_index] + + source_tensor = self._load_image(pair.source.path) + output_tensor = self._load_image(pair.output.path) + + _SIGNATURE_INDEX[signature] = state.index + + print( + "[PortraitUtils_PairedImageLoader] " + f"{pair.display_name} ({next_index + 1}/{len(pairs)})" + ) + + return source_tensor, output_tensor, pair.source.name + + @classmethod + def _scan_directories( + cls, + source_dir: Path, + output_dir: Path, + strip_trailing_numbers: bool, + ) -> Tuple[List[_Pair], Tuple, Dict[str, object]]: + source_entries = cls._gather_entries(source_dir, strip_trailing_numbers) + output_entries = cls._gather_entries(output_dir, strip_trailing_numbers) + + source_map: Dict[str, List[_FileEntry]] = {} + output_map: Dict[str, List[_FileEntry]] = {} + collisions: Dict[str, List[str]] = {} + + for entry in source_entries: + source_map.setdefault(entry.normalized_key, []).append(entry) + for entry in output_entries: + output_map.setdefault(entry.normalized_key, []).append(entry) + + if strip_trailing_numbers: + for mapping in (source_map, output_map): + for key, entries in mapping.items(): + if len(entries) > 1: + existing = collisions.setdefault(key, []) + existing.extend(e.name for e in entries) + + for mapping in (source_map, output_map): + for key in mapping: + mapping[key].sort(key=lambda e: _natural_sort_key(e.name)) + + shared_keys = sorted( + set(source_map).intersection(output_map), + key=lambda k: _natural_sort_key( + source_map.get(k, output_map[k])[0].stem + ), + ) + + pairs: List[_Pair] = [] + extension_mismatch: List[str] = [] + + for key in shared_keys: + src_candidates = source_map[key] + out_candidates = output_map[key] + src_entry, out_entry, matched_ext = cls._select_pair( + src_candidates, out_candidates + ) + if matched_ext is None: + extension_mismatch.append(src_entry.stem) + display_name = src_entry.stem or out_entry.stem + pairs.append(_Pair(key, display_name, src_entry, out_entry)) + + source_only = cls._flatten_unmatched(source_map, shared_keys) + output_only = cls._flatten_unmatched(output_map, shared_keys) + + signature = ( + strip_trailing_numbers, + str(source_dir), + str(output_dir), + tuple(sorted((e.name.lower(), e.mtime_ns, e.size) for e in source_entries)), + tuple(sorted((e.name.lower(), e.mtime_ns, e.size) for e in output_entries)), + ) + + warnings = { + "source_only": source_only, + "output_only": output_only, + "collisions": collisions, + "extension_mismatch": extension_mismatch, + } + + return pairs, signature, warnings + + @staticmethod + def _gather_entries( + directory: Path, + strip_trailing_numbers: bool, + ) -> List[_FileEntry]: + entries: List[_FileEntry] = [] + with os.scandir(directory) as iterator: + for item in iterator: + if not item.is_file(): + continue + name = item.name + ext = Path(name).suffix.lower() + if ext not in _IMAGE_EXTENSIONS: + continue + stat = item.stat() + stem = Path(name).stem + normalized_key = _normalize_base(stem, strip_trailing_numbers) + entry = _FileEntry( + path=directory / name, + name=name, + stem=stem, + normalized_key=normalized_key, + ext=ext, + mtime_ns=stat.st_mtime_ns, + size=stat.st_size, + ) + entries.append(entry) + return entries + + @staticmethod + def _select_pair( + source_candidates: List[_FileEntry], + output_candidates: List[_FileEntry], + ) -> Tuple[_FileEntry, _FileEntry, str | None]: + source_by_ext: Dict[str, List[_FileEntry]] = {} + output_by_ext: Dict[str, List[_FileEntry]] = {} + + for entry in source_candidates: + source_by_ext.setdefault(entry.ext, []).append(entry) + for entry in output_candidates: + output_by_ext.setdefault(entry.ext, []).append(entry) + + for entries in source_by_ext.values(): + entries.sort(key=lambda e: _natural_sort_key(e.name)) + for entries in output_by_ext.values(): + entries.sort(key=lambda e: _natural_sort_key(e.name)) + + for ext in sorted(output_by_ext): + if ext in source_by_ext: + return source_by_ext[ext][0], output_by_ext[ext][0], ext + + return ( + source_candidates[0], + output_candidates[0], + None, + ) + + @staticmethod + def _flatten_unmatched( + mapping: Dict[str, List[_FileEntry]], + matched_keys: Iterable[str], + ) -> List[str]: + matched = set(matched_keys) + leftovers: List[str] = [] + for key, entries in mapping.items(): + if key in matched: + continue + leftovers.extend(entry.name for entry in entries) + leftovers.sort(key=_natural_sort_key) + return leftovers + + @staticmethod + def _emit_warnings( + warnings: Dict[str, object], + source_dir: Path, + output_dir: Path, + ) -> None: + source_only = warnings.get("source_only", []) + output_only = warnings.get("output_only", []) + collisions = warnings.get("collisions", {}) + ext_mismatch = warnings.get("extension_mismatch", []) + + def _summarize(names: List[str]) -> str: + if len(names) <= 5: + return ", ".join(names) + head = ", ".join(names[:5]) + return f"{head}, … (+{len(names) - 5} more)" + + if source_only: + print( + "[PortraitUtils_PairedImageLoader] Unmatched source files " + f"({len(source_only)}) in {source_dir}: {_summarize(source_only)}" + ) + if output_only: + print( + "[PortraitUtils_PairedImageLoader] Unmatched output files " + f"({len(output_only)}) in {output_dir}: {_summarize(output_only)}" + ) + if collisions: + for key, files in collisions.items(): + print( + "[PortraitUtils_PairedImageLoader] Normalised name collision " + f"for '{key}': {_summarize(sorted(set(files), key=_natural_sort_key))}" + ) + if ext_mismatch: + print( + "[PortraitUtils_PairedImageLoader] Chosen fallback matches with differing " + f"extensions ({len(ext_mismatch)}): {_summarize(ext_mismatch)}" + ) + + @staticmethod + def _load_image(path: Path) -> torch.Tensor: + img = node_helpers.pillow(Image.open, path) + try: + img = ImageOps.exif_transpose(img) + if img.mode != "RGB": + if "A" in img.getbands(): + img = img.convert("RGBA") + else: + img = img.convert("RGB") + if img.mode == "RGBA": + img = img.convert("RGB") + array = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(array)[None, ...] + return tensor + finally: + img.close() + + @classmethod + def IS_CHANGED( + cls, + source_dir, + output_dir, + reverse=False, + strip_trailing_numbers=False, + unique_id=None, + ): + source_dir = _coerce_str(source_dir).strip() + output_dir = _coerce_str(output_dir).strip() + if isinstance(reverse, list): + reverse = reverse[0] + if isinstance(strip_trailing_numbers, list): + strip_trailing_numbers = strip_trailing_numbers[0] + + reverse = bool(reverse) + strip_trailing_numbers = bool(strip_trailing_numbers) + + digest = hashlib.sha256() + digest.update(f"reverse={reverse}".encode("utf-8")) + + if not source_dir or not output_dir: + digest.update(f"|missing_dirs|{source_dir}|{output_dir}|".encode("utf-8")) + return digest.hexdigest() + + try: + src_path = _resolve_directory(source_dir) + out_path = _resolve_directory(output_dir) + except Exception as exc: + digest.update(f"|resolve_error|{exc}|".encode("utf-8")) + return digest.hexdigest() + + if not src_path.exists() or not src_path.is_dir(): + digest.update(f"|src_missing|{src_path}|".encode("utf-8")) + return digest.hexdigest() + if not out_path.exists() or not out_path.is_dir(): + digest.update(f"|out_missing|{out_path}|".encode("utf-8")) + return digest.hexdigest() + + pairs, signature, _ = cls._scan_directories( + src_path, out_path, strip_trailing_numbers + ) + + digest.update(str(signature).encode("utf-8")) + digest.update(f"|pairs={len(pairs)}|".encode("utf-8")) + last_index = _SIGNATURE_INDEX.get(signature, -1) + digest.update(f"|last_index={last_index}|".encode("utf-8")) + return digest.hexdigest() + + +NODE_CLASS_MAPPINGS = { + "PortraitUtils_PairedImageLoader": PortraitUtils_PairedImageLoader, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PortraitUtils_PairedImageLoader": "Paired Image Loader", +} diff --git a/smart_crop.py b/smart_crop.py new file mode 100644 index 0000000..1413cb8 --- /dev/null +++ b/smart_crop.py @@ -0,0 +1,433 @@ +import math +import numpy as np +import torch + + +def _to_numpy_mask(m): + if torch is not None and isinstance(m, torch.Tensor): + arr = m.detach().cpu().numpy() + elif isinstance(m, dict) and "mask" in m: + v = m["mask"] + arr = v.detach().cpu().numpy() if torch is not None and isinstance(v, torch.Tensor) else np.asarray(v) + else: + arr = np.asarray(m) + arr = np.squeeze(arr) + if arr.ndim == 3: + arr = arr.max(axis=0) if arr.shape[0] in (1, 3, 4) else arr.max(axis=-1) + if arr.ndim != 2: + h, w = arr.shape[-2], arr.shape[-1] + arr = arr.reshape(h, w) + return arr.astype(np.float32, copy=False) + + +def _normalize01(m): + mmax = float(m.max()) if m.size else 0.0 + if mmax <= 0: + return m + if mmax > 1.5: + return (m / 255.0).astype(np.float32, copy=False) + return m + + +def _border_mass_ratio(m, border_px=8): + H, W = m.shape + border_px = int(max(1, min(border_px, min(H, W) // 4))) + total = float(m.sum()) + 1e-12 + top = m[:border_px, :].sum() + bottom = m[-border_px:, :].sum() + left = m[:, :border_px].sum() + right = m[:, -border_px:].sum() + corners = ( + m[:border_px, :border_px].sum() + + m[:border_px, -border_px:].sum() + + m[-border_px:, :border_px].sum() + + m[-border_px:, -border_px:].sum() + ) + border = top + bottom + left + right - corners + return float(border) / total + + +def _auto_pick_foreground(m, mode="auto"): + m = _normalize01(m) + if mode == "false": + fg, pick = m, "fg" + elif mode == "true": + fg, pick = 1.0 - m, "bg" + else: + r_fg = _border_mass_ratio(m) + r_bg = _border_mass_ratio(1.0 - m) + if r_bg < r_fg: + fg, pick = 1.0 - m, "auto->bg" + else: + fg, pick = m, "auto->fg" + return fg, pick + + +def _quantile_bounds(mass, qL, qR, qT, qB, min_span): + H, W = mass.shape + eps = 1e-8 + c = mass.sum(axis=0) + eps + r = mass.sum(axis=1) + eps + cc = np.cumsum(c) + rc = np.cumsum(r) + ct = float(cc[-1]) + rt = float(rc[-1]) + + def qidx(cum, total, q): + return int(np.clip(np.searchsorted(cum, total * float(q), "left"), 0, len(cum) - 1)) + + x0 = qidx(cc, ct, qL) + x1 = qidx(cc, ct, qR) + y0 = qidx(rc, rt, qT) + y1 = qidx(rc, rt, qB) + if x1 < x0: + x0, x1 = x1, x0 + if y1 < y0: + y0, y1 = y1, y0 + if (x1 - x0 + 1) < min_span: + cx = (x0 + x1) // 2 + half = max(1, min_span // 2) + x0 = max(0, cx - half) + x1 = min(W - 1, cx + half) + if (y1 - y0 + 1) < min_span: + cy = (y0 + y1) // 2 + half = max(1, min_span // 2) + y0 = max(0, cy - half) + y1 = min(H - 1, cy + half) + return int(x0), int(y0), int(x1 - x0 + 1), int(y1 - y0 + 1) + + +def _pad_rect(x, y, w, h, pad, W, H): + x2 = max(0, x - pad) + y2 = max(0, y - pad) + w2 = min(W - x2, w + 2 * pad) + h2 = min(H - y2, h + 2 * pad) + return int(x2), int(y2), int(w2), int(h2) + + +def _container_for_ar(x, y, w, h, R, W, H): + """ + Smallest rectangle of aspect R (=w/h) that contains (x,y,w,h), centered, + then clamped to image. If it can't fit (too large), return None. + """ + wA = max(w, int(np.ceil(h * R))) + hA = int(np.ceil(wA / R)) + hB = max(h, int(np.ceil(w / R))) + wB = int(np.ceil(hB * R)) + cand = [] + if wA <= W and hA <= H: + cand.append((wA, hA)) + if wB <= W and hB <= H: + cand.append((wB, hB)) + if not cand: + return None + cw, ch = min(cand, key=lambda t: t[0] * t[1]) + + cx = x + w / 2.0 + cy = y + h / 2.0 + X = int(np.round(cx - cw / 2.0)) + Y = int(np.round(cy - ch / 2.0)) + X = max(0, min(W - cw, X)) + Y = max(0, min(H - ch, Y)) + return X, Y, int(cw), int(ch) + + +class MQBBoxMin: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "invert_mask": (["auto", "false", "true"], {"default": "auto"}), + "q_left": ("FLOAT", {"default": 0.005, "min": 0.0, "max": 0.2, "step": 0.001}), + "q_right": ("FLOAT", {"default": 0.995, "min": 0.8, "max": 1.0, "step": 0.001}), + "q_top": ("FLOAT", {"default": 0.005, "min": 0.0, "max": 0.2, "step": 0.001}), + "q_bottom": ("FLOAT", {"default": 0.995, "min": 0.8, "max": 1.0, "step": 0.001}), + "min_span_px": ("INT", {"default": 8, "min": 1, "max": 2048}), + "tight_crop": ( + "BOOLEAN", + { + "default": True, + "tooltip": "ON: tight subject bbox with smart AR; OFF: full frame (minimal AR crop downstream).", + }, + ), + } + } + + RETURN_TYPES = ("INT", "INT", "INT", "INT", "STRING") + RETURN_NAMES = ("x", "y", "w", "h", "debug") + FUNCTION = "run" + CATEGORY = "utils/debug" + + def run(self, mask, invert_mask, q_left, q_right, q_top, q_bottom, min_span_px, tight_crop=True): + m = _to_numpy_mask(mask) + H, W = m.shape + + if not tight_crop: + dbg = f"tight=False full=({0},{0},{W},{H})" + return (0, 0, int(W), int(H), dbg) + + # --- original tight bbox --- + fg, picked = _auto_pick_foreground(m, invert_mask) + x, y, w, h = _quantile_bounds(fg, q_left, q_right, q_top, q_bottom, int(min_span_px)) + + # --- smart AR (minimal-area container), internal only --- + pad_px = max(8, int(round(0.02 * min(H, W)))) + sx, sy, sw, sh = _pad_rect(x, y, w, h, pad_px, W, H) + + # Baseline: container using current image AR + R0 = float(W) / max(1.0, float(H)) + base = _container_for_ar(sx, sy, sw, sh, R0, W, H) + if base is None: + dbg = f"tight=True pick={picked} mq=({x},{y},{w},{h})" + return (x, y, w, h, dbg) + bx, by, bw, bh = base + base_area = bw * bh + + # Try common ARs in BOTH orientations (includes 9:16) + ratios = [ + (1, 1), + (2, 3), + (3, 2), + (3, 4), + (4, 3), + (9, 16), + (16, 9), + ] + + best_rect = base + best_area = base_area + best_info = f"AR={W}:{H}" + + for aw, ah in ratios: + R = float(aw) / float(ah) + cand = _container_for_ar(sx, sy, sw, sh, R, W, H) + if cand is None: + continue + cx, cy, cw, ch = cand + area = cw * ch + if area < best_area: + best_area = area + best_rect = cand + best_info = f"AR={aw}:{ah}" + + rx, ry, rw, rh = best_rect + dbg = f"tight=True pick={picked} mq=({x},{y},{w},{h}) -> smart({rx},{ry},{rw},{rh}) {best_info}, pad={pad_px}" + return (int(rx), int(ry), int(rw), int(rh), dbg) + + +def _get_hw(image): + t = image["image"] if isinstance(image, dict) and "image" in image else image + s = tuple(getattr(t, "shape", np.asarray(t).shape)) + if len(s) == 4: + H, W = (s[1], s[2]) if s[-1] in (1, 3, 4) else (s[2], s[3]) + elif len(s) == 3: + H, W = (s[0], s[1]) if s[0] not in (1, 3, 4) else (s[1], s[2]) + else: + H, W = s[-2], s[-1] + return int(H), int(W) + + +def _parse_aspects(s): + out = [] + for p in s.split(","): + p = p.strip() + if not p: + continue + if ":" in p: + a, b = p.split(":", 1) + else: + a, b = p.replace("X", "x").split("x", 1) + try: + a = float(a) + b = float(b) + if a > 0 and b > 0: + out.append((a / b, f"{int(a)}:{int(b)}")) + except Exception: + pass + return out or [(2 / 3.0, "2:3")] + + +def _closest_ratio(target_ratio, candidates): + tr = float(target_ratio) + best = None + for r, tag in candidates: + d = abs(math.log(r) - math.log(tr)) + if best is None or d < best[0]: + best = (d, r, tag) + return best[1], best[2] + + +def _cover_min_rect(x0, y0, w0, h0, ratio): + """ + Smallest rectangle of aspect 'ratio' that fully covers the rect (x0,y0,w0,h0). + Returns floats (xa,ya,wa,ha) without clamping to image. + """ + if w0 <= 0 or h0 <= 0: + return x0, y0, w0, h0 + r0 = w0 / max(1.0, h0) + if r0 >= ratio: + wa = w0 + ha = w0 / ratio + else: + ha = h0 + wa = h0 * ratio + cx = x0 + w0 / 2.0 + cy = y0 + h0 / 2.0 + return cx - wa / 2.0, cy - ha / 2.0, wa, ha + + +class FitAspectHeadSafe: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "x": ("INT",), + "y": ("INT",), + "w": ("INT",), + "h": ("INT",), + "aspects_csv": ( + "STRING", + {"default": "2:3,3:4,1:1,9:16,16:9,5:8,8:5"}, + ), + "match_to": (["image", "mq_box"], {"default": "mq_box"}), + "headroom_ratio": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 0.5, "step": 0.01}), + "footroom_ratio": ("FLOAT", {"default": 0.06, "min": 0.0, "max": 0.5, "step": 0.01}), + "side_margin_ratio": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 0.5, "step": 0.01}), + "bottom_priority": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 1.0, "step": 0.05}), + "horiz_gravity": (["center", "left", "right"], {"default": "center"}), + } + } + + RETURN_TYPES = ("INT", "INT", "INT", "INT", "STRING", "STRING") + RETURN_NAMES = ("w", "h", "x", "y", "aspect_used", "debug") + FUNCTION = "run" + CATEGORY = "utils/geometry" + + def run( + self, + image, + x, + y, + w, + h, + aspects_csv, + match_to, + headroom_ratio, + footroom_ratio, + side_margin_ratio, + bottom_priority, + horiz_gravity, + ): + + H, W = _get_hw(image) + img_ratio = W / max(1.0, H) + mq_ratio = w / max(1.0, h) + target_ratio = img_ratio if match_to == "image" else mq_ratio + + aspects = _parse_aspects(aspects_csv) + ratio, tag = _closest_ratio(target_ratio, aspects) + + req_x0 = x + req_y0 = y + req_x1 = x + w + req_y1 = y + h + + guess_h = H + head_px = headroom_ratio * guess_h + foot_px = footroom_ratio * guess_h + side_px = side_margin_ratio * (guess_h * ratio) + + req_x0p = req_x0 - side_px + req_x1p = req_x1 + side_px + req_y0p = req_y0 - head_px + req_y1p = req_y1 + foot_px + + xa, ya, wa, ha = _cover_min_rect( + req_x0p, + req_y0p, + max(1.0, req_x1p - req_x0p), + max(1.0, req_y1p - req_y0p), + ratio, + ) + + head_px = headroom_ratio * ha + foot_px = footroom_ratio * ha + side_px = side_margin_ratio * wa + req_x0p = req_x0 - side_px + req_x1p = req_x1 + side_px + req_y0p = req_y0 - head_px + req_y1p = req_y1 + foot_px + xa, ya, wa, ha = _cover_min_rect( + req_x0p, + req_y0p, + max(1.0, req_x1p - req_x0p), + max(1.0, req_y1p - req_y0p), + ratio, + ) + + if horiz_gravity == "left": + xa = 0.0 + elif horiz_gravity == "right": + xa = W - wa + else: + cx = x + w / 2.0 + xa = cx - wa / 2.0 + xa = max(0.0, min(xa, W - wa)) + + y_head = y + y_feet = y + h + y_top_des = y_head - head_px + y_bot_des = y_feet + foot_px + ya = y_top_des + ya = max(0.0, min(ya, H - ha)) + + def violations(Y): + top_gap = (Y) - (y_head - head_px) + bot_gap = (Y + ha) - (y_feet + foot_px) + loss_top = max(0.0, -top_gap) + loss_bot = max(0.0, -bot_gap) + return loss_top, loss_bot, (1.0 - bottom_priority) * loss_top + bottom_priority * loss_bot + + candidates = [ + max(0.0, min(y_top_des, H - ha)), + max(0.0, min(y_bot_des - ha, H - ha)), + 0.0, + float(H - ha), + ] + best = None + for Yc in candidates: + lt, lb, cost = violations(Yc) + if best is None or cost < best[0]: + best = (cost, Yc, lt, lb) + _, ya, lt, lb = best + ya = max(0.0, min(ya, H - ha)) + + X = int(round(xa)) + Y = int(round(ya)) + Wc = int(round(wa)) + Hc = int(round(ha)) + X = max(0, min(X, W - 1)) + Y = max(0, min(Y, H - 1)) + Wc = max(1, min(Wc, W - X)) + Hc = max(1, min(Hc, H - Y)) + + dbg = ( + f"match_to={match_to}, chosen={tag}({ratio:.4f}), " + f"img={W}x{H}, mq=({x},{y},{w},{h}), " + f"margins(px): head={head_px:.1f}, foot={foot_px:.1f}, side≈{side_px:.1f}, " + f"viol: top={lt:.1f}, bot={lb:.1f}; final=({X},{Y},{Wc},{Hc})" + ) + return (Wc, Hc, X, Y, tag, dbg) + + +NODE_CLASS_MAPPINGS = { + "MQBBoxMin": MQBBoxMin, + "FitAspectHeadSafe": FitAspectHeadSafe, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MQBBoxMin": "MQ BBox (MIN)", + "FitAspectHeadSafe": "Fit Aspect (Head-Safe) - Closest AR + Tight Cover", +} diff --git a/stitch_image_by_mask.py b/stitch_image_by_mask.py new file mode 100644 index 0000000..d48a8c9 --- /dev/null +++ b/stitch_image_by_mask.py @@ -0,0 +1,213 @@ +import torch +import torch.nn.functional as F + +print("[StitchByMask] loaded") + + +def _to_bhwc(img: torch.Tensor) -> torch.Tensor: + # Accept [H,W,C] or [B,H,W,C]; ensure float32 0..1 + if img.dim() == 3: + img = img.unsqueeze(0) + if img.dim() != 4 or img.size(-1) not in (1, 3, 4): + raise ValueError(f"IMAGE must be [B,H,W,C] (C=1/3/4), got {tuple(img.shape)}") + return img.to(torch.float32).clamp(0.0, 1.0) + + +def _resize_bhwc(x: torch.Tensor, h: int, w: int, mode: str) -> torch.Tensor: + t = x.permute(0, 3, 1, 2) + t = F.interpolate( + t, + size=(h, w), + mode=mode, + align_corners=False if mode in ("bilinear", "bicubic") else None, + ) + return t.permute(0, 2, 3, 1) + + +def _mask_to_bhw1(m: torch.Tensor) -> torch.Tensor: + """ + MASK arrives as [B,H,W] or [B,H,W,1] or [H,W] or [H,W,1] (float 0..1). + Return [B,H,W,1] float32 0..1. + """ + if m.dim() == 2: # [H,W] + m = m.unsqueeze(0).unsqueeze(-1) + elif m.dim() == 3: + # Could be [B,H,W] or [H,W,1] + if m.shape[-1] == 1: # [H,W,1] + m = m.unsqueeze(0) + else: # [B,H,W] + m = m.unsqueeze(-1) + elif m.dim() == 4: + if m.shape[-1] != 1: + raise ValueError("MASK 4D must be [B,H,W,1]") + else: + raise ValueError(f"Unexpected MASK shape {tuple(m.shape)}") + return m.to(torch.float32).clamp(0.0, 1.0) + + +def _make_disk_kernel(radius: int, device) -> torch.Tensor: + if radius <= 0: + k = torch.ones((1, 1, 1, 1), device=device) + return k + d = 2 * radius + 1 + yy, xx = torch.meshgrid( + torch.arange(d, device=device), + torch.arange(d, device=device), + indexing="ij", + ) + cy = cx = radius + disk = (((yy - cy) ** 2 + (xx - cx) ** 2) <= radius * radius).float() + if disk.sum() <= 0: + disk = torch.ones_like(disk) + disk = disk / disk.sum() + return disk.view(1, 1, d, d) + + +def _gaussian_blur_chw(x: torch.Tensor, sigma: float) -> torch.Tensor: + if sigma <= 0: + return x + # Separable approx + rad = max(1, int(torch.ceil(torch.tensor(3.0 * max(0.1, sigma))).item())) + xs = torch.arange(-rad, rad + 1, device=x.device, dtype=torch.float32) + k = torch.exp(-(xs ** 2) / (2 * sigma * sigma)) + k = k / k.sum().clamp_min(1e-6) + kx = k.view(1, 1, 1, -1) + ky = k.view(1, 1, -1, 1) + C = x.shape[1] + y = F.conv2d(x, kx.expand(C, 1, 1, kx.shape[-1]), padding=(0, rad), groups=C) + y = F.conv2d(y, ky.expand(C, 1, ky.shape[-2], 1), padding=(rad, 0), groups=C) + return y + + +class StitchByMask: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_a": ("IMAGE",), + "image_b": ("IMAGE",), + "mask": ("MASK", {"default": None}), # optional when bypassing + "invert_mask": ("BOOLEAN", {"default": False}), + "bypass_mask": ("BOOLEAN", {"default": False}), + "opacity": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + "feather_radius": ( + "INT", + {"default": 5, "min": 0, "max": 100, "step": 1}, + ), + "force_size": ("BOOLEAN", {"default": False}), + "target_width": ( + "INT", + {"default": 1344, "min": 16, "max": 8192, "step": 1}, + ), + "target_height": ( + "INT", + {"default": 768, "min": 16, "max": 8192, "step": 1}, + ), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("stitched", "processed_mask") + FUNCTION = "blend" + CATEGORY = "Image/Composite" + + def blend( + self, + image_a: torch.Tensor, + image_b: torch.Tensor, + mask: torch.Tensor = None, + invert_mask: bool = False, + bypass_mask: bool = False, + opacity: float = 1.0, + feather_radius: int = 5, + force_size: bool = True, + target_width: int = 1440, + target_height: int = 1080, + ): + + a = _to_bhwc(image_a) + b = _to_bhwc(image_b) + if mask is not None: + m = _mask_to_bhw1(mask) # [B,H,W,1] + else: + m = None + + # Align sizes (do before feather so radius is in target pixels when forcing size) + if force_size: + Ht, Wt = target_height, target_width + a = _resize_bhwc(a, Ht, Wt, mode="bilinear") + b = _resize_bhwc(b, Ht, Wt, mode="bilinear") + if m is not None: + # nearest for mask to keep edges before feathering + m = _resize_bhwc(m, Ht, Wt, mode="nearest") + else: + if not (a.shape[1:3] == b.shape[1:3]): + raise ValueError( + f"Size mismatch: A {a.shape[1:3]} B {b.shape[1:3]}" + ) + if m is not None and m.shape[1:3] != a.shape[1:3]: + raise ValueError( + f"Mask size mismatch: expected {a.shape[1:3]}, got {m.shape[1:3]}" + ) + + if bypass_mask: + mask_scalar = torch.full( + (a.shape[0], a.shape[1], a.shape[2], 1), + float(opacity), + device=a.device, + dtype=a.dtype, + ) + blend_mask = ( + mask_scalar + if a.size(-1) == 1 + else mask_scalar.expand(-1, -1, -1, a.size(-1)) + ) + out = blend_mask * b + (1.0 - blend_mask) * a + return (out.clamp(0.0, 1.0), mask_scalar.clamp(0.0, 1.0)) + + if m is None: + raise ValueError("Mask input required unless 'bypass_mask' is enabled.") + + # Invert & apply opacity + if invert_mask: + m = 1.0 - m + if opacity < 1.0: + m = m * float(opacity) + + # ----- Edge-safe feathering ----- + # Make a binary core from original mask (protect original coverage) + core = (m >= 0.5).float() # [B,H,W,1] + # Dilate (expand) by feather_radius using a disk kernel + if feather_radius > 0: + k = _make_disk_kernel(feather_radius, device=m.device) # [1,1,d,d] + core_chw = core.permute(0, 3, 1, 2) # [B,1,H,W] + # convolution-based dilation approx: if any pixel inside disk -> >0 after conv + dil = F.conv2d(core_chw, k, padding=feather_radius) # normalized sum + expanded = (dil > 0).float() # binary expanded + # Feather by Gaussian blur with sigma ~ radius/2 + sigma = max(0.5, feather_radius / 2.0) + feathered = _gaussian_blur_chw(expanded, sigma=sigma).clamp(0.0, 1.0) + m_feather = feathered.permute(0, 2, 3, 1) + # Guarantee we never unmask what was masked: take max with original soft mask + m_safe = torch.maximum(m_feather, m) + else: + m_safe = m + + # Broadcast mask to match image channels + if a.size(-1) != 1 and m_safe.size(-1) == 1: + m_safe = m_safe.expand(-1, -1, -1, a.size(-1)) + + out = m_safe * b + (1.0 - m_safe) * a + return (out.clamp(0.0, 1.0), m_safe.clamp(0.0, 1.0)) + + +NODE_CLASS_MAPPINGS = { + "StitchByMask": StitchByMask, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "StitchByMask": "Stitch Two Images by Mask", +} diff --git a/workflow_config.py b/workflow_config.py new file mode 100644 index 0000000..402705a --- /dev/null +++ b/workflow_config.py @@ -0,0 +1,329 @@ +import json +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Iterable, Optional + +import folder_paths +import random + + +@dataclass(frozen=True) +class _Preset: + name: str + values: Dict[str, Any] + + +_BUILT_IN_PRESETS: Iterable[_Preset] = ( + _Preset( + "Default", + { + "initial_upscale": 1.0, + "final_upscale": 4.0, + "tight_crop": False, + "stitch_bypass_mask": False, + "stitch_opacity": 1.0, + "crop_fuzz_percent": 5.0, + "crop_pad_px": 0, + "seed": 0, + "job_notes": "", + }, + ), +) + + +def _preset_directory() -> Path: + input_root = Path(folder_paths.get_input_directory()).resolve() + target = input_root.parent / "user" / "default" / "workflows" / "workflow_config_presets" + target.mkdir(parents=True, exist_ok=True) + return target + + +def _built_in_options() -> Dict[str, Dict[str, Any]]: + return {preset.name: dict(preset.values) for preset in _BUILT_IN_PRESETS} + + +def _list_saved_presets() -> Iterable[str]: + for path in _preset_directory().glob("*.json"): + if path.is_file(): + yield path.stem + + +def _load_saved_payload(file_path: Optional[Path]) -> Dict[str, Any]: + if not file_path: + return {} + try: + with file_path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + except Exception: + return {} + + if isinstance(data, dict) and "blend_factor" in data and "stitch_opacity" not in data: + data["stitch_opacity"] = data["blend_factor"] + return data if isinstance(data, dict) else {} + + +def _persist_snapshot(target_name: Optional[str], payload: Dict[str, Any]) -> None: + filename = (target_name or "").strip() + if not filename: + return + if not filename.lower().endswith(".json"): + filename = f"{filename}.json" + destination = _preset_directory() / filename + try: + with destination.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, ensure_ascii=False) + except Exception: + pass + + +def _log(prefix: str, message: str) -> None: + print(f"[{prefix}] {message}") + + +SEED_MIN = 0 +SEED_MAX = 0xFFFFFFFF + + +def _clamp_seed(value: Any) -> int: + try: + candidate = int(value) + except Exception: + return SEED_MIN + return max(SEED_MIN, min(SEED_MAX, candidate)) + +_initial_state = random.getstate() +random.seed(datetime.now().timestamp()) +_seed_rng_state = random.getstate() +random.setstate(_initial_state) + +LOG_PREFIX = "WorkflowConfig" + + +def _new_random_seed() -> int: + global _seed_rng_state # pylint: disable=global-statement + prev_state = random.getstate() + random.setstate(_seed_rng_state) + seed_value = random.randint(SEED_MIN, SEED_MAX) + _seed_rng_state = random.getstate() + random.setstate(prev_state) + return seed_value + + +def _update_workflow_widgets( + extra_pnginfo: Dict[str, Any], + node_id: Any, + original_seed: Any, + new_seed: int, +) -> None: + workflow = extra_pnginfo.get("workflow", {}) if isinstance(extra_pnginfo, dict) else {} + nodes = workflow.get("nodes", []) if isinstance(workflow.get("nodes", []), list) else [] + workflow_node = next((node for node in nodes if str(node.get("id")) == str(node_id)), None) + if not workflow_node or "widgets_values" not in workflow_node: + _log(LOG_PREFIX, "Unable to store seed in workflow metadata (node not found).") + else: + for index, value in enumerate(workflow_node["widgets_values"]): + if value == original_seed: + workflow_node["widgets_values"][index] = new_seed + + +def _update_prompt_inputs(prompt_nodes: Dict[str, Any], node_id: Any, new_seed: int) -> None: + prompt_node = prompt_nodes.get(str(node_id)) if isinstance(prompt_nodes, dict) else None + if not prompt_node or "inputs" not in prompt_node or "seed" not in prompt_node["inputs"]: + _log(LOG_PREFIX, "Unable to store seed in prompt metadata (node not found).") + else: + prompt_node["inputs"]["seed"] = new_seed + + +def _resolve_seed( + seed: int, + prompt: Optional[Dict[str, Any]], + extra_pnginfo: Optional[Dict[str, Any]], + unique_id: Optional[Any], +) -> int: + if seed in (-1, -2, -3): + _log(LOG_PREFIX, f'Received special seed "{seed}". Generating a new random seed server-side.') + if seed in (-2, -3): + action = "increment" if seed == -2 else "decrement" + _log(LOG_PREFIX, f"Cannot {action} without prior seed; using random seed.") + + original_seed = seed + seed = _new_random_seed() + _log(LOG_PREFIX, f"Generated random seed {seed} and will attempt to persist it.") + + if unique_id is None: + _log(LOG_PREFIX, "Node unique_id was not provided; cannot persist seed.") + else: + if extra_pnginfo is None: + _log(LOG_PREFIX, "Workflow metadata not provided; cannot persist seed.") + else: + _update_workflow_widgets(extra_pnginfo, unique_id, original_seed, seed) + + if prompt is None: + _log(LOG_PREFIX, "Prompt metadata not provided; cannot persist seed.") + else: + _update_prompt_inputs(prompt, unique_id, seed) + + return seed + + +class WorkflowConfig: + """ + Presents workflow-level controls (upscale amounts, crop policy, stitch options, seed control) + and emits both socket values and a serialized JSON bundle for downstream nodes. + """ + + @classmethod + def INPUT_TYPES(cls): + built_in = sorted(_built_in_options().keys()) + saved_files = [""] + sorted(set(_list_saved_presets())) + return { + "required": { + "template_name": (built_in, {"default": built_in[0]}), + "saved_preset": (saved_files, {"default": saved_files[0]}), + "load_saved_preset": ("BOOLEAN", {"default": False}), + "initial_upscale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1}), + "final_upscale": ("FLOAT", {"default": 4.0, "min": 0.1, "max": 100.0, "step": 0.1}), + "tight_crop": ("BOOLEAN", {"default": False}), + "stitch_bypass_mask": ("BOOLEAN", {"default": False}), + "stitch_opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "crop_fuzz_percent": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 50.0, "step": 0.1}), + "crop_pad_px": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}), + "seed": ("INT", {"default": 0, "min": SEED_MIN, "max": SEED_MAX}), + "job_notes": ("STRING", {"default": "", "multiline": True, "placeholder": "Notes"}), + }, + "optional": { + "save_preset_as": ("STRING", {"default": "", "multiline": False, "placeholder": "preset name (e.g., portrait)"}), + "write_preset": ("BOOLEAN", {"default": False}), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ( + "FLOAT", + "FLOAT", + "BOOLEAN", + "BOOLEAN", + "FLOAT", + "FLOAT", + "INT", + "INT", + "STRING", + "STRING", + ) + RETURN_NAMES = ( + "initial_upscale", + "final_upscale", + "tight_crop", + "stitch_bypass_mask", + "stitch_opacity", + "crop_fuzz_percent", + "crop_pad_px", + "seed", + "job_notes", + "config_bundle", + ) + FUNCTION = "configure" + CATEGORY = "config" + + @classmethod + def IS_CHANGED( + cls, + template_name, + saved_preset, + load_saved_preset, + initial_upscale, + final_upscale, + tight_crop, + stitch_bypass_mask, + stitch_opacity, + crop_fuzz_percent, + crop_pad_px, + seed, + job_notes, + save_preset_as=None, + write_preset=False, + prompt=None, + extra_pnginfo=None, + unique_id=None, + ): + if seed in (-1, -2, -3): + return _new_random_seed() + return seed + + def configure( + self, + template_name, + saved_preset, + load_saved_preset, + initial_upscale, + final_upscale, + tight_crop, + stitch_bypass_mask, + stitch_opacity, + crop_fuzz_percent, + crop_pad_px, + seed=0, + job_notes="", + save_preset_as=None, + write_preset=False, + prompt=None, + extra_pnginfo=None, + unique_id=None, + ): + templates = _built_in_options() + template_payload = dict(templates.get(template_name, templates["Default"])) + + preset_payload: Dict[str, Any] = {} + if load_saved_preset and saved_preset and saved_preset != "": + preset_path = _preset_directory() / f"{saved_preset}.json" + preset_payload = _load_saved_payload(preset_path) + + merged = {**template_payload, **preset_payload} + + values: Dict[str, Any] = { + "initial_upscale": float(initial_upscale if initial_upscale is not None else merged["initial_upscale"]), + "final_upscale": float(final_upscale if final_upscale is not None else merged["final_upscale"]), + "tight_crop": bool(tight_crop if tight_crop is not None else merged["tight_crop"]), + "stitch_bypass_mask": bool(stitch_bypass_mask if stitch_bypass_mask is not None else merged["stitch_bypass_mask"]), + "stitch_opacity": float(stitch_opacity if stitch_opacity is not None else merged["stitch_opacity"]), + "crop_fuzz_percent": float(crop_fuzz_percent if crop_fuzz_percent is not None else merged["crop_fuzz_percent"]), + "crop_pad_px": int(crop_pad_px if crop_pad_px is not None else merged["crop_pad_px"]), + "seed": _clamp_seed(seed if seed is not None else merged.get("seed", 0)), + "job_notes": str(job_notes if job_notes is not None else merged.get("job_notes", "")), + } + + seed_value = _clamp_seed(_resolve_seed(values["seed"], prompt, extra_pnginfo, unique_id)) + values["seed"] = seed_value + + if write_preset and save_preset_as: + to_store = dict(values) + _persist_snapshot(save_preset_as, to_store) + + bundle = json.dumps(values, ensure_ascii=False) + + return ( + values["initial_upscale"], + values["final_upscale"], + values["tight_crop"], + values["stitch_bypass_mask"], + values["stitch_opacity"], + values["crop_fuzz_percent"], + values["crop_pad_px"], + seed_value, + values["job_notes"], + bundle, + ) + + +NODE_CLASS_MAPPINGS = { + "WorkflowConfig": WorkflowConfig, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WorkflowConfig": "Workflow Config (central)", +}