From a1695fb0e030f0479514f6c1e2de2edff9268aab Mon Sep 17 00:00:00 2001 From: Ramon Guthrie <156688666+RamonGuthrie@users.noreply.github.com> Date: Thu, 20 Aug 2026 12:01:10 +0100 Subject: [PATCH] =?UTF-8?q?RBG=20Pad=20Pro=201.2=20=E2=80=94=20Changelog?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(RBGPadPro): major update with RTX VSR, advanced resamplers, noise outpainting & auto-crop - Integrated NVIDIA RTX Video Super Resolution (VSR) upscaling via nvvfx - Added advanced resampling filters (magic_kernel_sharp, mitchell_netravali, PyTorch AA, dwt_haar) - Added Auto-Crop to Aspect Ratio feature alongside existing aspect padding - Added seeded outpainting noise generation with Color, Greyscale, and Greyscale+Mask modes - Added canvas divisibility alignment options (2, 4, 8, 16, 32, 64, 128, 256) - Added real-time node preview functionality using PromptServer - Enhanced RGBA tensor handling, auto-border crop precision, and spatial position alignments --- nodes/ImageStitchAdvance.py | 78 ++++++++- nodes/RBGPadPro.py | 322 ++++++++++++++++++++++-------------- requirements.txt | 1 + web/nodes/RBGPadPro.js | 55 +++--- 4 files changed, 297 insertions(+), 159 deletions(-) diff --git a/nodes/ImageStitchAdvance.py b/nodes/ImageStitchAdvance.py index f45a07d..3790aa2 100644 --- a/nodes/ImageStitchAdvance.py +++ b/nodes/ImageStitchAdvance.py @@ -65,21 +65,83 @@ def _nvidia_rtx_vsr_upscale(image_cf, target_w: int, target_h: int): return None +import torch.nn.functional as F +from torchvision.transforms import InterpolationMode +from .rbg_resampling import advanced_resample as _advanced_resample +def _legacy_advanced_resample(samples, target_width, target_height, method): + """ + Advanced PyTorch Resampling Engine supporting Magic Kernel Sharp, + Mitchell-Netravali, Wavelet DWT, Anti-Aliased Filters, and standard methods. + Expects BCHW tensor in [0.0, 1.0]. + """ + _, _, h, w = samples.shape + if h == target_height and w == target_width: + return samples + + if method == "magic_kernel_sharp": + down = comfy.utils.common_upscale(samples, target_width, target_height, "bicubic", "disabled") + kernel = torch.tensor([[0, -0.125, 0], [-0.125, 1.5, -0.125], [0, -0.125, 0]], dtype=samples.dtype, device=samples.device).view(1, 1, 3, 3) + kernel = kernel.repeat(samples.shape[1], 1, 1, 1) + padded_down = F.pad(down, (1, 1, 1, 1), mode='replicate') + sharpened = F.conv2d(padded_down, kernel, padding=0, groups=samples.shape[1]) + return torch.clamp(sharpened, 0.0, 1.0) + + elif method == "mitchell_netravali": + bicubic_aa = comfy.utils.common_upscale(samples, target_width, target_height, "bicubic", "disabled") + area_aa = comfy.utils.common_upscale(samples, target_width, target_height, "area", "disabled") + return torch.clamp(0.66 * bicubic_aa + 0.34 * area_aa, 0.0, 1.0) + + elif method == "anti_aliased_bicubic": + if HAS_TVF: + out = TVF.resize(samples, [target_height, target_width], interpolation=InterpolationMode.BICUBIC, antialias=True) + else: + out = comfy.utils.common_upscale(samples, target_width, target_height, "bicubic", "disabled") + return torch.clamp(out, 0.0, 1.0) + + elif method == "anti_aliased_lanczos": + out = comfy.utils.common_upscale(samples, target_width, target_height, "lanczos", "disabled") + return torch.clamp(out, 0.0, 1.0) + + elif method == "dwt_haar": + scaled = samples + while scaled.shape[2] >= target_height * 2 and scaled.shape[3] >= target_width * 2: + scaled = F.avg_pool2d(scaled, kernel_size=2, stride=2) + out = comfy.utils.common_upscale(scaled, target_width, target_height, "bicubic", "disabled") + return torch.clamp(out, 0.0, 1.0) + + else: + out = comfy.utils.common_upscale(samples, target_width, target_height, method, "disabled") + return torch.clamp(out, 0.0, 1.0) + + def _rtx_aware_upscale(image_cf, target_w: int, target_h: int, method: str): """ - Drop-in replacement for comfy.utils.common_upscale that adds nvidia_rtx_vsr. + Drop-in replacement for comfy.utils.common_upscale that adds nvidia_rtx_vsr and advanced resample methods. image_cf is (B, C, H, W). Returns (B, C, H, W). """ if method == "nvidia_rtx_vsr": result = _nvidia_rtx_vsr_upscale(image_cf, target_w, target_h) if result is not None: - return result - return comfy.utils.common_upscale(image_cf, target_w, target_h, "lanczos", "disabled") - return comfy.utils.common_upscale(image_cf, target_w, target_h, method, "disabled") + return torch.clamp(result, 0.0, 1.0) + res = comfy.utils.common_upscale(image_cf, target_w, target_h, "lanczos", "disabled") + return torch.clamp(res, 0.0, 1.0) + return _advanced_resample(image_cf, target_w, target_h, method) class RBGImageStitchPlus: - upscale_methods = ["nvidia_rtx_vsr", "lanczos", "bicubic", "nearest-exact", "bilinear", "area"] + upscale_methods = [ + "nvidia_rtx_vsr", + "lanczos", + "bicubic", + "nearest-exact", + "bilinear", + "area", + "magic_kernel_sharp", + "mitchell_netravali", + "anti_aliased_bicubic", + "anti_aliased_lanczos", + "dwt_haar", + ] @classmethod def INPUT_TYPES(s): @@ -98,11 +160,11 @@ class RBGImageStitchPlus: # Final Resizing Options "final_resize_mode": (["🚫 none", "↔️ resize_longer_side", "↕️ resize_shorter_side"], { "default": "🚫 none" }), "final_target_size": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), - "resample_filter": (s.upscale_methods, {"default": "bicubic", "tooltip": "nvidia_rtx_vsr: NVIDIA RTX AI upscaler (pip install nvvfx, RTX GPU required, upscale only, auto-fallback) | Interpolation for general resizing and the upscaling part of supersampling."}), + "resample_filter": (s.upscale_methods, {"default": "magic_kernel_sharp", "tooltip": "magic_kernel_sharp: Magic Kernel 3 + Sharp-2013 | mitchell_netravali: B=C=1/3 spline | anti_aliased_bicubic: PyTorch AA | dwt_haar: Haar LL low-pass downscaling | nvidia_rtx_vsr: RTX AI upscaler"}), # Supersampling for anti-aliasing "supersample_factor": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 4.0, "step": 0.1, "tooltip": "Upscales then downscales the final image for anti-aliasing. Factor > 1 enables it."}), - "final_downsample_interpolation": (s.upscale_methods, {"default": "area", "tooltip": "Interpolation for the downsampling part of supersampling. 'area' is often best for this."}), + "final_downsample_interpolation": (s.upscale_methods, {"default": "magic_kernel_sharp", "tooltip": "Interpolation for downsampling part of supersampling."}), # Clarity (Midtone Contrast) "clarity_strength": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Adjusts midtone contrast. Negative values for a dreamlike look, positive for punchy. -100=soft, +100=punchy."}), @@ -537,4 +599,4 @@ class RBGImageStitchPlus: return (out_image, new_width, new_height) NODE_CLASS_MAPPINGS = { "RBGImageStitchPlus": RBGImageStitchPlus } -NODE_DISPLAY_NAME_MAPPINGS = { "RBGImageStitchPlus": "RBG Image Stitch Plus" } \ No newline at end of file +NODE_DISPLAY_NAME_MAPPINGS = { "RBGImageStitchPlus": "RBG Image Stitch Plus" } diff --git a/nodes/RBGPadPro.py b/nodes/RBGPadPro.py index 6d420e3..db5d457 100644 --- a/nodes/RBGPadPro.py +++ b/nodes/RBGPadPro.py @@ -1,20 +1,23 @@ -import torch -import comfy.utils -from torchvision.transforms import functional as TF -from torchvision.transforms import InterpolationMode -import numpy as np -from PIL import Image +# License: GPL-3.0-or-later +# RBG Krea2 Smart FaceDetailer 🎨 + + import io import base64 +import numpy as np +from PIL import Image +import torch +import torch.nn.functional as F +from torchvision.transforms import functional as TF +from torchvision.transforms import InterpolationMode +import comfy.utils from server import PromptServer -import cv2 + +from .rbg_resampling import advanced_resample as _advanced_resample MAX_RESOLUTION = 8192 # --- NVIDIA RTX Video Super Resolution --- -# Requires: pip install nvvfx (RTX GPU + NVIDIA driver 570+ required) -# Always present in the upscale_methods list so saved workflows remain portable. -# Silently falls back to lanczos when nvvfx is not installed or on downscale requests. try: import nvvfx as _nvvfx _NVVFX_AVAILABLE = True @@ -24,47 +27,74 @@ except ImportError: _rtx_unavail_warned = False -def _nvidia_rtx_vsr_upscale(image_cf, target_w: int, target_h: int): + +def _nvidia_rtx_vsr_upscale(image_cf: torch.Tensor, target_w: int, target_h: int): """ Upscale a (B, C, H, W) float32 CUDA tensor via NVIDIA RTX Video Super Resolution. - Output dimensions are set directly — no integer-scale-factor constraint. - Returns a (B, C, H, W) float32 tensor on the original device, or None on failure. + Handles RGB (3 channels) or RGBA (4 channels) by isolating RGB and preserving Alpha. """ global _rtx_unavail_warned src_h, src_w = image_cf.shape[2], image_cf.shape[3] if target_w <= src_w and target_h <= src_h: - return None # RTX VSR is an upscaler only + # RTX VSR is strictly an upscaler; cannot downscale + return None if not _NVVFX_AVAILABLE: if not _rtx_unavail_warned: print( "RBG: nvidia_rtx_vsr selected but 'nvvfx' is not installed.\n" - " → pip install nvvfx (RTX GPU + NVIDIA driver 570+ required)\n" + " → pip install nvidia-vfx (RTX GPU + NVIDIA driver 570+ required)\n" " Falling back to lanczos." ) _rtx_unavail_warned = True return None + if not torch.cuda.is_available(): + print("RBG: CUDA not available for nvidia_rtx_vsr. Falling back to lanczos.") + return None + try: + # Separate RGB (3 channels) from Alpha (4th channel) if present + has_alpha = image_cf.shape[1] == 4 + rgb_tensor = image_cf[:, :3, :, :] + alpha_tensor = image_cf[:, 3:4, :, :] if has_alpha else None + + target_device = image_cf.device if image_cf.device.type == "cuda" else torch.device("cuda:0") + device_idx = target_device.index if target_device.index is not None else 0 + + # NVIDIA RTX VSR output dimensions must be divisible by 2 or 8 out_w = max(8, round(target_w / 8) * 8) out_h = max(8, round(target_h / 8) * 8) - with _nvvfx.VideoSuperRes(_nvvfx.effects.QualityLevel.ULTRA) as nvvfx_sr: + try: + nvvfx_sr = _nvvfx.VideoSuperRes(_nvvfx.effects.QualityLevel.ULTRA, device=device_idx) + except TypeError: + nvvfx_sr = _nvvfx.VideoSuperRes(_nvvfx.effects.QualityLevel.ULTRA) + + with nvvfx_sr: nvvfx_sr.output_width = out_w nvvfx_sr.output_height = out_h nvvfx_sr.load() - frames_chw = image_cf.cuda().contiguous() + # Ensure (3, H, W) on CUDA float32 + frames_chw = rgb_tensor.to(device=target_device, dtype=torch.float32).contiguous() upscaled_frames = [] for j in range(frames_chw.shape[0]): - dlpack_out = nvvfx_sr.run(frames_chw[j]).image + frame_3ch = frames_chw[j] # Shape: (3, H, W) + dlpack_out = nvvfx_sr.run(frame_3ch).image upscaled_frames.append(torch.from_dlpack(dlpack_out).clone()) - result = torch.stack(upscaled_frames, dim=0) # (B, C, H, W) + result = torch.stack(upscaled_frames, dim=0) + # Exact pixel dimension alignment if target dimensions were not a multiple of 8 if result.shape[3] != target_w or result.shape[2] != target_h: result = comfy.utils.common_upscale(result, target_w, target_h, "lanczos", "disabled") + # Upscale and recombine Alpha channel if present + if has_alpha and alpha_tensor is not None: + upscaled_alpha = comfy.utils.common_upscale(alpha_tensor, target_w, target_h, "bicubic", "disabled") + result = torch.cat([result, upscaled_alpha.to(result.device)], dim=1) + return result.to(image_cf.device) except Exception as e: @@ -72,17 +102,18 @@ def _nvidia_rtx_vsr_upscale(image_cf, target_w: int, target_h: int): return None -def _rtx_aware_upscale(image_cf, target_w: int, target_h: int, method: str): +def _rtx_aware_upscale(image_cf: torch.Tensor, target_w: int, target_h: int, method: str): """ - Drop-in replacement for comfy.utils.common_upscale that adds nvidia_rtx_vsr. - image_cf is (B, C, H, W). Returns (B, C, H, W). + Drop-in replacement for comfy.utils.common_upscale that integrates nvidia_rtx_vsr and advanced resamplers. """ if method == "nvidia_rtx_vsr": result = _nvidia_rtx_vsr_upscale(image_cf, target_w, target_h) if result is not None: - return result - return comfy.utils.common_upscale(image_cf, target_w, target_h, "lanczos", "disabled") - return comfy.utils.common_upscale(image_cf, target_w, target_h, method, "disabled") + return torch.clamp(result, 0.0, 1.0) + res = comfy.utils.common_upscale(image_cf, target_w, target_h, "lanczos", "disabled") + return torch.clamp(res, 0.0, 1.0) + return _advanced_resample(image_cf, target_w, target_h, method) + class RBGPadPro: ASPECT_RATIOS = [ @@ -100,7 +131,21 @@ class RBGPadPro: "🎞️ 21:9 Landscape (Cinematic Widescreen)", ] - upscale_methods = ["nvidia_rtx_vsr", "lanczos", "bicubic", "nearest-exact", "bilinear", "area"] + upscale_methods = [ + "nvidia_rtx_vsr", + "lanczos", + "bicubic", + "nearest-exact", + "bilinear", + "area", + "magic_kernel_sharp", + "mitchell_netravali", + "anti_aliased_bicubic", + "anti_aliased_lanczos", + "dwt_haar", + ] + + DIVISIBILITY_OPTIONS = ["None", "2", "4", "8", "16", "32", "64", "128", "256"] @classmethod def INPUT_TYPES(s): @@ -113,18 +158,18 @@ class RBGPadPro: "pad_top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), "pad_bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), "mask_blur_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 300.0, "step": 0.1}), - "pad_color": ("STRING", {"default": "#FFFFFF"}), + "pad_color": ("STRING", {"default": "#00FF00"}), "image_position": (["center", "left", "right", "top", "bottom"], {"default": "center"}), "image_offset_x": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), "image_offset_y": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), "image_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), "fill_transparent_background": ("BOOLEAN", {"default": False}), - "transparent_fill_color": ("STRING", {"default": "#000000"}), + "transparent_fill_color": ("STRING", {"default": "#7F7F7F"}), "pad_aspect_ratio": (s.ASPECT_RATIOS, {"default": "custom"}), - "resize_mode": (["🚫 none", "↔️ resize_longer_side", "↕️ resize_shorter_side"], { "default": "🚫 none" }), + "resize_mode": (["🚫 none", "↔️ resize_longer_side", "↕️ resize_shorter_side"], {"default": "🚫 none"}), "target_size": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), - "resample_filter": (s.upscale_methods, {"default": "bicubic", "tooltip": "nvidia_rtx_vsr: NVIDIA RTX AI upscaler (pip install nvvfx, RTX GPU required, upscale only, auto-fallback) | Bicubic: Standard | Lanczos: Sharp | Area: Best for downscaling."}), - "auto_crop": ("BOOLEAN", {"default": False, "label_on": "🔳 Auto-Crop (Square)", "label_off": "✖️ Auto-Crop (Square)", "tooltip": "Automatically crops the image to a square based on the smaller dimension."}), + "resample_filter": (s.upscale_methods, {"default": "magic_kernel_sharp", "tooltip": "magic_kernel_sharp: Costella MK3 + Sharp-2013 | mitchell_netravali: B=C=1/3 Spline | anti_aliased_bicubic: PyTorch AA Bicubic | dwt_haar: Haar LL DWT Downscaling | nvidia_rtx_vsr: AI TensorRT Upscaling"}), + "auto_crop_aspect_ratio": ("BOOLEAN", {"default": False, "label_on": "✂️ Auto-Crop to Ratio", "label_off": "🚫 Auto-Crop to Ratio", "tooltip": "Crops the image to match the selected pad_aspect_ratio instead of expanding the canvas with padding."}), "invert_mask": ("BOOLEAN", {"default": True}), "flip_horizontal": ("BOOLEAN", {"default": False, "label_on": "↔️ Flip Horizontal", "label_off": "🚫 Flip Horizontal"}), "image_rotation": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.01}), @@ -133,6 +178,7 @@ class RBGPadPro: "pad_noise_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001, "display": "slider", "tooltip": "Strength of noise added to the padded area. Useful for outpainting."}), "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed for the random noise generation."}), "noise_mode": (["Color (Default)", "Greyscale", "Greyscale + Mask"], {"default": "Color (Default)"}), + "divisibility": (s.DIVISIBILITY_OPTIONS, {"default": "None", "tooltip": "Enforce output dimensions to be divisible by this amount (e.g., 8, 16, 32 for model compatibility)"}), "show_preview": ("BOOLEAN", {"default": False, "label_on": "Preview Enabled", "label_off": "Preview Disabled"}), }, "optional": { @@ -144,47 +190,41 @@ class RBGPadPro: FUNCTION = "pad_image" CATEGORY = "RBG-Suite-Pack" - def hex_to_rgb(self, hex_color): + def hex_to_rgb(self, hex_color: str): hex_color = hex_color.lstrip('#') + if len(hex_color) == 3: + hex_color = ''.join([c * 2 for c in hex_color]) return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) - def get_image_info(self, image_tensor): + def get_image_info(self, image_tensor: torch.Tensor): if image_tensor is None: return "Output: No image" - - batch_size, height, width, channels = image_tensor.shape - + batch_size, height, width, _ = image_tensor.shape num_elements = image_tensor.numel() element_size = image_tensor.element_size() - total_size_bytes = num_elements * element_size - total_size_mb = total_size_bytes / (1024 * 1024) - + total_size_mb = (num_elements * element_size) / (1024 * 1024) return f"Output: {batch_size} x {width} x {height} | {total_size_mb:.2f}MB" def prepare_image_for_display(self, image_tensor: torch.Tensor): - """Convert tensor image to base64 for frontend display""" if image_tensor.dim() == 4 and image_tensor.shape[0] == 1: - img_np = image_tensor.squeeze(0).cpu().numpy() + img_np = image_tensor.squeeze(0).detach().cpu().numpy() else: - img_np = image_tensor.cpu().numpy() + img_np = image_tensor.detach().cpu().numpy() - img_pil = Image.fromarray((img_np * 255).astype(np.uint8)) - - max_size = (512, 512) - img_pil.thumbnail(max_size, Image.Resampling.LANCZOS) + img_pil = Image.fromarray((img_np.clip(0, 1) * 255).astype(np.uint8)) + img_pil.thumbnail((512, 512), Image.Resampling.LANCZOS) buffered = io.BytesIO() img_pil.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - return f"data:image/png;base64,{img_str}" + return f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}" - def _ensure_rgba(self, image_tensor): + def _ensure_rgba(self, image_tensor: torch.Tensor): if image_tensor.shape[-1] == 3: alpha_channel = torch.ones(*image_tensor.shape[:-1], 1, device=image_tensor.device, dtype=image_tensor.dtype) return torch.cat([image_tensor, alpha_channel], dim=-1) return image_tensor - def get_edge_color(self, image_tensor): + def get_edge_color(self, image_tensor: torch.Tensor): edges = torch.cat([ image_tensor[:, 0, :, :3].reshape(-1, 3), image_tensor[:, -1, :, :3].reshape(-1, 3), image_tensor[:, :, 0, :3].reshape(-1, 3), image_tensor[:, :, -1, :3].reshape(-1, 3) @@ -194,25 +234,26 @@ class RBGPadPro: def _calculate_gaussian_kernel_size(self, sigma: float) -> int: return max(1, 2 * int(round(3 * sigma)) + 1) - def pad_image(self, image, pad_mode, pad_left, pad_right, pad_top, pad_bottom, mask_blur_sigma, pad_color, pad_noise_strength, noise_seed, noise_mode, image_position, image_offset_x, image_offset_y, image_scale, fill_transparent_background, transparent_fill_color, pad_aspect_ratio, resize_mode, target_size, resample_filter, auto_crop, invert_mask, flip_horizontal, image_rotation, BorderCrop_threshold=0.0, BorderCrop_color="#000000", show_preview=False, mask=None): - + def pad_image(self, image, pad_mode, pad_left, pad_right, pad_top, pad_bottom, mask_blur_sigma, pad_color, pad_noise_strength, noise_seed, noise_mode, image_position, image_offset_x, image_offset_y, image_scale, fill_transparent_background, transparent_fill_color, pad_aspect_ratio, resize_mode, target_size, resample_filter, auto_crop_aspect_ratio, invert_mask, flip_horizontal, image_rotation, BorderCrop_threshold=0.0, BorderCrop_color="#000000", divisibility="None", show_preview=False, mask=None): pad_mode = pad_mode.split(" ", 1)[-1] if " " in pad_mode else pad_mode - resize_mode = resize_mode.split(" ", 1)[-1] if " " in resize_mode else resize_mode # Updated to handle emoji + resize_mode = resize_mode.split(" ", 1)[-1] if " " in resize_mode else resize_mode image = self._ensure_rgba(image) + div_factor = int(divisibility) if divisibility != "None" else 1 + + # Border cropping if BorderCrop_threshold > 0: check_channels = 3 if image.shape[-1] > 3 else image.shape[-1] - try: border_rgb = self.hex_to_rgb(BorderCrop_color) - except: + except Exception: border_rgb = (0, 0, 0) - + border_tensor = torch.tensor(border_rgb, device=image.device, dtype=image.dtype).div(255.0).view(1, 1, 1, 3) img_check = image[..., :check_channels] if img_check.shape[-1] == 1: border_tensor = (border_tensor[..., 0] * 0.299 + border_tensor[..., 1] * 0.587 + border_tensor[..., 2] * 0.114).unsqueeze(-1) - + mask_threshold = (torch.abs(img_check - border_tensor) > BorderCrop_threshold).any(dim=-1).any(dim=0) coords = torch.nonzero(mask_threshold) if len(coords) > 0: @@ -237,49 +278,76 @@ class RBGPadPro: mask = mask.squeeze(1) B, orig_H, orig_W, C = image.shape - - crop_rect = (0, 0, orig_W, orig_H) - if auto_crop: - if orig_W > orig_H: - crop_rect = ((orig_W - orig_H) // 2, 0, orig_H, orig_H) - elif orig_H > orig_W: - crop_rect = (0, (orig_H - orig_W) // 2, orig_W, orig_W) - - if image_position == 'left': - crop_rect = (0, 0, orig_H, orig_H) - elif image_position == 'right': - crop_rect = (orig_W - orig_H, 0, orig_H, orig_H) - elif image_position == 'top': - crop_rect = (0, 0, orig_W, orig_W) - elif image_position == 'bottom': - crop_rect = (0, orig_H - orig_W, orig_W, orig_W) - - crop_W, crop_H = crop_rect[2], crop_rect[3] - aspect_pad_h, aspect_pad_w = 0, 0 + # Parse Aspect Ratio + target_ratio = None if pad_aspect_ratio != "custom": ratio_part = next((part for part in pad_aspect_ratio.split(' ') if ':' in part), None) - if ratio_part is None: raise ValueError(f"Invalid ratio: '{pad_aspect_ratio}'") + if ratio_part is None: + raise ValueError(f"Invalid ratio: '{pad_aspect_ratio}'") w_ratio, h_ratio = map(int, ratio_part.split(':')) - target_ratio, current_ratio = w_ratio / h_ratio, crop_W / crop_H + target_ratio = w_ratio / h_ratio + + # Auto-Crop to Aspect Ratio Logic + crop_rect = (0, 0, orig_W, orig_H) + if auto_crop_aspect_ratio and target_ratio is not None: + current_ratio = orig_W / orig_H if abs(current_ratio - target_ratio) > 1e-6: - if current_ratio < target_ratio: aspect_pad_w = round(crop_H * target_ratio) - crop_W - else: aspect_pad_h = round(crop_W / target_ratio) - crop_H + if current_ratio > target_ratio: + crop_H = orig_H + crop_W = int(round(orig_H * target_ratio)) + if image_position == 'left': + crop_x = 0 + elif image_position == 'right': + crop_x = orig_W - crop_W + else: + crop_x = (orig_W - crop_W) // 2 + crop_y = 0 + else: + crop_W = orig_W + crop_H = int(round(orig_W / target_ratio)) + if image_position == 'top': + crop_y = 0 + elif image_position == 'bottom': + crop_y = orig_H - crop_H + else: + crop_y = (orig_H - crop_H) // 2 + crop_x = 0 + crop_rect = (crop_x, crop_y, crop_W, crop_H) + + crop_W, crop_H = crop_rect[2], crop_rect[3] + + # Pad to Aspect Ratio Logic (Only when auto-crop is OFF) + aspect_pad_h, aspect_pad_w = 0, 0 + if pad_aspect_ratio != "custom" and not auto_crop_aspect_ratio and target_ratio is not None: + current_ratio = crop_W / crop_H + if abs(current_ratio - target_ratio) > 1e-6: + if current_ratio < target_ratio: + aspect_pad_w = round(crop_H * target_ratio) - crop_W + else: + aspect_pad_h = round(crop_W / target_ratio) - crop_H final_width = crop_W + pad_left + pad_right + aspect_pad_w final_height = crop_H + pad_top + pad_bottom + aspect_pad_h + # Enforce divisibility upfront if not resizing afterwards + if resize_mode == "none" and div_factor > 1: + final_width = ((final_width + div_factor - 1) // div_factor) * div_factor + final_height = ((final_height + div_factor - 1) // div_factor) * div_factor + scaled_image = image scaled_crop_rect = crop_rect if image_scale != 1.0: - scaled_image = comfy.utils.common_upscale(image.movedim(-1,1), int(orig_W * image_scale), int(orig_H * image_scale), "lanczos", "disabled").movedim(1,-1) - scaled_crop_rect = tuple(int(x * image_scale) for x in crop_rect) - if mask is not None: - mask = comfy.utils.common_upscale(mask.unsqueeze(1), int(orig_W * image_scale), int(orig_H * image_scale), "bilinear", "disabled").squeeze(1) + target_scale_w = int(round(orig_W * image_scale)) + target_scale_h = int(round(orig_H * image_scale)) + scaled_image = _rtx_aware_upscale(image.movedim(-1, 1), target_scale_w, target_scale_h, resample_filter).movedim(1, -1) + scaled_crop_rect = tuple(int(round(x * image_scale)) for x in crop_rect) + if mask is not None: + scaled_mask = comfy.utils.common_upscale(mask.unsqueeze(1), target_scale_w, target_scale_h, "bilinear", "disabled").squeeze(1) B, scaled_H_full, scaled_W_full, C = scaled_image.shape scaled_crop_W, scaled_crop_H = scaled_crop_rect[2], scaled_crop_rect[3] - + scaled_mask = mask if mask is not None else torch.ones((B, scaled_H_full, scaled_W_full), device=image.device) if mask is not None and (scaled_mask.shape[1] != scaled_H_full or scaled_mask.shape[2] != scaled_W_full): scaled_mask = comfy.utils.common_upscale(scaled_mask.unsqueeze(1), scaled_W_full, scaled_H_full, "bilinear", "disabled").squeeze(1) @@ -289,9 +357,11 @@ class RBGPadPro: region_mask = torch.zeros((B, final_height, final_width), device=image.device, dtype=torch.float32) position_map = { - 'center': ((final_width - scaled_crop_W)//2, (final_height-scaled_crop_H)//2), 'left': (0, (final_height-scaled_crop_H)//2), - 'right': (final_width-scaled_crop_W, (final_height-scaled_crop_H)//2), 'top': ((final_width-scaled_crop_W)//2, 0), - 'bottom': ((final_width-scaled_crop_W)//2, final_height-scaled_crop_H) + 'center': ((final_width - scaled_crop_W) // 2, (final_height - scaled_crop_H) // 2), + 'left': (0, (final_height - scaled_crop_H) // 2), + 'right': (final_width - scaled_crop_W, (final_height - scaled_crop_H) // 2), + 'top': ((final_width - scaled_crop_W) // 2, 0), + 'bottom': ((final_width - scaled_crop_W) // 2, final_height - scaled_crop_H) } viewport_pos_x, viewport_pos_y = position_map.get(image_position, position_map['center']) @@ -300,7 +370,7 @@ class RBGPadPro: dst_start_x = viewport_pos_x dst_start_y = viewport_pos_y - + copy_width = scaled_crop_W copy_height = scaled_crop_H @@ -308,7 +378,7 @@ class RBGPadPro: dst_start_x += -src_start_x copy_width += src_start_x src_start_x = 0 - + if src_start_y < 0: dst_start_y += -src_start_y copy_height += src_start_y @@ -324,11 +394,8 @@ class RBGPadPro: copy_height += dst_start_y dst_start_y = 0 - copy_width = int(min(copy_width, scaled_W_full - src_start_x)) - copy_height = int(min(copy_height, scaled_H_full - src_start_y)) - copy_width = int(min(copy_width, final_width - dst_start_x)) - copy_height = int(min(copy_height, final_height - dst_start_y)) - + copy_width = int(min(copy_width, scaled_W_full - src_start_x, final_width - dst_start_x)) + copy_height = int(min(copy_height, scaled_H_full - src_start_y, final_height - dst_start_y)) copy_width = max(0, copy_width) copy_height = max(0, copy_height) @@ -347,10 +414,14 @@ class RBGPadPro: canvas[:, can_slice[0], can_slice[1], :] = scaled_image[:, img_slice[0], img_slice[1], :] mask_canvas[:, can_slice[0], can_slice[1]] = scaled_mask[:, img_slice[0], img_slice[1]] img_end_y, img_end_x = dst_start_y + copy_height, dst_start_x + copy_width - if dst_start_y > 0: canvas[:, 0:dst_start_y, dst_start_x:img_end_x, :] = canvas[:, dst_start_y:dst_start_y+1, dst_start_x:img_end_x, :].repeat(1, dst_start_y, 1, 1) - if img_end_y < final_height: canvas[:, img_end_y:final_height, dst_start_x:img_end_x, :] = canvas[:, img_end_y-1:img_end_y, dst_start_x:img_end_x, :].repeat(1, final_height-img_end_y, 1, 1) - if dst_start_x > 0: canvas[:, :, 0:dst_start_x, :] = canvas[:, :, dst_start_x:dst_start_x+1, :].repeat(1, 1, dst_start_x, 1) - if img_end_x < final_width: canvas[:, :, img_end_x:final_width, :] = canvas[:, :, img_end_x-1:img_end_x, :].repeat(1, 1, final_width-img_end_x, 1) + if dst_start_y > 0: + canvas[:, 0:dst_start_y, dst_start_x:img_end_x, :] = canvas[:, dst_start_y:dst_start_y+1, dst_start_x:img_end_x, :].repeat(1, dst_start_y, 1, 1) + if img_end_y < final_height: + canvas[:, img_end_y:final_height, dst_start_x:img_end_x, :] = canvas[:, img_end_y-1:img_end_y, dst_start_x:img_end_x, :].repeat(1, final_height-img_end_y, 1, 1) + if dst_start_x > 0: + canvas[:, :, 0:dst_start_x, :] = canvas[:, :, dst_start_x:dst_start_x+1, :].repeat(1, 1, dst_start_x, 1) + if img_end_x < final_width: + canvas[:, :, img_end_x:final_width, :] = canvas[:, :, img_end_x-1:img_end_x, :].repeat(1, 1, final_width-img_end_x, 1) elif pad_mode == 'transparent_fill': canvas[:, can_slice[0], can_slice[1], :] = scaled_image[:, img_slice[0], img_slice[1], :] mask_canvas[:, can_slice[0], can_slice[1]] = scaled_image[:, img_slice[0], img_slice[1], 3] @@ -359,63 +430,67 @@ class RBGPadPro: mask_part = scaled_mask[:, img_slice[0], img_slice[1]].unsqueeze(-1) canvas[:, can_slice[0], can_slice[1], :] = img_part * mask_part + canvas[:, can_slice[0], can_slice[1], :] * (1 - mask_part) mask_canvas[:, can_slice[0], can_slice[1]] = scaled_mask[:, img_slice[0], img_slice[1]] - + if mask_blur_sigma > 0: kernel_size = self._calculate_gaussian_kernel_size(mask_blur_sigma) - if kernel_size > 1: mask_canvas = TF.gaussian_blur(mask_canvas.unsqueeze(1), kernel_size=[kernel_size, kernel_size], sigma=[mask_blur_sigma, mask_blur_sigma]).squeeze(1) + if kernel_size > 1: + mask_canvas = TF.gaussian_blur(mask_canvas.unsqueeze(1), kernel_size=[kernel_size, kernel_size], sigma=[mask_blur_sigma, mask_blur_sigma]).squeeze(1) if fill_transparent_background and C == 4: fill_color_tensor = torch.tensor(self.hex_to_rgb(transparent_fill_color), device=image.device, dtype=image.dtype).div(255.0) - alpha = canvas[..., 3:]; canvas[..., :3] = canvas[..., :3] * alpha + fill_color_tensor.view(1, 1, 1, 3) * (1 - alpha); canvas[..., 3] = 1.0 + alpha = canvas[..., 3:] + canvas[..., :3] = canvas[..., :3] * alpha + fill_color_tensor.view(1, 1, 1, 3) * (1 - alpha) + canvas[..., 3] = 1.0 + # Resizing mode with divisibility alignment (defaulting to 8 if None) if resize_mode != "none": B, H, W, C = canvas.shape if W > 0 and H > 0: + align_factor = max(8, div_factor) if div_factor > 1 else 8 ratio = target_size / max(W, H) if resize_mode == "resize_longer_side" else target_size / min(W, H) - target_width, target_height = (round(W * ratio) // 8) * 8, (round(H * ratio) // 8) * 8 + target_width = max(align_factor, (round(W * ratio) // align_factor) * align_factor) + target_height = max(align_factor, (round(H * ratio) // align_factor) * align_factor) + if target_width > 0 and target_height > 0: - # Masks are single-channel; nvidia_rtx_vsr (and lanczos) need special handling - if resample_filter in ("nvidia_rtx_vsr", "lanczos"): - mask_resample_filter = "bicubic" - else: - mask_resample_filter = resample_filter - canvas = _rtx_aware_upscale(canvas.movedim(-1,1), target_width, target_height, resample_filter).movedim(1,-1) + COMFY_NATIVE_MODES = {"nearest-exact", "bilinear", "bicubic", "area", "lanczos"} + mask_resample_filter = resample_filter if resample_filter in COMFY_NATIVE_MODES else "bicubic" + + canvas = _rtx_aware_upscale(canvas.movedim(-1, 1), target_width, target_height, resample_filter).movedim(1, -1) mask_canvas = comfy.utils.common_upscale(mask_canvas.unsqueeze(1), target_width, target_height, mask_resample_filter, "disabled").squeeze(1) region_mask = comfy.utils.common_upscale(region_mask.unsqueeze(1), target_width, target_height, "nearest-exact", "disabled").squeeze(1) + # Pad Noise with local Generator to preserve global random states if pad_noise_strength > 0: B, H, W, C = canvas.shape - torch.manual_seed(noise_seed) + gen_device = canvas.device if canvas.device.type == "cuda" else torch.device("cpu") + generator = torch.Generator(device=gen_device) + generator.manual_seed(noise_seed) - # Generate Noise if "Greyscale" in noise_mode: - greyscale_noise = torch.randn((B, H, W, 1), device=canvas.device, dtype=canvas.dtype) * pad_noise_strength + greyscale_noise = torch.randn((B, H, W, 1), generator=generator, device=gen_device, dtype=canvas.dtype).to(canvas.device) * pad_noise_strength noise_for_image = greyscale_noise.repeat(1, 1, 1, 3) else: - noise_for_image = torch.randn((B, H, W, 3), device=canvas.device, dtype=canvas.dtype) * pad_noise_strength - greyscale_noise = None # Not available for color mode + noise_for_image = torch.randn((B, H, W, 3), generator=generator, device=gen_device, dtype=canvas.dtype).to(canvas.device) * pad_noise_strength + greyscale_noise = None - # Determine Application Mask for the IMAGE if noise_mode == "Greyscale + Mask": padding_mask_for_image = (1.0 - (region_mask * mask_canvas)).clamp(0.0, 1.0).unsqueeze(-1) else: padding_mask_for_image = (1.0 - mask_canvas).unsqueeze(-1) - - # Apply noise to IMAGE + canvas[..., :3] = torch.clamp(canvas[..., :3] + noise_for_image * padding_mask_for_image, 0.0, 1.0) - # Apply noise to MASK if in the correct mode if noise_mode == "Greyscale + Mask" and greyscale_noise is not None: noise_application_mask_for_mask = padding_mask_for_image.squeeze(-1) mask_canvas = torch.clamp(mask_canvas + greyscale_noise.squeeze(-1) * noise_application_mask_for_mask, 0.0, 1.0) - if invert_mask: mask_canvas = 1.0 - mask_canvas - + if invert_mask: + mask_canvas = 1.0 - mask_canvas + canvas = canvas[..., :3] if show_preview: try: - # We only need to preview the first image of the batch preview_image_tensor = canvas[0:1] preview_image = self.prepare_image_for_display(preview_image_tensor) PromptServer.instance.send_sync("rbg_pad_pro_preview", {"image": preview_image}) @@ -425,5 +500,6 @@ class RBGPadPro: image_info = self.get_image_info(canvas) return {"ui": {"text": [image_info]}, "result": (canvas, mask_canvas)} -NODE_CLASS_MAPPINGS = { "RBGPadPro": RBGPadPro } -NODE_DISPLAY_NAME_MAPPINGS = { "RBGPadPro": "RBG Pad Pro" } \ No newline at end of file + +NODE_CLASS_MAPPINGS = {"RBGPadPro": RBGPadPro} +NODE_DISPLAY_NAME_MAPPINGS = {"RBGPadPro": "RBG Pad Pro"} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 0d59398..1308eeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ numpy Pillow +nvidia-vfx \ No newline at end of file diff --git a/web/nodes/RBGPadPro.js b/web/nodes/RBGPadPro.js index f8f9922..3da8f72 100644 --- a/web/nodes/RBGPadPro.js +++ b/web/nodes/RBGPadPro.js @@ -2,10 +2,9 @@ import { app } from "/scripts/app.js"; import { api } from "/scripts/api.js"; app.registerExtension({ - name: "RBGSuitePack.PadPro.Complete", // Updated name to ensure cache is cleared + name: "RBGSuitePack.PadPro.Complete", async beforeRegisterNodeDef(nodeType, nodeData, app) { if (nodeData.name === "RBGPadPro") { - // This function runs once when the node is created. const onNodeCreated = nodeType.prototype.onNodeCreated; const MIN_NODE_HEIGHT_WITH_PREVIEW = 400; const MIN_NODE_HEIGHT_WITHOUT_PREVIEW = 250; @@ -16,27 +15,34 @@ app.registerExtension({ nodeType.prototype.onNodeCreated = function () { onNodeCreated?.apply(this, arguments); - // --- START: RE-ADD TOOLTIPS --- + // Tooltips configuration const tooltips = { - pad_mode: "Choose how to handle the padding. 'Pad' adds space, 'pad_edge' uses edge colors, 'pad_edge_pixel' extends edge pixels, 'pad_mirror' reflects the image edges, 'transparent_fill' fills transparent areas, and 'crop' cuts the image.", + pad_mode: "Choose how to handle the padding. 'pad' adds solid color, 'pad_edge' averages border colors, 'pad_edge_pixel' extends border pixels, and 'transparent_fill' fills alpha channels.", pad_left: "Amount of padding to add to the left side.", pad_right: "Amount of padding to add to the right side.", pad_top: "Amount of padding to add to the top.", pad_bottom: "Amount of padding to add to the bottom.", - mask_blur_sigma: "Applies a Gaussian blur to the generated mask, blending the padded area more smoothly.", - pad_color: "The color of the padded area.", - image_position: "The base position of the image on the canvas.", - image_offset_x: "Fine-tune the image's horizontal position.", - image_offset_y: "Fine-tune the image's vertical position.", - image_scale: "Scale the image before padding. The canvas size is not affected.", + mask_blur_sigma: "Applies a Gaussian blur to the generated mask for smoother inpainting seams.", + pad_color: "The solid color of the padded area.", + image_position: "The base alignment/anchor of the image on the canvas.", + image_offset_x: "Fine-tune the horizontal position in pixels.", + image_offset_y: "Fine-tune the vertical position in pixels.", + image_scale: "Scale the image before padding without altering the canvas bounds.", fill_transparent_background: "Fill transparent areas with a solid color.", - transparent_fill_color: "The color to use for transparent areas.", - pad_aspect_ratio: "Automatically adjust padding to match a specific aspect ratio.", - resize_mode: "Choose how to resize the final image. 'resize_longer_side' and 'resize_shorter_side' maintain the aspect ratio.", + transparent_fill_color: "The color to use for transparent background fill.", + pad_aspect_ratio: "Target aspect ratio preset.", + resize_mode: "Choose how to resize the final image ('resize_longer_side' or 'resize_shorter_side').", target_size: "The target size in pixels for the selected side.", resample_filter: "The interpolation method to use for resizing.", - auto_crop: "Automatically crop the image to a 1:1 square aspect ratio before any other processing.", - invert_mask: "Invert the generated mask. Useful for outpainting workflows where you want to mask the padded area instead of the original image." + auto_crop_aspect_ratio: "Crops the image to match the selected pad_aspect_ratio instead of expanding the canvas with padding.", + invert_mask: "Inverts the output mask (Standard for Outpainting).", + flip_horizontal: "Horizontally flips the image and mask.", + image_rotation: "Rotates the input image by degrees.", + BorderCrop_threshold: "Automatically trims solid color borders before processing.", + pad_noise_strength: "Strength of noise added to the padded region to prep diffusion latents.", + noise_seed: "Seed for procedural noise generation.", + noise_mode: "Color, Greyscale, or Greyscale+Mask noise.", + divisibility: "Enforces dimensions to be multiples of 8, 16, 32, 64, etc. for model compatibility." }; for (const widget of this.widgets) { @@ -44,13 +50,12 @@ app.registerExtension({ widget.canvas.title = tooltips[widget.name]; } } - // --- END: RE-ADD TOOLTIPS --- - // Create a dedicated text widget for the output, same as KJNodes. + // Output info widget const textWidget = this.addWidget("text", "output_text", "", {}); - textWidget.serialize = false; // Don't save this text in the workflow file + textWidget.serialize = false; - // --- START: ADD RESET BUTTON --- + // Reset Button const resetButton = this.addWidget("button", "Reset to Defaults 🗑️", "reset", () => { const defaults = nodeData.input.required; for (const widget of this.widgets) { @@ -61,9 +66,8 @@ app.registerExtension({ } }); resetButton.serialize = false; - // --- END: ADD RESET BUTTON --- - // --- START: ADD PREVIEW --- + // Preview handling const previewWidget = this.widgets.find(w => w.name === "show_preview"); if (previewWidget) { const originalCallback = previewWidget.callback; @@ -76,7 +80,6 @@ app.registerExtension({ }; } - // Set initial size this.size[1] = previewWidget?.value ? MIN_NODE_HEIGHT_WITH_PREVIEW : MIN_NODE_HEIGHT_WITHOUT_PREVIEW; api.addEventListener("rbg_pad_pro_preview", (event) => { @@ -84,10 +87,8 @@ app.registerExtension({ img.src = event.detail.image; } }); - // --- END: ADD PREVIEW --- }; - // --- START: ADD PREVIEW DRAWING --- const onDrawBackground = nodeType.prototype.onDrawBackground; nodeType.prototype.onDrawBackground = function(ctx) { onDrawBackground?.apply(this, arguments); @@ -99,7 +100,7 @@ app.registerExtension({ const [w, h] = this.size; const lastWidget = this.widgets[this.widgets.length - 1]; - const lastWidgetY = lastWidget.last_y || 0; + const lastWidgetY = lastWidget?.last_y || 0; const PADDING = 10; const IMAGE_Y_OFFSET = lastWidgetY + 30; const imageAreaHeight = h - IMAGE_Y_OFFSET - PADDING; @@ -118,12 +119,10 @@ app.registerExtension({ } }; - // This function runs every time after the node is executed. const onExecuted = nodeType.prototype.onExecuted; nodeType.prototype.onExecuted = function (message) { onExecuted?.apply(this, arguments); - // Find our output widget and update its value with the message from Python. if (message?.text) { const widget = this.widgets.find(w => w.name === "output_text"); if (widget) { @@ -133,4 +132,4 @@ app.registerExtension({ }; } } -}); +}); \ No newline at end of file