import torch import torch.nn.functional as F import os import math import urllib.request from PIL import Image, ImageOps import numpy as np import folder_paths import comfy.model_management as mm from comfy.utils import load_torch_file, ProgressBar try: import spandrel SPANDREL_AVAILABLE = True except ImportError: SPANDREL_AVAILABLE = False try: from transformers import AutoModelForImageSegmentation from torchvision import transforms import torchvision.transforms.functional as TF TRANSFORMERS_AVAILABLE = True except ImportError: TRANSFORMERS_AVAILABLE = False try: import onnxruntime as ort ONNXRUNTIME_AVAILABLE = True except ImportError: ONNXRUNTIME_AVAILABLE = False class QuickMergeNode: def __init__(self): self.cached_bg_model = None self.cached_bg_model_name = None self.cached_upscaler = None self.cached_upscaler_name = None self.onnx_session = None @classmethod def INPUT_TYPES(cls): upscale_models = folder_paths.get_filename_list("upscale_models") auto_dl_name = "RealESRGAN_x2plus.pth (Auto-Download)" if not upscale_models: upscale_models = [auto_dl_name] elif "RealESRGAN_x2plus.pth" not in upscale_models: upscale_models.insert(0, auto_dl_name) resize_modes = ["Crop to Fill (Preserve Aspect)", "Stretch to Fit", "Fit Inside (Letterbox)", "Fit Inside (Pad with Blur)"] align_modes = ["center", "top", "bottom", "left", "right"] return { "required": {}, "optional": { "Background": ("IMAGE",), "Foreground (optional)": ("IMAGE",), "mask (optional)": ("MASK",), # --- 0. BATCH PROCESSING --- "📁 batch_processing": ("BOOLEAN", { "default": False, "tooltip": "CRITICAL: Ensure all images processed in a folder batch share identical resolutions and aspect ratios to prevent platform layout shifts." }), "📁 foreground_batch": ("BOOLEAN", {"default": False}), "📁 fg_batch_folder": ("STRING", {"default": "📁 Type or paste folder path here..."}), "📁 background_batch": ("BOOLEAN", {"default": False}), "📁 bg_batch_folder": ("STRING", {"default": "📁 Type or paste folder path here..."}), "📁 batch_pairing_mode": (["1-to-1 Match (Sequential)", "Cross-Multiply (All Combinations)"],), "📁 batch_chunk_size": ("INT", { "default": 4, "min": 1, "max": 128, "step": 1, "tooltip": "Max images processed at once through the heavy pipeline. 0 = no limit." }), # --- 1. INITIAL RESIZE --- "📏 initial_resize": ("BOOLEAN", {"default": False}), "📏 fg_resize_mode": (resize_modes, {"default": "Crop to Fill (Preserve Aspect)"}), "📏 fg_crop_align": (align_modes, {"default": "center"}), "📏 fg_initial_width": ("INT", {"default": 720, "min": 8, "max": 16384, "step": 8}), "📏 fg_initial_height": ("INT", {"default": 1280, "min": 8, "max": 16384, "step": 8}), "📏 bg_resize_mode": (resize_modes, {"default": "Crop to Fill (Preserve Aspect)"}), "📏 bg_crop_align": (align_modes, {"default": "center"}), "📏 bg_initial_width": ("INT", {"default": 1920, "min": 8, "max": 16384, "step": 8}), "📏 bg_initial_height": ("INT", {"default": 1080, "min": 8, "max": 16384, "step": 8}), # --- 2. FOREGROUND MERGE --- "✂️ merge_image": ("BOOLEAN", {"default": False}), "✂️ cutout_model": (["BiRefNet_lite (Fast/Accurate)", "RMBG-2.0 (New/Accurate)", "RMBG-1.4 (Fast)", "BiRefNet (High Quality)", "BiRefNet_HR (4K/8K High Quality)", "None (Use Alpha)"],), "✂️ auto-position (WiP)": ("BOOLEAN", {"default": False}), "✂️ foreground_size": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), "✂️ foreground_left_right": ("INT", {"default": 0, "min": -8192, "max": 8192, "step": 1}), "✂️ foreground_up_down": ("INT", {"default": 0, "min": -8192, "max": 8192, "step": 1}), "✂️ flip_horizontal": ("BOOLEAN", {"default": False}), "✂️ flip_vertical": ("BOOLEAN", {"default": False}), "✂️ rotation_degrees": ("INT", {"default": 0, "min": -360, "max": 360, "step": 1}), "✂️ edge_shrink_grow": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}), "✂️ edge_softness": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), "✂️ foreground_opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), "✂️ blend_mode": (["Normal", "Multiply", "Screen", "Overlay", "Soft Light"],), "✂️ blend_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), # --- 3. AI SUPERSAMPLING --- "⚡ ai_upscale": ("BOOLEAN", {"default": False}), "⚡ pipeline_order": (["upscale composite", "upscale foreground only", "upscale background only"],), "⚡ choose_upscale_model": (upscale_models,), "⚡ upscale_by": ("FLOAT", { "default": 2.0, "min": 0.25, "max": 32.0, "step": 0.25, }), "⚡ tile_size (0 means off/native)": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 32}), "⚡ post_upscale_edge_soften": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), # --- 4. COLORS & FILTERS --- "🎨 colors_and_filters": ("BOOLEAN", {"default": False}), "🎨 filter_target": (["Foreground", "Background", "Both"],), "🎨 filter_type": (["Gaussian Blur", "Auto-Color Match", "Laplacian Sharpen", "Edge Detect", "Sepia", "Invert", "None"],), "🎨 filter_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05}), "🎨 brightness": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.05}), "🎨 contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}), "🎨 saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}), "🎨 filter_size": ("INT", {"default": 3, "min": 1, "max": 127, "step": 2}), "🎨 background_blur": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), # --- 5. SHADOWS --- "🌗 shadows": ("BOOLEAN", {"default": False}), "🌗 shadow_color_hex": ("STRING", {"default": "#000000"}), "🌗 shadow_darkness": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "🌗 shadow_softness": ("INT", {"default": 15, "min": 1, "max": 255, "step": 1}), "🌗 shadow_position_x": ("INT", {"default": 10, "min": -1000, "max": 1000, "step": 1}), "🌗 shadow_position_y": ("INT", {"default": 10, "min": -1000, "max": 1000, "step": 1}), # --- 6. TARGET RESOLUTION --- "🎯 target_resolution": ("BOOLEAN", {"default": False}), "🎯 target_resize_mode": (resize_modes,), "🎯 target_crop_align": (align_modes,), "🎯 target_width": ("INT", {"default": 1920, "min": 64, "max": 16384, "step": 8}), "🎯 target_height": ("INT", {"default": 1080, "min": 64, "max": 16384, "step": 8}), # --- GLOBAL --- "🖼️ pixel_smoothing": (["bicubic", "lanczos", "bilinear", "nearest"],), } } RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "IMAGE") RETURN_NAMES = ("composited_image", "light_mask", "relight_fg_crop", "relight_bg_crop") FUNCTION = "process" CATEGORY = "image/compositing" def color_match(self, fg, bg, mask): mask_sum = mask.sum(dim=(2, 3), keepdim=True) + 1e-6 mu_fg = (fg * mask).sum(dim=(2, 3), keepdim=True) / mask_sum std_fg = torch.sqrt(torch.sum((((fg - mu_fg) * mask) ** 2), dim=(2, 3), keepdim=True) / mask_sum) + 1e-6 mu_bg = bg.mean(dim=(2, 3), keepdim=True) std_bg = bg.std(dim=(2, 3), keepdim=True) + 1e-6 matched_fg = (fg - mu_fg) * (std_bg / std_fg) + mu_bg return torch.clamp(matched_fg, 0.0, 1.0) def load_bg_model(self, model_choice, device): if not TRANSFORMERS_AVAILABLE: raise Exception("Please install transformers via ComfyUI Manager") if self.cached_bg_model_name == model_choice: return self.cached_bg_model mm.soft_empty_cache() original_getattr = torch.nn.Module.__getattr__ def safe_getattr(module_self, name): if name == "all_tied_weights_keys": return {} return original_getattr(module_self, name) try: torch.nn.Module.__getattr__ = safe_getattr if model_choice == "RMBG-1.4 (Fast)": model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True) elif model_choice == "RMBG-2.0 (New/Accurate)": model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True) elif model_choice == "BiRefNet_lite (Fast/Accurate)": model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet_lite", trust_remote_code=True) elif model_choice == "BiRefNet (High Quality)": model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True) elif model_choice == "BiRefNet_HR (4K/8K High Quality)": model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet_HR", trust_remote_code=True) finally: torch.nn.Module.__getattr__ = original_getattr model.to(device) model.eval() self.cached_bg_model = model self.cached_bg_model_name = model_choice return model def apply_bg_removal(self, image, model_choice, device): model = self.load_bg_model(model_choice, device) transform = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) img_t = image.permute(0, 3, 1, 2) if img_t.shape[1] == 4: img_t = img_t[:, :3, :, :] model_dtype = next(model.parameters()).dtype b_size = img_t.shape[0] out_masks = [] for i in range(b_size): chunk = img_t[i:i+1].to(device) input_tensor = transform(chunk).to(dtype=model_dtype) with torch.no_grad(): outputs = model(input_tensor) def get_all_tensors(obj): if isinstance(obj, torch.Tensor): return [obj] elif isinstance(obj, (list, tuple)): res = [] for item in obj: res.extend(get_all_tensors(item)) return res elif isinstance(obj, dict): res = [] for val in obj.values(): res.extend(get_all_tensors(val)) return res return [] all_tensors = get_all_tensors(outputs) preds_tensor = all_tensors[0] if "RMBG" in model_choice else all_tensors[-1] preds_tensor = preds_tensor.float().sigmoid() preds_tensor = (preds_tensor - preds_tensor.min()) / (preds_tensor.max() - preds_tensor.min() + 1e-8) preds_tensor = torch.where(preds_tensor < 0.05, torch.zeros_like(preds_tensor), preds_tensor) mask = F.interpolate(preds_tensor, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) if mask.shape[1] > 1: mask = mask[:, 0:1, :, :] out_masks.append(mask.cpu()) mm.soft_empty_cache() return torch.cat(out_masks, dim=0) def tiled_upscale(self, model, tensor, tile_size=512, overlap=32, pbar=None, start_pct=0, end_pct=100): b, c, h, w = tensor.shape device = tensor.device is_mask = (c == 1) if is_mask: tensor = tensor.repeat(1, 3, 1, 1) with torch.no_grad(): dummy = torch.zeros((1, 3, 16, 16), device=device) scale = model(dummy).shape[-1] // 16 if tile_size == 0: with torch.no_grad(): result = model(tensor) if is_mask: result = result.mean(dim=1, keepdim=True) if pbar: pbar.update_absolute(end_pct) return result out_h, out_w = h * scale, w * scale output = torch.zeros((b, 3, out_h, out_w), device=device) weights = torch.zeros((b, 1, out_h, out_w), device=device) stride = max(1, tile_size - overlap) y_range = list(range(0, h, stride)) x_range = list(range(0, w, stride)) total_tiles = max(1, len(y_range) * len(x_range)) tiles_done = 0 for y in y_range: for x in x_range: y1 = max(0, min(y, h - tile_size)) if h > tile_size else 0 x1 = max(0, min(x, w - tile_size)) if w > tile_size else 0 y2 = min(y1 + tile_size, h) x2 = min(x1 + tile_size, w) tile = tensor[:, :, y1:y2, x1:x2] tile_h, tile_w = tile.shape[2], tile.shape[3] gy = torch.linspace(-1.0, 1.0, tile_h, device=device) gx = torch.linspace(-1.0, 1.0, tile_w, device=device) yy, xx = torch.meshgrid(gy, gx, indexing='ij') g_map = torch.exp(-(xx ** 2 + yy ** 2) / (2 * 0.5 ** 2)) g_map = g_map.unsqueeze(0).unsqueeze(0) with torch.no_grad(): up_tile = model(tile) up_th, up_tw = up_tile.shape[2], up_tile.shape[3] g_map_up = F.interpolate(g_map, size=(up_th, up_tw), mode='bilinear', align_corners=False) out_y1, out_x1 = y1 * scale, x1 * scale out_y2, out_x2 = out_y1 + up_th, out_x1 + up_tw output[:, :, out_y1:out_y2, out_x1:out_x2] += up_tile * g_map_up weights[:, :, out_y1:out_y2, out_x1:out_x2] += g_map_up del tile, up_tile, g_map, g_map_up tiles_done += 1 if pbar: pbar.update_absolute(int(start_pct + (tiles_done / total_tiles) * (end_pct - start_pct))) final_out = output / weights.clamp(min=1e-6) if is_mask: final_out = final_out.mean(dim=1, keepdim=True) return final_out def process_alpha(self, mask, dilate_erode, feather): if dilate_erode > 0: mask = F.max_pool2d(mask, kernel_size=dilate_erode*2+1, stride=1, padding=dilate_erode) elif dilate_erode < 0: erode_val = abs(dilate_erode) mask = -F.max_pool2d(-mask, kernel_size=erode_val*2+1, stride=1, padding=erode_val) if feather > 0: kernel_size = feather * 2 + 1 sigma = 0.3 * ((kernel_size - 1) * 0.5 - 1) + 0.8 x = torch.arange(kernel_size, dtype=torch.float32, device=mask.device) - kernel_size // 2 x = x.expand(kernel_size, kernel_size) y = x.t() kernel = torch.exp(-(x**2 + y**2) / (2 * sigma**2)) kernel = kernel / kernel.sum() kernel = kernel.view(1, 1, kernel_size, kernel_size) mask = F.conv2d(mask, kernel, padding=feather) return torch.clamp(mask, 0.0, 1.0) def generate_gaussian_blur(self, t_img, blur_amount, device): k_size = blur_amount if blur_amount % 2 != 0 else blur_amount + 1 sigma = max(0.3 * ((k_size - 1) * 0.5 - 1) + 0.8, 0.1) coords = torch.arange(k_size, dtype=torch.float32, device=device) - k_size // 2 x_g = coords.expand(k_size, k_size) y_g = x_g.t() g_kernel = torch.exp(-(x_g**2 + y_g**2) / (2 * sigma**2)) g_kernel = (g_kernel / g_kernel.sum()).view(1, 1, k_size, k_size).expand(t_img.shape[1], 1, k_size, k_size) return F.conv2d(t_img, g_kernel, padding=k_size//2, groups=t_img.shape[1]) # --- RE-INSTATED METHODS: RESIZE AND GRADIENT ENGINE SLOTS --- def generate_smart_light_gradient(self, mask_tensor, target_h, target_w, device): if mask_tensor is None: return torch.zeros((1, target_h, target_w), dtype=torch.float32).cpu() b_size = mask_tensor.shape[0] tiny_size = 128 out_masks = [] for i in range(b_size): m_tiny = F.interpolate(mask_tensor[i:i+1].float().to(device), size=(tiny_size, tiny_size), mode='area') m_blurred = self.generate_gaussian_blur(m_tiny, 31, device) m_final = F.interpolate(m_blurred, size=(target_h, target_w), mode='bicubic', align_corners=False) out_masks.append(torch.clamp(m_final, 0.0, 1.0).squeeze(1).cpu()) return torch.cat(out_masks, dim=0) def resize_tensor(self, tensor, target_w, target_h, mode, align, interpolation): b, c, h, w = tensor.shape target_ratio = target_w / target_h tensor_ratio = w / h if mode == "Crop to Fill (Preserve Aspect)": if abs(tensor_ratio - target_ratio) > 0.001: if tensor_ratio > target_ratio: crop_w = int(h * target_ratio) if align == "left": start_x = 0 elif align == "right": start_x = w - crop_w else: start_x = (w - crop_w) // 2 tensor = tensor[:, :, :, start_x:start_x + crop_w] else: crop_h = int(w / target_ratio) if align == "top": start_y = 0 elif align == "bottom": start_y = h - crop_h else: start_y = (h - crop_h) // 2 tensor = tensor[:, :, start_y:start_y + crop_h, :] return F.interpolate(tensor, size=(target_h, target_w), mode=interpolation, align_corners=False) elif mode == "Stretch to Fit": return F.interpolate(tensor, size=(target_h, target_w), mode=interpolation, align_corners=False) elif mode == "Fit Inside (Letterbox)": if tensor_ratio > target_ratio: new_w = target_w new_h = int(target_w / tensor_ratio) else: new_h = target_h new_w = int(target_h * tensor_ratio) tensor_scaled = F.interpolate(tensor, size=(max(1, new_h), max(1, new_w)), mode=interpolation, align_corners=False) padded = torch.zeros((b, c, target_h, target_w), device=tensor.device) start_y = (target_h - new_h) // 2 start_x = (target_w - new_w) // 2 padded[:, :, start_y:start_y+new_h, start_x:start_x+new_w] = tensor_scaled return padded elif mode == "Fit Inside (Pad with Blur)": bg_fill = F.interpolate(tensor, size=(target_h, target_w), mode=interpolation, align_corners=False) bg_blur = self.generate_gaussian_blur(bg_fill.to(mm.get_torch_device()), 51, mm.get_torch_device()).cpu() if tensor.device.type == 'cpu' else self.generate_gaussian_blur(bg_fill, 51, tensor.device) if tensor_ratio > target_ratio: new_w = target_w new_h = int(target_w / tensor_ratio) else: new_h = target_h new_w = int(target_h * tensor_ratio) tensor_scaled = F.interpolate(tensor, size=(max(1, new_h), max(1, new_w)), mode=interpolation, align_corners=False) start_y = (target_h - new_h) // 2 start_x = (target_w - new_w) // 2 bg_blur[:, :, start_y:start_y+new_h, start_x:start_x+new_w] = tensor_scaled return bg_blur # --- SECURE PATH VALIDATION ENGINE --- def get_secure_path(self, user_path): """ Validates that the provided folder path is strictly inside ComfyUI's input directory. Prevents directory traversal (../) and symlink security exploits. """ if not user_path or user_path == "📁 Type or paste folder path here...": return None # Use realpath to resolve malicious symlinks pointing outside the directory base_dir = os.path.realpath(folder_paths.get_input_directory()) full_path = os.path.realpath(os.path.join(base_dir, user_path)) try: if os.path.commonpath([base_dir, full_path]) != base_dir: print(f"[Quick Merge] SECURITY BLOCK: Attempted path traversal outside input directory -> {user_path}") return None return full_path except ValueError: # Handles edge cases where paths exist on different drives (Windows) print(f"[Quick Merge] SECURITY BLOCK: Invalid path (different drive) -> {user_path}") return None def load_folder_images(self, folder_path, device, is_active): if not is_active or not os.path.isdir(folder_path): return None valid_exts = ['.jpg', '.jpeg', '.png', '.webp'] processed_tensors = [] first_w, first_h = None, None for file in sorted(os.listdir(folder_path)): if any(file.lower().endswith(ext) for ext in valid_exts): img_path = os.path.join(folder_path, file) try: i = Image.open(img_path) i = ImageOps.exif_transpose(i) i = i.convert("RGBA") img_arr = np.array(i).astype(np.float32) / 255.0 t = torch.from_numpy(img_arr).unsqueeze(0).cpu() t_perm = t.permute(0, 3, 1, 2) if first_w is None: first_w, first_h = t_perm.shape[3], t_perm.shape[2] elif t_perm.shape[3] != first_w or t_perm.shape[2] != first_h: t_perm = self.resize_tensor(t_perm, first_w, first_h, "Fit Inside (Letterbox)", "center", "bicubic") processed_tensors.append(t_perm.permute(0, 2, 3, 1)) except Exception as e: print(f"Quick Merge Image Pack Alignment Failed: {file} - {e}") if not processed_tensors: return None return torch.cat(processed_tensors, dim=0) def process(self, **kwargs): device = mm.get_torch_device() alpha = None fg = None b_size = 1 effective_chunk = 1 num_chunks = 1 multi_chunk = False pre_transform_bg_h = 1080 pre_transform_bg_w = 1920 pbar = ProgressBar(100) # --- MAP SETTINGS BLOCK CONTEXT --- batch_processing = kwargs.get("📁 batch_processing", False) foreground_batch = kwargs.get("📁 foreground_batch", False) background_batch = kwargs.get("📁 background_batch", False) fg_batch_folder = kwargs.get("📁 fg_batch_folder", "").strip() bg_batch_folder = kwargs.get("📁 bg_batch_folder", "").strip() batch_pairing_mode = kwargs.get("📁 batch_pairing_mode", "1-to-1 Match (Sequential)") batch_chunk_size = kwargs.get("📁 batch_chunk_size", 4) initial_resize = kwargs.get("📏 initial_resize", False) fg_resize_mode = kwargs.get("📏 fg_resize_mode", "Crop to Fill (Preserve Aspect)") fg_crop_align = kwargs.get("📏 fg_crop_align", "center") fg_initial_width = kwargs.get("📏 fg_initial_width", 720) fg_initial_height = kwargs.get("📏 fg_initial_height", 1280) bg_resize_mode = kwargs.get("📏 bg_resize_mode", "Crop to Fill (Preserve Aspect)") bg_crop_align = kwargs.get("📏 bg_crop_align", "center") bg_initial_width = kwargs.get("📏 bg_initial_width", 1920) bg_initial_height = kwargs.get("📏 bg_initial_height", 1080) cutout_model = kwargs.get("✂️ cutout_model", "BiRefNet_lite (Fast/Accurate)") foreground_size = kwargs.get("✂️ foreground_size", 1.0) foreground_left_right = kwargs.get("✂️ foreground_left_right", 0) foreground_up_down = kwargs.get("✂️ foreground_up_down", 0) edge_shrink_grow = kwargs.get("✂️ edge_shrink_grow", 0) edge_softness = kwargs.get("✂️ edge_softness", 0) flip_horizontal = kwargs.get("✂️ flip_horizontal", False) flip_vertical = kwargs.get("✂️ flip_vertical", False) rotation_degrees = kwargs.get("✂️ rotation_degrees", 0) foreground_opacity = kwargs.get("✂️ foreground_opacity", 1.0) blend_mode = kwargs.get("✂️ blend_mode", "Normal") blend_strength = kwargs.get("✂️ blend_strength", 1.0) auto_position = False ai_upscale = kwargs.get("⚡ ai_upscale", False) pipeline_order = kwargs.get("⚡ pipeline_order", "upscale composite") ai_upscale_model = kwargs.get("⚡ choose_upscale_model", "None Found") upscale_by = kwargs.get("⚡ upscale_by", 2.0) ai_upscale_tile_size = kwargs.get("⚡ tile_size (0 means off/native)", 512) post_upscale_edge_soften = kwargs.get("⚡ post_upscale_edge_soften", 0) colors_and_filters = kwargs.get("🎨 colors_and_filters", False) filter_target = kwargs.get("🎨 filter_target", "Foreground") color_and_filters = kwargs.get("🎨 filter_type", "Auto-Color Match") filter_strength = kwargs.get("🎨 filter_strength", 1.0) filter_size = kwargs.get("🎨 filter_size", 3) brightness = kwargs.get("🎨 brightness", 0.0) contrast = kwargs.get("🎨 contrast", 1.0) saturation = kwargs.get("🎨 saturation", 1.0) background_blur = kwargs.get("🎨 background_blur", 0) shadows = kwargs.get("🌗 shadows", False) shadow_color_hex = kwargs.get("🌗 shadow_color_hex", "#000000").lstrip('#') if len(shadow_color_hex) != 6: shadow_color_hex = "000000" shadow_R = int(shadow_color_hex[0:2], 16) / 255.0 shadow_G = int(shadow_color_hex[2:4], 16) / 255.0 shadow_B = int(shadow_color_hex[4:6], 16) / 255.0 shadow_darkness = kwargs.get("🌗 shadow_darkness", 0.5) shadow_softness = kwargs.get("🌗 shadow_softness", 15) shadow_position_x = kwargs.get("🌗 shadow_position_x", 10) shadow_position_y = kwargs.get("🌗 shadow_position_y", 10) target_resolution = kwargs.get("🎯 target_resolution", False) target_resize_mode = kwargs.get("🎯 target_resize_mode", "Crop to Fill (Preserve Aspect)") target_crop_align = kwargs.get("🎯 target_crop_align", "center") target_width = kwargs.get("🎯 target_width", 1920) target_height = kwargs.get("🎯 target_height", 1080) pixel_smoothing = kwargs.get("🖼️ pixel_smoothing", "bicubic") interpolation = pixel_smoothing if pixel_smoothing != "lanczos" else "bicubic" # --- 1. DIRECTORY CONDITIONAL PROCESSING OVERRIDES (SECURED) --- bg_input = kwargs.get("Background", None) fg_input = kwargs.get("Foreground (optional)", None) mask_input = kwargs.get("mask (optional)", None) default_folder_text = "📁 Type or paste folder path here..." if fg_batch_folder == default_folder_text: fg_batch_folder = "" if bg_batch_folder == default_folder_text: bg_batch_folder = "" if batch_processing: if background_batch and bg_batch_folder: secure_bg = self.get_secure_path(bg_batch_folder) if secure_bg and os.path.isdir(secure_bg): bg_input = self.load_folder_images(secure_bg, device, True) if foreground_batch and fg_batch_folder: secure_fg = self.get_secure_path(fg_batch_folder) if secure_fg and os.path.isdir(secure_fg): fg_input = self.load_folder_images(secure_fg, device, True) has_bg = bg_input is not None has_fg = fg_input is not None has_mask = mask_input is not None if not has_bg and not has_fg and not has_mask: empty_img = torch.zeros((1, 64, 64, 3), dtype=torch.float32).cpu() empty_msk = torch.zeros((1, 64, 64), dtype=torch.float32).cpu() return (empty_img, empty_msk, empty_img, empty_img) if has_bg: bg = bg_input.clone().detach().cpu().permute(0, 3, 1, 2) if bg_input.shape[3] in [3, 4] else bg_input.clone().detach().cpu() if bg.shape[1] == 4: bg = bg[:, :3, :, :] if has_fg: fg = fg_input.clone().detach().cpu().permute(0, 3, 1, 2) if fg_input.shape[3] in [3, 4] else fg_input.clone().detach().cpu() if has_mask: m_tensor = mask_input.clone().detach().cpu() if len(m_tensor.shape) == 3: m_tensor = m_tensor.unsqueeze(1) merge_image = kwargs.get("✂️ merge_image", True) if not has_fg: merge_image = False is_real_bg = has_bg if has_bg: b_size = bg.shape[0] elif has_fg: if merge_image: b_size = fg.shape[0] bg = torch.zeros((b_size, 3, fg.shape[2], fg.shape[3]), device='cpu') has_bg = True else: bg = fg.clone() has_bg = True is_real_bg = True b_size = bg.shape[0] elif has_mask: b_size = m_tensor.shape[0] bg = torch.zeros((b_size, 3, m_tensor.shape[2], m_tensor.shape[3]), device='cpu') has_bg = True if not merge_image: has_fg = False # --- PRE-PROCESS FOREGROUND CUTOUT ENGINE --- if has_fg and merge_image: if cutout_model != "None (Use Alpha)": pass_fg = fg.permute(0, 2, 3, 1) alpha = self.apply_bg_removal(pass_fg, cutout_model, device) del pass_fg if fg.shape[1] == 4: fg = fg[:, :3, :, :] else: if fg.shape[1] == 4: alpha = fg[:, 3:4, :, :] fg = fg[:, :3, :, :] else: alpha = torch.ones((fg.shape[0], 1, fg.shape[2], fg.shape[3]), dtype=torch.float32, device='cpu') alpha = self.process_alpha(alpha, edge_shrink_grow, edge_softness) if flip_horizontal: fg, alpha = torch.flip(fg, [3]), torch.flip(alpha, [3]) if flip_vertical: fg, alpha = torch.flip(fg, [2]), torch.flip(alpha, [2]) if rotation_degrees != 0: fg = TF.rotate(fg, rotation_degrees, interpolation=TF.InterpolationMode.BILINEAR, expand=True) alpha = TF.rotate(alpha, rotation_degrees, interpolation=TF.InterpolationMode.BILINEAR, expand=True) # --- COMBINATORIAL PAIRING SYNCS --- if is_real_bg and merge_image and has_fg: b_bg = bg.shape[0] b_fg = fg.shape[0] if batch_pairing_mode == "1-to-1 Match (Sequential)": b_size = max(b_bg, b_fg) if b_bg < b_size: repeats_bg = math.ceil(b_size / b_bg) bg = bg.repeat(repeats_bg, 1, 1, 1)[:b_size] if b_fg < b_size: repeats_fg = math.ceil(b_size / b_fg) fg = fg.repeat(repeats_fg, 1, 1, 1)[:b_size] if alpha is not None: alpha = alpha.repeat(repeats_fg, 1, 1, 1)[:b_size] else: b_size = b_bg * b_fg bg = bg.repeat_interleave(b_fg, dim=0) fg = fg.repeat(b_bg, 1, 1, 1) if alpha is not None: alpha = alpha.repeat(b_bg, 1, 1, 1) # --- UNCONDITIONAL CHUNK EVALUATION METRICS --- pre_transform_bg_h, pre_transform_bg_w = bg.shape[2], bg.shape[3] effective_chunk = batch_chunk_size if (batch_chunk_size > 0 and b_size > batch_chunk_size) else b_size num_chunks = math.ceil(b_size / effective_chunk) multi_chunk = num_chunks > 1 if has_mask: if m_tensor.shape[0] != b_size: m_tensor = m_tensor.repeat(math.ceil(b_size / m_tensor.shape[0]), 1, 1, 1)[:b_size] if initial_resize: bg = self.resize_tensor(bg, bg_initial_width, bg_initial_height, bg_resize_mode, bg_crop_align, interpolation) if merge_image and has_fg: fg = self.resize_tensor(fg, fg_initial_width, fg_initial_height, fg_resize_mode, fg_crop_align, interpolation) alpha = self.resize_tensor(alpha, fg_initial_width, fg_initial_height, fg_resize_mode, fg_crop_align, interpolation) pbar.update_absolute(5) # ---- CLOSURES ---- def execute_filters(t_fg, t_alpha, t_bg): if not colors_and_filters: return t_fg if brightness != 0.0 or contrast != 1.0: t_fg = t_fg + brightness t_fg = torch.clamp((t_fg - 0.5) * contrast + 0.5, 0.0, 1.0) if saturation != 1.0: lum = 0.299 * t_fg[:, 0:1] + 0.587 * t_fg[:, 1:2] + 0.114 * t_fg[:, 2:3] t_fg = torch.clamp(torch.lerp(lum.repeat(1, 3, 1, 1), t_fg, saturation), 0.0, 1.0) if color_and_filters == "Laplacian Sharpen": lap_kernel = torch.tensor([[[[0., 1., 0.], [1., -4., 1.], [0., 1., 0.]]]], device=device).expand(t_fg.shape[1], 1, 3, 3) edges = F.conv2d(t_fg, lap_kernel, padding=1, groups=t_fg.shape[1]) return torch.clamp(torch.lerp(t_fg, t_fg - (edges * filter_strength), min(filter_strength, 1.0)), 0.0, 1.0) elif color_and_filters == "Gaussian Blur": blur_fg = self.generate_gaussian_blur(t_fg, filter_size, device) return torch.lerp(t_fg, blur_fg, min(filter_strength, 1.0)) elif color_and_filters == "Edge Detect": sobel_x = torch.tensor([[[[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]]]], device=device).expand(t_fg.shape[1], 1, 3, 3) sobel_y = torch.tensor([[[[-1., -2., -1.], [0., 0., 0.], [1., 2., 1.]]]], device=device).expand(t_fg.shape[1], 1, 3, 3) edge_x = F.conv2d(t_fg, sobel_x, padding=1, groups=t_fg.shape[1]) edge_y = F.conv2d(t_fg, sobel_y, padding=1, groups=t_fg.shape[1]) edges = torch.clamp(torch.sqrt(edge_x**2 + edge_y**2), 0.0, 1.0) return torch.lerp(t_fg, edges, min(filter_strength, 1.0)) elif color_and_filters == "Sepia": sepia_w = torch.tensor([[0.393, 0.769, 0.189], [0.349, 0.686, 0.168], [0.272, 0.534, 0.131]], device=device) sepia_fg = torch.einsum('b c h w, d c -> b d h w', t_fg, sepia_w) return torch.lerp(t_fg, torch.clamp(sepia_fg, 0.0, 1.0), min(filter_strength, 1.0)) elif color_and_filters == "Invert": return torch.lerp(t_fg, 1.0 - t_fg, min(filter_strength, 1.0)) elif color_and_filters == "Auto-Color Match": if t_alpha is not None and is_real_bg: matched_fg = self.color_match(t_fg, t_bg, t_alpha) return torch.lerp(t_fg, matched_fg, min(filter_strength, 1.0)) return t_fg return t_fg def execute_upscale(t_img, t_is_mask=False, start_pct=0, end_pct=100): if not ai_upscale or ai_upscale_model == "None Found": pbar.update_absolute(end_pct) return t_img if not SPANDREL_AVAILABLE: print("[Quick Merge] Spandrel not available. Skipping upscale.") pbar.update_absolute(end_pct) return t_img chunk_b, orig_c, orig_h, orig_w = t_img.shape target_out_h = max(8, round(orig_h * upscale_by)) target_out_w = max(8, round(orig_w * upscale_by)) mm.soft_empty_cache() model_name_clean = ai_upscale_model.replace(" (Auto-Download)", "") if self.cached_upscaler_name != model_name_clean: model_path = folder_paths.get_full_path("upscale_models", model_name_clean) if not model_path: if model_name_clean == "RealESRGAN_x2plus.pth": dl_dir = os.path.join(folder_paths.models_dir, "upscale_models") os.makedirs(dl_dir, exist_ok=True) model_path = os.path.join(dl_dir, "RealESRGAN_x2plus.pth") try: print(f"[Quick Merge] Downloading RealESRGAN_x2plus.pth... Please wait, the UI may pause.") urllib.request.urlretrieve("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth", model_path) print(f"[Quick Merge] Download complete!") except Exception: pbar.update_absolute(end_pct) return t_img else: pbar.update_absolute(end_pct) return t_img try: sd = load_torch_file(model_path, safe_load=True) self.cached_upscaler = spandrel.ModelLoader().load_from_state_dict(sd).eval().to(device) self.cached_upscaler_name = model_name_clean except Exception: pbar.update_absolute(end_pct) return t_img tile = max(64, ai_upscale_tile_size) if ai_upscale_tile_size > 0 else 0 out_upscaled_slices = [] for i in range(chunk_b): frame_slice = t_img[i:i+1].to(device) slice_pct_s = start_pct + (i / chunk_b) * (end_pct - start_pct) slice_pct_e = start_pct + ((i + 1) / chunk_b) * (end_pct - start_pct) up_slice = self.tiled_upscale(self.cached_upscaler, frame_slice, tile_size=tile, pbar=pbar, start_pct=slice_pct_s, end_pct=slice_pct_e) if up_slice.shape[2] != target_out_h or up_slice.shape[3] != target_out_w: up_slice = F.interpolate(up_slice, size=(target_out_h, target_out_w), mode=interpolation, align_corners=False) out_upscaled_slices.append(up_slice.cpu()) del frame_slice, up_slice mm.soft_empty_cache() up_img = torch.cat(out_upscaled_slices, dim=0).to(device) if t_is_mask and post_upscale_edge_soften > 0: up_img = self.process_alpha(up_img, 0, post_upscale_edge_soften) return up_img def execute_composite(t_fg, t_alpha, t_bg, canvas_scale_x=1.0, canvas_scale_y=1.0, chunk_offset_start=0): chunk_b = t_bg.shape[0] canvas_h, canvas_w = t_bg.shape[2], t_bg.shape[3] result_img = t_bg.clone() final_alpha = torch.zeros((chunk_b, 1, canvas_h, canvas_w), device=device) relight_fg_out = torch.zeros_like(t_bg) relight_bg_out = t_bg.clone() pre_size_fg_w = max(8, int(t_fg.shape[3] * canvas_scale_x)) pre_size_fg_h = max(8, int(t_fg.shape[2] * canvas_scale_y)) t_fg = F.interpolate(t_fg, size=(pre_size_fg_h, pre_size_fg_w), mode=interpolation, align_corners=False) t_alpha = F.interpolate(t_alpha, size=(pre_size_fg_h, pre_size_fg_w), mode=interpolation, align_corners=False) target_fg_w = max(8, int(pre_size_fg_w * foreground_size)) target_fg_h = max(8, int(pre_size_fg_h * foreground_size)) t_fg = F.interpolate(t_fg, size=(target_fg_h, target_fg_w), mode=interpolation, align_corners=False) t_alpha = F.interpolate(t_alpha, size=(target_fg_h, target_fg_w), mode=interpolation, align_corners=False) center_x = (canvas_w // 2) - (target_fg_w // 2) center_y = (canvas_h // 2) - (target_fg_h // 2) scaled_lr = int(round(foreground_left_right * canvas_scale_x)) scaled_ud = int(round(foreground_up_down * canvas_scale_y)) final_x = center_x + scaled_lr final_y = center_y + scaled_ud x1, y1 = max(0, final_x), max(0, final_y) x2, y2 = min(canvas_w, final_x + target_fg_w), min(canvas_h, final_y + target_fg_h) fg_x1, fg_y1 = max(0, -final_x), max(0, -final_y) fg_x2, fg_y2 = fg_x1 + (x2 - x1), fg_y1 + (y2 - y1) for b in range(chunk_b): if (x2 > x1) and (y2 > y1): crop_fg = t_fg[b:b+1, :, fg_y1:fg_y2, fg_x1:fg_x2] crop_alpha = t_alpha[b:b+1, :, fg_y1:fg_y2, fg_x1:fg_x2] bg_slice = result_img[b:b+1, :, y1:y2, x1:x2] relight_fg_out[b:b+1, :, y1:y2, x1:x2] = crop_fg * crop_alpha active_alpha = crop_alpha * foreground_opacity if shadows and shadow_darkness > 0.0: s_blur = shadow_softness if shadow_softness % 2 != 0 else shadow_softness + 1 full_shadow_mask = F.avg_pool2d(t_alpha[b:b+1] * foreground_opacity, kernel_size=s_blur, stride=1, padding=s_blur // 2) scaled_spx = int(round(shadow_position_x * canvas_scale_x)) scaled_spy = int(round(shadow_position_y * canvas_scale_y)) sx1, sy1 = max(0, final_x + scaled_spx), max(0, final_y + scaled_spy) sx2, sy2 = min(canvas_w, final_x + target_fg_w + scaled_spx), min(canvas_h, final_y + target_fg_h + scaled_spy) sfg_x1, sfg_y1 = max(0, -(final_x + scaled_spx)), max(0, -(final_y + scaled_spy)) sfg_x2, sfg_y2 = sfg_x1 + (sx2 - sx1), sfg_y1 + (sy2 - sy1) if (sx2 > sx1) and (sy2 > sy1): crop_smask = full_shadow_mask[:, :, sfg_y1:sfg_y2, sfg_x1:sfg_x2] s_color = torch.tensor([shadow_R, shadow_G, shadow_B], device=device).view(1, 3, 1, 1) s_color = s_color.expand_as(result_img[b:b+1, :, sy1:sy2, sx1:sx2]) result_img[b:b+1, :, sy1:sy2, sx1:sx2] = torch.lerp( result_img[b:b+1, :, sy1:sy2, sx1:sx2], s_color, crop_smask * shadow_darkness) if blend_mode == "Multiply": blend_math = bg_slice * crop_fg elif blend_mode == "Screen": blend_math = 1.0 - (1.0 - bg_slice) * (1.0 - crop_fg) elif blend_mode == "Overlay": mask_lt = (bg_slice < 0.5).float() blend_math = mask_lt * (2.0 * bg_slice * crop_fg) + (1.0 - mask_lt) * (1.0 - 2.0 * (1.0 - bg_slice) * (1.0 - crop_fg)) elif blend_mode == "Soft Light": blend_math = (1.0 - 2.0 * crop_fg) * (bg_slice ** 2) + 2.0 * crop_fg * bg_slice else: blend_math = crop_fg active_fg = torch.lerp(crop_fg, blend_math, blend_strength) result_img[b:b+1, :, y1:y2, x1:x2] = torch.lerp(bg_slice, active_fg, active_alpha) final_alpha[b:b+1, :, y1:y2, x1:x2] = active_alpha return result_img, final_alpha, relight_fg_out, relight_bg_out # ========================================== # PATH A: BACKGROUND-ONLY BYPASS # ========================================== if not merge_image: out_images = [] for c_idx in range(num_chunks): c_start = c_idx * effective_chunk c_end = min(c_start + effective_chunk, b_size) c_bg = bg[c_start:c_end].clone().to(device) if colors_and_filters and filter_target in ["Background", "Both"]: c_bg = execute_filters(c_bg, None, c_bg) if ai_upscale: pct_s = 50 + int(c_start / b_size * 45) pct_e = 50 + int(c_end / b_size * 45) c_bg = execute_upscale(c_bg, False, start_pct=pct_s, end_pct=pct_e) if target_resolution: c_bg = self.resize_tensor(c_bg, target_width, target_height, target_resize_mode, target_crop_align, interpolation) out_images.append(c_bg.cpu() if multi_chunk else c_bg) del c_bg if multi_chunk: mm.soft_empty_cache() pbar.update_absolute(100) bg_final = torch.cat(out_images, dim=0) final_h, final_w = bg_final.shape[2], bg_final.shape[3] if has_mask: out_light_mask = self.generate_smart_light_gradient(m_tensor, final_h, final_w, device) else: out_light_mask = torch.zeros((b_size, final_h, final_w), dtype=torch.float32).cpu() empty_fg_tensor = torch.zeros((b_size, final_h, final_w, 3), dtype=torch.float32).cpu() return (bg_final.permute(0, 2, 3, 1).cpu(), out_light_mask, empty_fg_tensor, bg_final.permute(0, 2, 3, 1).cpu()) # ========================================== # PATH B: COMPOSITING CHUNK LOOPS # ========================================== all_fi, all_fm, all_rf, all_rb = [], [], [], [] for c_idx in range(num_chunks): c_start = c_idx * effective_chunk c_end = min(c_start + effective_chunk, b_size) pct_lo = 30 + int( c_idx / num_chunks * 70) pct_hi = 30 + int((c_idx + 1) / num_chunks * 70) c_bg = bg[c_start:c_end].clone().to(device) c_fg = fg[c_start:c_end].clone().to(device) c_alpha = alpha[c_start:c_end].clone().to(device) if target_resolution and pipeline_order != "upscale composite": c_bg = self.resize_tensor(c_bg, target_width, target_height, target_resize_mode, target_crop_align, interpolation) if colors_and_filters and background_blur > 0: c_bg = self.generate_gaussian_blur(c_bg, background_blur, device) if colors_and_filters and filter_target in ["Background", "Both"]: c_bg = execute_filters(c_bg, None, c_bg) if colors_and_filters and filter_target in ["Foreground", "Both"]: c_fg = execute_filters(c_fg, c_alpha, c_bg) _sx = c_bg.shape[3] / pre_transform_bg_w _sy = c_bg.shape[2] / pre_transform_bg_h if pipeline_order == "upscale foreground only": mid = pct_lo + (pct_hi - pct_lo) // 3 c_fg = execute_upscale(c_fg, False, start_pct=pct_lo, end_pct=mid) c_alpha = execute_upscale(c_alpha, True, start_pct=mid, end_pct=pct_lo + (pct_hi - pct_lo) * 2 // 3) c_fi, c_fm, c_rf, c_rb = execute_composite(c_fg, c_alpha, c_bg, _sx, _sy, c_start) elif pipeline_order == "upscale background only": c_bg = execute_upscale(c_bg, False, start_pct=pct_lo, end_pct=pct_lo + (pct_hi - pct_lo) * 2 // 3) _sx = c_bg.shape[3] / pre_transform_bg_w _sy = c_bg.shape[2] / pre_transform_bg_h c_fi, c_fm, c_rf, c_rb = execute_composite(c_fg, c_alpha, c_bg, _sx, _sy, c_start) else: # upscale composite c_fi, c_fm, c_rf, c_rb = execute_composite(c_fg, c_alpha, c_bg, _sx, _sy, c_start) mid = pct_lo + (pct_hi - pct_lo) // 2 c_fi = execute_upscale(c_fi, False, start_pct=pct_lo, end_pct=mid) c_fm = execute_upscale(c_fm, True, start_pct=mid, end_pct=pct_hi) up_h, up_w = c_fi.shape[2], c_fi.shape[3] c_rf = F.interpolate(c_rf, size=(up_h, up_w), mode=interpolation, align_corners=False) c_rb = F.interpolate(c_rb, size=(up_h, up_w), mode=interpolation, align_corners=False) if target_resolution: c_fi = self.resize_tensor(c_fi, target_width, target_height, target_resize_mode, target_crop_align, interpolation) c_fm = self.resize_tensor(c_fm, target_width, target_height, target_resize_mode, target_crop_align, interpolation) c_rf = self.resize_tensor(c_rf, target_width, target_height, target_resize_mode, target_crop_align, interpolation) c_rb = self.resize_tensor(c_rb, target_width, target_height, target_resize_mode, target_crop_align, interpolation) all_fi.append(c_fi.cpu() if multi_chunk else c_fi) all_fm.append(c_fm.cpu() if multi_chunk else c_fm) all_rf.append(c_rf.cpu() if multi_chunk else c_rf) all_rb.append(c_rb.cpu() if multi_chunk else c_rb) del c_bg, c_fg, c_alpha, c_fi, c_fm, c_rf, c_rb if multi_chunk: mm.soft_empty_cache() pbar.update_absolute(100) final_img = torch.cat(all_fi, dim=0) final_mask = torch.cat(all_fm, dim=0) relight_fg = torch.cat(all_rf, dim=0) relight_bg = torch.cat(all_rb, dim=0) if multi_chunk: final_img = final_img.to(device) final_mask = final_mask.to(device) relight_fg = relight_fg.to(device) relight_bg = relight_bg.to(device) # ========================================== # LIGHT MASK & OPTION A ALPHA GENERATION # ========================================== final_h, final_w = final_img.shape[2], final_img.shape[3] if has_mask: m_resized = F.interpolate(m_tensor.float(), size=(final_h, final_w), mode='bilinear', align_corners=False) out_light_mask = self.generate_gaussian_blur(m_resized.to(device), 51, device).squeeze(1).cpu() elif has_fg and not is_real_bg: out_light_mask = final_mask.squeeze(1).cpu() else: out_light_mask = torch.zeros((b_size, final_h, final_w), dtype=torch.float32).cpu() out_images = final_img.permute(0, 2, 3, 1).cpu() out_relight_fg = relight_fg.permute(0, 2, 3, 1).cpu() out_relight_bg = relight_bg.permute(0, 2, 3, 1).cpu() ui_msg = "[Quick Merge Engine] Auto-Position Disabled (WiP Phase Lock)" return { "ui": {"text": [ui_msg]}, "result": (out_images, out_light_mask, out_relight_fg, out_relight_bg) } NODE_CLASS_MAPPINGS = { "QuickMergeNode": QuickMergeNode } NODE_DISPLAY_NAME_MAPPINGS = { "QuickMergeNode": "🎨 Quick Merge" }