diff --git a/__pycache__/__init__.cpython-311.pyc b/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..a17a5ff Binary files /dev/null and b/__pycache__/__init__.cpython-311.pyc differ diff --git a/__pycache__/quick_merge_node.cpython-311.pyc b/__pycache__/quick_merge_node.cpython-311.pyc new file mode 100644 index 0000000..a281674 Binary files /dev/null and b/__pycache__/quick_merge_node.cpython-311.pyc differ diff --git a/pyproject.toml b/pyproject.toml index f9fc7e1..c2ebe36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ComfyUI-Quick-Merge" -version = "1.0.1" +version = "1.0.2" description = "High-performance modular layer blending and folder batch processing for ComfyUI." license = { text = "MIT License" } authors = [ diff --git a/quick_merge_node.py b/quick_merge_node.py index 1de4f58..dd66014 100644 --- a/quick_merge_node.py +++ b/quick_merge_node.py @@ -8,7 +8,12 @@ import numpy as np import folder_paths import comfy.model_management as mm from comfy.utils import load_torch_file, ProgressBar -import spandrel + +try: + import spandrel + SPANDREL_AVAILABLE = True +except ImportError: + SPANDREL_AVAILABLE = False try: from transformers import AutoModelForImageSegmentation @@ -24,6 +29,7 @@ try: except ImportError: ONNXRUNTIME_AVAILABLE = False + class QuickMergeNode: def __init__(self): self.cached_bg_model = None @@ -391,6 +397,29 @@ class QuickMergeNode: 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 @@ -443,8 +472,8 @@ class QuickMergeNode: 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", "") - bg_batch_folder = kwargs.get("📁 bg_batch_folder", "") + 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) @@ -510,7 +539,7 @@ class QuickMergeNode: pixel_smoothing = kwargs.get("🖼️ pixel_smoothing", "bicubic") interpolation = pixel_smoothing if pixel_smoothing != "lanczos" else "bicubic" - # --- 1. DIRECTORY CONDITIONAL PROCESSING OVERRIDES --- + # --- 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) @@ -520,10 +549,14 @@ class QuickMergeNode: if bg_batch_folder == default_folder_text: bg_batch_folder = "" if batch_processing: - if background_batch and bg_batch_folder and os.path.isdir(bg_batch_folder): - bg_input = self.load_folder_images(bg_batch_folder, device, True) - if foreground_batch and fg_batch_folder and os.path.isdir(fg_batch_folder): - fg_input = self.load_folder_images(fg_batch_folder, device, True) + 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 @@ -670,6 +703,11 @@ class QuickMergeNode: 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)) @@ -686,7 +724,9 @@ class QuickMergeNode: 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