""" Koolook Matte — ComfyUI nodes for mask-guided one-step video matting. Part of the Koolook pack (see matte/NOTICE for the model-weight licenses). It pairs an independent implementation of the published VideoMaMa method (pipeline.py) with: * temporal windowing (peak VRAM scales with the window, not the clip length), * subject-focused ROI cropping + paste-back (process only the subject, native res), * shortest-side / longest-axis resolution control, /64-aligned, with auto OOM resolution-backoff that preserves the temporal window (no new flicker). It loads the VideoMaMa UNet + SVD VAE *weights* at runtime; those weights carry their own (non-commercial) licenses and are the user's responsibility. """ import os import numpy as np import torch from PIL import Image from pathlib import Path from comfy.utils import ProgressBar from .pipeline import KoolookMattePipeline def _ensure_local_model(path: str, repo_id: str, label: str) -> str: """Return a usable local model dir. If `path` is missing/empty, try to fetch `repo_id` from the Hugging Face Hub (optional dep). Weights are never bundled.""" if os.path.isdir(path) and os.listdir(path): return path try: from huggingface_hub import snapshot_download except Exception: raise RuntimeError( f"{label} not found at {path} and huggingface_hub isn't installed. " f"Place the model there, or `pip install huggingface_hub`." ) print(f"[KoolookMatte] downloading {label} ({repo_id}) -> {path}") os.makedirs(path, exist_ok=True) snapshot_download(repo_id=repo_id, local_dir=path) return path class KoolookMatteLoader: """ Loads the VideoMaMa pipeline with base model and fine-tuned UNet checkpoint. """ @classmethod def INPUT_TYPES(cls): return { "required": { "base_model_path": ("STRING", { "default": "checkpoints/stabilityai/stable-video-diffusion-img2vid-xt", "multiline": False }), "unet_checkpoint_path": ("STRING", { "default": "checkpoints/VideoMaMa", "multiline": False }), "precision": (["fp16", "bf16"], { "default": "fp16" }), "enable_model_cpu_offload": ("BOOLEAN", { "default": True }), "vae_encode_chunk_size": ("INT", { "default": 1, "min": 1, "max": 25, "step": 1, "tooltip": "Frames the VAE encodes/decodes at once. 1 = desktop (Sammie) default and lowest VRAM; raising it makes the VAE decode spike proportionally bigger and is the usual cause of OOM at high resolution / 4K. Leave at 1 unless you have VRAM to spare and want a small speedup." }), "attention_mode": (["auto", "xformers", "sdpa", "none"], { "default": "auto" }), "enable_vae_tiling": ("BOOLEAN", { "default": False }), "enable_vae_slicing": ("BOOLEAN", { "default": True }), } } RETURN_TYPES = ("KOOLOOK_MATTE_PIPELINE",) FUNCTION = "load_pipeline" CATEGORY = "Koolook Matte" def load_pipeline( self, base_model_path: str, unet_checkpoint_path: str, precision: str, enable_model_cpu_offload: bool, vae_encode_chunk_size: int, attention_mode: str, enable_vae_tiling: bool, enable_vae_slicing: bool ): # fp16 is the default; only an explicit bf16 opts out. Anything else # (incl. a stale "auto" widget from an older node version) -> fp16. weight_dtype = torch.bfloat16 if precision == "bf16" else torch.float16 device = "cuda" if torch.cuda.is_available() else "cpu" if device == "cpu": # Half precision is unreliable/unsupported for many ops on CPU, and # SVD effectively needs a GPU anyway. Force float32 so a CPU fallback # at least runs correctly (slowly) instead of erroring on fp16 kernels. print("[KoolookMatte] CUDA unavailable -> CPU/float32 (very slow; a " "GPU is strongly recommended).") weight_dtype = torch.float32 node_dir = Path(__file__).parent # relative paths resolve next to this pack; absolute paths pass through base = base_model_path if os.path.isabs(base_model_path) else str(node_dir / base_model_path) unet = unet_checkpoint_path if os.path.isabs(unet_checkpoint_path) else str(node_dir / unet_checkpoint_path) base = _ensure_local_model(base, "stabilityai/stable-video-diffusion-img2vid-xt", "SVD base model") unet = _ensure_local_model(unet, "SammyLim/VideoMaMa", "VideoMaMa UNet") print(f"[KoolookMatte] loading pipeline | base={base} unet={unet} | {precision} " f"offload={enable_model_cpu_offload} vae_chunk={vae_encode_chunk_size} attn={attention_mode}") try: pipeline = KoolookMattePipeline( base_model_path=base, unet_checkpoint_path=unet, weight_dtype=weight_dtype, device=device, enable_model_cpu_offload=enable_model_cpu_offload, vae_encode_chunk_size=vae_encode_chunk_size, attention_mode=attention_mode, enable_vae_tiling=enable_vae_tiling, enable_vae_slicing=enable_vae_slicing, ) print(f"[KoolookMatte] pipeline loaded ({precision})") return (pipeline,) except Exception as e: raise RuntimeError(f"Failed to load Koolook Matte pipeline: {e}") def _ensure_2d_mask(mask_np: np.ndarray) -> np.ndarray: """Reduce a mask array to 2D (H, W) for PIL grayscale mode. Handles the standard ComfyUI ``[H, W]`` plus the extra-dim layouts some sources emit (e.g. MatAnyone ``[C, H, W]`` / ``[H, W, C]``). Singleton dims are squeezed first; a genuine multi-channel mask has identical channels, so we collapse to the first. Anything that can't be unambiguously reduced raises rather than silently returning the wrong geometry (e.g. the old code turned ``[H, W, 3]`` into ``[W, 3]``). """ mask_np = np.squeeze(mask_np) # [1,H,W] / [H,W,1] / [1,1,H,W] -> ... if mask_np.ndim == 2: return mask_np if mask_np.ndim == 3: if mask_np.shape[-1] in (3, 4): # [H, W, C] (channels last) return mask_np[:, :, 0] if mask_np.shape[0] in (3, 4): # [C, H, W] (channels first) return mask_np[0] raise ValueError( f"_ensure_2d_mask: cannot reduce mask of shape {mask_np.shape} to 2D; " "expected [H,W], [H,W,C], [C,H,W] or singleton-padded variants." ) class KoolookMatteSampler: """ Runs one-step mask-guided matting on video frames with mask conditioning. Expects ComfyUI IMAGE format: [B, H, W, C] tensors with values in [0, 1]. """ @classmethod def INPUT_TYPES(cls): return { "required": { "pipeline": ("KOOLOOK_MATTE_PIPELINE",), "images": ("IMAGE",), "masks": ("MASK",), "seed": ("INT", { "default": 42, "min": 0, "max": 0xffffffffffffffff }), "max_resolution": ("INT", { "default": 1024, "min": 256, "max": 2048, "step": 8 }), "fps": ("INT", { "default": 7, "min": 1, "max": 60 }), "motion_bucket_id": ("INT", { "default": 127, "min": 1, "max": 255 }), "noise_aug_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01 }), "window_size": ("INT", { "default": 16, "min": 0, "max": 128, "step": 1, "tooltip": "Frames sent through the UNet per pass. Lower = far less VRAM, so you can raise max_resolution (e.g. 1440) on 24GB. 0 = process the whole clip in one pass (legacy behaviour)." }), "window_overlap": ("INT", { "default": 2, "min": 0, "max": 32, "step": 1, "tooltip": "Frames re-processed and blended at each window boundary for temporal continuity. 0 disables overlap feedback/blending. Must be < window_size." }), "crop_to_mask": ("BOOLEAN", { "default": True, "tooltip": "Process only the union bounding box of the masks (Sammie-style ROI). Far fewer pixels through the UNet => much faster and lets max_resolution apply to just the subject. The output matte is pasted back into the full frame (zero alpha outside)." }), "crop_padding": ("INT", { "default": 16, "min": 0, "max": 512, "step": 8, "tooltip": "Pixels of margin added around the mask bounding box before cropping, so soft edges/motion aren't clipped." }), "resize_mode": (["shortest_side", "longest_axis"], { "default": "shortest_side", "tooltip": "shortest_side = desktop/Sammie parity: cap the shorter side at max_resolution, never upscale, keep the long axis at full detail (sharpest matte, more VRAM). longest_axis = force the longer side to max_resolution (lower VRAM, softer edges)." }), } } RETURN_TYPES = ("MASK",) FUNCTION = "run_inference" CATEGORY = "Koolook Matte" # SVD's UNet downsamples the latent 8x, so pixel dims should divide by 64. # Anything only /8 gets padded/cropped inside the UNet bottleneck, which softens # edges — so we align the model working size to 64. ALIGN = 64 @classmethod def _compute_target_size(cls, width: int, height: int, max_resolution: int, mode: str = "shortest_side"): """Compute the model working size, aligned to a multiple of 64 (SVD-friendly). mode="shortest_side" (desktop / Sammie parity): cap the SHORTER side at max_resolution and NEVER upscale. A crop smaller than max_resolution is processed at native pixels, so the long axis keeps full detail — this is what makes the desktop app's edges sharp. Uses more VRAM than longest_axis for the same number (the long side is unbounded). mode="longest_axis": force the LONGER side to exactly max_resolution (upscaling small crops). Lower VRAM, but downscales the detail axis => softer. """ a = cls.ALIGN if max_resolution <= 0: return max((width // a) * a, a), max((height // a) * a, a) if mode == "shortest_side": min_side = min(width, height) scale = (max_resolution / min_side) if min_side > max_resolution else 1.0 new_width = int(width * scale) new_height = int(height * scale) else: # longest_axis if width >= height: new_width = max_resolution new_height = int(height * max_resolution / width) else: new_height = max_resolution new_width = int(width * max_resolution / height) # Align to multiple of 64 (SVD UNet requirement — avoids internal padding softness) new_width = max((new_width // a) * a, a) new_height = max((new_height // a) * a, a) return new_width, new_height @staticmethod def _generate_windows(total_frames: int, window_size: int, overlap: int): """Generate sliding-window (start, end) pairs covering total_frames (end exclusive). window_size <= 0 or >= total_frames processes the whole clip in one pass. overlap that is >= window_size is clamped so the window always advances. Standard sliding-window scheme (independent implementation). """ if window_size <= 0 or window_size >= total_frames: return [(0, total_frames)] step = window_size - overlap if step <= 0: print(f"[KoolookMatte] window_overlap ({overlap}) >= window_size ({window_size}); " f"clamping so windows advance by 1 frame") step = 1 windows = [] pos = 0 while pos < total_frames: end = min(pos + window_size, total_frames) windows.append((pos, end)) if end >= total_frames: break pos += step return windows @staticmethod def _to_gray01(frame_pil) -> np.ndarray: """Convert a pipeline output frame (PIL, RGB or L) to float32 grayscale in [0, 1].""" arr = np.array(frame_pil).astype(np.float32) / 255.0 if arr.ndim == 3: arr = arr.mean(axis=-1) return arr @staticmethod def _compute_union_bbox(masks_full, padding: int, orig_w: int, orig_h: int, threshold: int = 127): """Union bounding box (x1, y1, x2, y2; x2/y2 exclusive) over every mask in the clip, padded and clamped to the frame. `masks_full` is a list of 2D uint8 arrays already resampled to the image resolution (orig_h, orig_w), so bbox, image crop and mask crop all share one coordinate space. Returns None if no mask is present (caller then processes the full frame). Mirrors Sammie-Roto-2's compute_mask_bounding_box: one rect for the entire range, so every window/frame uses an identical spatial region. """ union = None for a in masks_full: b = a > threshold union = b if union is None else (union | b) if union is None or not union.any(): print("[KoolookMatte] guide mask is EMPTY across all frames -> " "falling back to full-frame. If unexpected, check the " "segmentation (a black matte usually means an empty guide).") return None rows = np.where(union.any(axis=1))[0] cols = np.where(union.any(axis=0))[0] y1, y2 = int(rows[0]), int(rows[-1]) + 1 x1, x2 = int(cols[0]), int(cols[-1]) + 1 # Pad and clamp to frame x1 = max(0, x1 - padding) y1 = max(0, y1 - padding) x2 = min(orig_w, x2 + padding) y2 = min(orig_h, y2 + padding) if x2 - x1 < 8 or y2 - y1 < 8: print(f"[KoolookMatte] guide-mask region is degenerate " f"({x2 - x1}x{y2 - y1}px) -> falling back to full-frame. " "Check the segmentation if unexpected.") return None # degenerate region; process full frame return (x1, y1, x2, y2) def run_inference( self, pipeline, images, masks, seed: int, max_resolution: int, fps: int, motion_bucket_id: int, noise_aug_strength: float, window_size: int = 0, window_overlap: int = 0, crop_to_mask: bool = False, crop_padding: int = 16, resize_mode: str = "shortest_side" ): """Run VideoMaMa inference, optionally split into temporal windows to cap VRAM. With window_size > 0 and < num_frames the clip is processed in sliding windows of `window_size` frames. The trailing `window_overlap` soft-alpha outputs of each window are fed back as the leading masks of the next window (temporal continuity), and the overlap region is linearly blended across the seam. Peak VRAM then scales with window_size instead of the full clip length, which is what lets max_resolution reach 1440 on a 24GB card. window_size=0 (or >= num_frames) reproduces the original single-pass behaviour exactly. """ import time num_frames = images.shape[0] if masks.shape[0] != num_frames: raise ValueError( f"Number of image frames ({num_frames}) must match " f"number of mask frames ({masks.shape[0]})" ) orig_h, orig_w = images.shape[1], images.shape[2] # --- Evaluation instrumentation: wall-clock + peak VRAM --- t_start = time.time() if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() # Bring every mask up to the IMAGE resolution first (nearest — it's only a # coarse guide). This lets you feed the node the FULL-RES plate while the mask # keeps coming from the cheap downscaled SAM2/RMBG path: image, bbox and mask # crop then all share one coordinate space, so no detail is thrown away before # the model sees it (this is how the desktop app stays crisp). masks_full = [] for i in range(num_frames): a = _ensure_2d_mask((masks[i].cpu().numpy() * 255).astype(np.uint8)) if a.shape[:2] != (orig_h, orig_w): a = np.array(Image.fromarray(a, mode='L').resize((orig_w, orig_h), Image.NEAREST)) masks_full.append(a) # --- Optional ROI crop: process only the subject's union bounding box --- crop_rect = self._compute_union_bbox(masks_full, crop_padding, orig_w, orig_h) if crop_to_mask else None if crop_rect is not None: cx1, cy1, cx2, cy2 = crop_rect src_w, src_h = cx2 - cx1, cy2 - cy1 else: cx1, cy1, cx2, cy2 = 0, 0, orig_w, orig_h src_w, src_h = orig_w, orig_h # Model working size for the processed region (see _compute_target_size). target_w, target_h = self._compute_target_size(src_w, src_h, max_resolution, resize_mode) # Build the model input frames at a given working size. Called once up front, # and again smaller if we have to back off on a CUDA OOM. Crops to the ROI here; # masks are binarized (crisp guide) and resized nearest. def _build_frames(tw, th): cond, msk = [], [] for i in range(num_frames): img_np = (images[i].cpu().numpy() * 255).astype(np.uint8) if crop_rect is not None: img_np = img_np[cy1:cy2, cx1:cx2] cond.append(Image.fromarray(img_np, mode='RGB').resize((tw, th), Image.LANCZOS)) m = masks_full[i] if crop_rect is not None: m = m[cy1:cy2, cx1:cx2] m = np.where(m > 127, np.uint8(255), np.uint8(0)) msk.append(Image.fromarray(m, mode='L').resize((tw, th), Image.NEAREST)) return cond, msk try: # Run the windowed inference. On a CUDA OOM we back off on RESOLUTION, NOT # window_size — the temporal window stays full so mattes don't start # flickering. A sharp lower-res matte beats a temporally-inconsistent one. # Catch RuntimeError broadly (torch.cuda.OutOfMemoryError subclasses # it, but some OOMs — cuBLAS/cuDNN/older paths — raise a plain # RuntimeError: CUDA out of memory). The message guard below filters # out any non-OOM RuntimeError so we never mask a real bug. oom_types = (RuntimeError,) min_side_floor = 256 while True: cond_frames, mask_frames = _build_frames(target_w, target_h) print(f"Input {orig_w}x{orig_h} | region {src_w}x{src_h}" f"{' (cropped to mask)' if crop_rect else ''} -> target {target_w}x{target_h} " f"(max_resolution={max_resolution}, {resize_mode})") try: output_accum = self._run_windows( pipeline, cond_frames, mask_frames, num_frames, window_size, window_overlap, seed, fps, motion_bucket_id, noise_aug_strength ) break except oom_types as oom: if "out of memory" not in str(oom).lower(): raise # a genuine (non-OOM) error — don't mask or retry it if torch.cuda.is_available(): torch.cuda.empty_cache() if min(target_w, target_h) <= min_side_floor: raise new_w = max((int(target_w * 0.85) // self.ALIGN) * self.ALIGN, self.ALIGN) new_h = max((int(target_h * 0.85) // self.ALIGN) * self.ALIGN, self.ALIGN) print(f"[KoolookMatte] CUDA OOM at {target_w}x{target_h}; cleared cache, " f"backing off resolution to {new_w}x{new_h} (window_size kept at " f"{window_size} for temporal stability) and retrying.") target_w, target_h = new_w, new_h # Resize each matte back to the region size and paste into the full # frame canvas (zero alpha outside the crop), then stack -> ComfyUI MASK output_masks = [] for idx in range(num_frames): g = output_accum[idx] if g is None: g = np.zeros((target_h, target_w), dtype=np.float32) g_pil = Image.fromarray((g * 255.0).clip(0, 255).astype(np.uint8), mode='L') g_pil = g_pil.resize((src_w, src_h), Image.LANCZOS) region = np.array(g_pil).astype(np.float32) / 255.0 if crop_rect is not None: full = np.zeros((orig_h, orig_w), dtype=np.float32) full[cy1:cy2, cx1:cx2] = region output_masks.append(full) else: output_masks.append(region) output_tensor = torch.from_numpy(np.stack(output_masks, axis=0)) elapsed = time.time() - t_start per_frame = elapsed / max(num_frames, 1) vram_str = "" if torch.cuda.is_available(): vram_str = f", peak VRAM {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB" print(f"KoolookMatte done: {output_tensor.shape} in {elapsed:.1f}s " f"({per_frame:.2f}s/frame){vram_str}") return (output_tensor,) except Exception as e: raise RuntimeError(f"Koolook Matte inference failed: {e}") def _run_windows(self, pipeline, cond_frames, mask_frames, num_frames, window_size, window_overlap, seed, fps, motion_bucket_id, noise_aug_strength): """Run the model over the clip in temporal windows and return a per-frame list of target-resolution grayscale float32 [0,1] mattes. Raises the CUDA OutOfMemoryError from pipeline.run unchanged so the caller can catch it and retry with a smaller window_size. """ windows = self._generate_windows(num_frames, window_size, window_overlap) single_pass = (len(windows) == 1) # Effective overlap (0 for a single pass, clamped so windows always advance) overlap = 0 if single_pass else min(window_overlap, window_size - 1) if single_pass: print(f"Running Koolook Matte inference on {num_frames} frames (single pass)...") else: print(f"Running Koolook Matte inference on {num_frames} frames in " f"{len(windows)} windows of up to {window_size} frames (overlap {overlap})...") # pipeline.run issues 4 progress updates per pass pbar = ProgressBar(4 * len(windows)) # output_accum[i] holds a target-resolution grayscale float32 [0,1] array output_accum = [None] * num_frames prev_overlap_soft = None # previous window's last `overlap` outputs, as PIL 'L' for w_idx, (w_start, w_end) in enumerate(windows): is_first = (w_idx == 0) win_cond = cond_frames[w_start:w_end] win_mask = list(mask_frames[w_start:w_end]) # Feedback injection: replace the leading `overlap` conditioning masks # with the previous window's soft alpha output. Kept soft (not # binarized) so the model sees graduated edge values across the seam. if not is_first and overlap > 0 and prev_overlap_soft is not None: for i in range(min(overlap, len(win_mask), len(prev_overlap_soft))): win_mask[i] = prev_overlap_soft[i] out_pil = pipeline.run( cond_frames=win_cond, mask_frames=win_mask, seed=seed, fps=fps, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, pbar=pbar ) out_gray = [self._to_gray01(f) for f in out_pil] # Capture soft tail for the next window's feedback loop if overlap > 0 and len(out_gray) >= overlap: prev_overlap_soft = [ Image.fromarray((g * 255.0).clip(0, 255).astype(np.uint8), mode='L') for g in out_gray[-overlap:] ] else: prev_overlap_soft = None # Commit outputs. On non-first windows the leading `overlap` frames # already have a committed prediction from the previous window, so we # linearly blend the two across the seam (weight ramps 0 -> 1). for i, g in enumerate(out_gray): abs_idx = w_start + i if is_first or i >= overlap: output_accum[abs_idx] = g elif output_accum[abs_idx] is not None: blend_w = (i + 1) / (overlap + 1) output_accum[abs_idx] = (1.0 - blend_w) * output_accum[abs_idx] + blend_w * g else: output_accum[abs_idx] = g return output_accum class KoolookMatteFocusCrop: """Crop a video batch to a single static 'hero' box around the subject. The box is the union of `coarse_mask` across every frame (plus padding), so the crop is identical for the whole clip — a subject moving inside a locked frame, NOT a per-frame match-move (which would break SAM3 / VideoMaMa temporal tracking). Run SAM3 and VideoMaMa on the crop at the subject's native pixels, then paste the result back with KoolookMatteFocusStitch. This is the "focused comp" pattern: crop once, do all the work in that high-res space, stitch back at the end. """ @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE",), "coarse_mask": ("MASK",), "padding": ("INT", { "default": 24, "min": 0, "max": 1024, "step": 8, "tooltip": "Pixels of margin around the subject's union bounding box. Give thin appendages (claws, tails) room so they fall INSIDE the crop, otherwise they can't be recovered downstream." }), "max_long_side": ("INT", { "default": 2560, "min": 0, "max": 8192, "step": 64, "tooltip": "VRAM safety cap: if the native crop's LONGEST side exceeds this, downscale the crop to fit (aspect kept). Never upscales. This same cropped image feeds BOTH the segmentation and the matte model, so it bounds the area both process. 0 = no cap (native crop). 2560 comfortably fits a 24GB card." }), "enabled": ("BOOLEAN", { "default": True, "tooltip": "Master toggle for the focused two-pass. ON = crop to the subject, process at native res, stitch back. OFF = pass the FULL frame straight through unchanged (identity focus), so the downstream seg/matte run single-pass full-frame. Promote this to a subgraph to get a one-click two-pass switch." }), } } RETURN_TYPES = ("IMAGE", "MASK", "FOCUS") RETURN_NAMES = ("cropped_images", "cropped_mask", "focus") FUNCTION = "crop" CATEGORY = "Koolook Matte" def crop(self, images, coarse_mask, enabled, padding, max_long_side): num = images.shape[0] if coarse_mask.shape[0] != num: # Mismatched counts would make cropped_images (num) and cropped_mask # (coarse_mask.shape[0]) diverge, silently dropping/misaligning frames # at FocusStitch. Fail like KoolookMatteSampler.run_inference does. raise ValueError( f"Number of image frames ({num}) must match number of " f"coarse_mask frames ({coarse_mask.shape[0]})" ) orig_h, orig_w = images.shape[1], images.shape[2] # --- Bypass: pass the full frame through unchanged with an identity focus, so # FocusStitch (which just pastes into the bbox) becomes a no-op passthrough. if not enabled: cm = [] for i in range(coarse_mask.shape[0]): a = _ensure_2d_mask((coarse_mask[i].cpu().numpy() * 255).astype(np.uint8)) if a.shape[:2] != (orig_h, orig_w): a = np.array(Image.fromarray(a, mode='L').resize((orig_w, orig_h), Image.NEAREST)) cm.append(a.astype(np.float32) / 255.0) cropped_mask = torch.from_numpy(np.stack(cm, axis=0)) if cm else torch.zeros((num, orig_h, orig_w)) focus = {"bbox": (0, 0, orig_w, orig_h), "orig_w": orig_w, "orig_h": orig_h, "crop_w": orig_w, "crop_h": orig_h} print(f"[FocusCrop] disabled -> full-frame passthrough ({orig_w}x{orig_h})") return (images, cropped_mask, focus) # Bring the coarse mask to image resolution so the bbox is in image space. masks_full = [] for i in range(coarse_mask.shape[0]): a = _ensure_2d_mask((coarse_mask[i].cpu().numpy() * 255).astype(np.uint8)) if a.shape[:2] != (orig_h, orig_w): a = np.array(Image.fromarray(a, mode='L').resize((orig_w, orig_h), Image.NEAREST)) masks_full.append(a) bbox = KoolookMatteSampler._compute_union_bbox(masks_full, padding, orig_w, orig_h) if bbox is None: bbox = (0, 0, orig_w, orig_h) # no mask found -> pass the whole frame x1, y1, x2, y2 = bbox crop_w0, crop_h0 = x2 - x1, y2 - y1 # VRAM cap: downscale the crop only if its longest side exceeds max_long_side # (never upscales). This one cropped image feeds BOTH the segmentation and the # matte model, so it bounds the area both process. long_side = max(crop_w0, crop_h0) if max_long_side and max_long_side > 0 and long_side > max_long_side: scale = max_long_side / long_side out_w = int(crop_w0 * scale) out_h = int(crop_h0 * scale) else: out_w, out_h = crop_w0, crop_h0 out_w = max((out_w // 8) * 8, 8) out_h = max((out_h // 8) * 8, 8) img_out = [] for i in range(num): im = (images[i].cpu().numpy() * 255).astype(np.uint8)[y1:y2, x1:x2] im = np.array(Image.fromarray(im, mode='RGB').resize((out_w, out_h), Image.LANCZOS)) img_out.append(im.astype(np.float32) / 255.0) msk_out = [] for m in masks_full: mm = m[y1:y2, x1:x2] mm = np.array(Image.fromarray(mm, mode='L').resize((out_w, out_h), Image.NEAREST)) msk_out.append(mm.astype(np.float32) / 255.0) cropped_images = torch.from_numpy(np.stack(img_out, axis=0)) cropped_mask = torch.from_numpy(np.stack(msk_out, axis=0)) if msk_out else torch.zeros((num, out_h, out_w)) focus = {"bbox": (x1, y1, x2, y2), "orig_w": orig_w, "orig_h": orig_h, "crop_w": crop_w0, "crop_h": crop_h0} print(f"[FocusCrop] {orig_w}x{orig_h} -> hero box {crop_w0}x{crop_h0} at " f"({x1},{y1}) -> output {out_w}x{out_h}") return (cropped_images, cropped_mask, focus) class KoolookMatteFocusStitch: """Paste a crop-space matte (produced inside KoolookMatteFocusCrop's region) back into the full-frame plate. Outside the hero box the matte is 0 (background).""" @classmethod def INPUT_TYPES(cls): return { "required": { "focus": ("FOCUS",), "cropped_mask": ("MASK",), } } RETURN_TYPES = ("MASK",) RETURN_NAMES = ("mask",) FUNCTION = "stitch" CATEGORY = "Koolook Matte" def stitch(self, focus, cropped_mask): x1, y1, x2, y2 = focus["bbox"] orig_w, orig_h = focus["orig_w"], focus["orig_h"] crop_w, crop_h = focus["crop_w"], focus["crop_h"] out = [] for i in range(cropped_mask.shape[0]): m = _ensure_2d_mask((cropped_mask[i].cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)) # Resize the matte back to the hero box's native pixels, then paste. if m.shape[:2] != (crop_h, crop_w): m = np.array(Image.fromarray(m, mode='L').resize((crop_w, crop_h), Image.LANCZOS)) canvas = np.zeros((orig_h, orig_w), dtype=np.float32) canvas[y1:y2, x1:x2] = m.astype(np.float32) / 255.0 out.append(canvas) full = torch.from_numpy(np.stack(out, axis=0)) print(f"[FocusStitch] pasted {crop_w}x{crop_h} matte back into {orig_w}x{orig_h}") return (full,) # Node class mappings NODE_CLASS_MAPPINGS = { "KoolookMatteLoader": KoolookMatteLoader, "KoolookMatteSampler": KoolookMatteSampler, "KoolookMatteFocusCrop": KoolookMatteFocusCrop, "KoolookMatteFocusStitch": KoolookMatteFocusStitch, } NODE_DISPLAY_NAME_MAPPINGS = { "KoolookMatteLoader": "Koolook Matte — Loader", "KoolookMatteSampler": "Koolook Matte — Sampler", "KoolookMatteFocusCrop": "Koolook Matte — Focus Crop", "KoolookMatteFocusStitch": "Koolook Matte — Focus Stitch", }