diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f54028..d246691 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ The format is inspired by Keep a Changelog and SemVer. ## [Unreleased] ### Added +- **Koolook Matte nodes (`KoolookMatteLoader` / `KoolookMatteSampler` / + `KoolookMatteFocusCrop` / `KoolookMatteFocusStitch`).** First-party, mask-guided + one-step video matting — an original GPL-3.0 reimplementation of the published + VideoMaMa method (single-step SVD, 12-channel UNet), with temporal windowing, + subject "hero-crop" ROI, and shortest-side/`/64` resolution control with OOM + auto-backoff. **Model caveat:** the nodes *load* (never bundle) the VideoMaMa + UNet (**CC BY-NC 4.0 — non-commercial**) and SVD's VAE (**Stability Community + License**) at runtime; the user is responsible for model-license compliance + (see `matte/NOTICE`). Maintainer notes: `docs/maintainers/matte-pipeline.md`. - **LTX Keyframe Soften Schedule node (`LTXKeyframeSoftenSchedule`).** Smooths motion between keyframes without adding any guide frames: a model patcher that sigma-schedules the keyframe denoise-mask pins — loosening them early (high diff --git a/README.md b/README.md index 79605c9..30c79d8 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,9 @@ The root Koolook nodes and the slim Radiance Koolook VAE wrappers have no extra Python dependencies beyond `torch`, which ComfyUI already requires. The video wrappers require Video Helper Suite, and the LTX Director fork is for workflows that already use the LTX / Prompt Relay node ecosystem. +The Koolook Matte nodes additionally need `diffusers` (already pulled in by the +LTX Director / Radiance forks) and use `huggingface_hub` to download the SVD + +VideoMaMa model weights at runtime. ## Install diff --git a/__init__.py b/__init__.py index d8ed732..a8f1348 100644 --- a/__init__.py +++ b/__init__.py @@ -143,6 +143,9 @@ else: _merge_node_group("Easy Video Combine", ".k_video_combine") _merge_node_group("Radiance Koolook VAE", ".forks.radiance_koolook") _merge_node_group("WhatDreamsCost LTX Director", ".forks.whatdreamscost_koolook") + # First-party matte nodes. Load diffusers (already present via the forks + # above); the guard skips them cleanly if diffusers is ever absent. + _merge_node_group("Koolook Matte", ".matte") # Register the Kforge Labs snapshot/preset endpoints. Installed # regardless of the node-import results above, so a broken node diff --git a/constraints-test.txt b/constraints-test.txt index a1f4257..64002fe 100644 --- a/constraints-test.txt +++ b/constraints-test.txt @@ -33,7 +33,7 @@ msgpack==1.2.1 multidict==6.7.1 packageurl-python==0.17.6 packaging==26.2 -pillow==12.2.0 +pillow==12.3.0 pip-api==0.0.34 pip_audit==2.10.1 pip-requirements-parser==32.0.1 diff --git a/docs/maintainers/matte-pipeline.md b/docs/maintainers/matte-pipeline.md new file mode 100644 index 0000000..52288d1 --- /dev/null +++ b/docs/maintainers/matte-pipeline.md @@ -0,0 +1,147 @@ +# Koolook Matte — pipeline internals, fidelity, and debugging + +`matte/` is a first-party (GPL-3.0, like the rest of the repo) **clean +reimplementation** of the published +**VideoMaMa** one-step, mask-guided SVD matting method. It loads (never bundles) +the SVD base VAE + the VideoMaMa-fine-tuned UNet and runs a single denoising step +at `timestep=1` with CLIP cross-attention zeroed, then decodes the latents to a +luminance alpha. + +Because it is an independent reimplementation, the trap is **subtle numerical +divergence from the reference** — the pipeline can run end-to-end, produce a +correctly-shaped alpha, and still be wrong (black, or blurry). This doc records +the two bugs that actually bit us, how to tell them apart, and the fast loop that +catches them. + +The upstream reference we validate against is okdalto's `ComfyUI-VideoMaMa` +(unlicensed — do **not** vendor it; it's a *comparison oracle* only). Our nodes +are an independent implementation that targets behavioural parity with it, so a +side-by-side A/B isolates any +regression to `matte/pipeline.py`. + +--- + +## 1. Latent scaling — feed the UNet RAW latents (×1) + +**This is the single most important invariant.** The VideoMaMa UNet was fine-tuned +on conditioning + mask latents at the **raw VAE scale** (`vae.encode(...).latent_dist.sample()`, +with **no** `scaling_factor` applied net). Feed it anything else and it goes +out-of-distribution. + +The confusing part: okdalto ships **two** pipeline classes that disagree, and only +one matches the shipped weights: + +| class | net cond/mask scale | result | +|---|---|---| +| `VideoInferencePipeline` (what the ComfyUI node uses) | `sample × s` then `/ s` → **×1 (raw)** | **correct / sharp** | +| `...OnestepWithMask` (a diffusers-style variant) | `sample` then `/ s` → **×5.49** | over-scaled | + +`s = vae.config.scaling_factor = 0.18215`, so `1/s ≈ 5.49`. + +**Our `_encode_latents` must return raw latents** (no `/ s`): + +```python +lat = torch.cat(parts, dim=0) # RAW — do NOT divide by scaling_factor +return lat.reshape(b, f, *lat.shape[1:]) +``` + +Decode is unchanged and matches both references: divide the UNet output by `s` +before `vae.decode` (`pred = pred.squeeze(0) / s`). + +### The regression this caused +While chasing a black output we diffed against the *wrong* okdalto class +(`OnestepWithMask`, ×5.49) and added `/ scaling_factor` to `_encode_latents`. +That over-scaled the conditioning ~5.5×, pushed the UNet OOD, and produced a +**systematically blurry** matte on every frame — subtle enough to look like a +"resolution" problem. Removing the division restored parity with the reference +(rated 5/5 vs the desktop/vanilla output on identical inputs). + +**Rule of thumb:** cond/mask latents in, alpha out — if the alpha is *soft on +every frame at every seed*, suspect a scale/normalization divergence in +`pipeline.py` before you suspect resolution. + +--- + +## 2. Black output ≠ scaling + +An all-black matte is almost always an **empty guide mask**, not the pipeline: + +- **`SAM3Propagate` `direction = "backward"` with `start_frame = 0`** — nothing to + propagate *before* frame 0, so the track is empty on frames 1…N. Use `forward` + (or `both`) from a seed frame. +- **Text prompt that doesn't match the subject** (e.g. `"goats"` on eagle footage) + → SAM3 finds nothing → empty mask. + +Confirm with the SAM3 output preview: white subject on black = good guide; all +black = fix the seg, not the matte. (The pixel test: our matte, the guide mask, +and the source should all be non-black before you blame the sampler.) + +--- + +## 3. Fidelity levers (in priority order) + +The matte decodes from an SVD VAE latent (8× downsample), so **effective pixels on +the subject** is the dominant detail lever. + +1. **Processing resolution.** `KoolookMatteSampler.max_resolution` + `resize_mode` + (`shortest_side` = desktop parity, never upscales). The sampler downscales the + input to fit `max_resolution`; the decoded alpha is upscaled back. Keep this at + or above the crop's native short side to avoid a needless downscale→upscale + round-trip. +2. **OOM auto-backoff reduces *resolution*** (`target × 0.85`) to fit VRAM — which + directly costs detail (we saw a `1408 → 1152` backoff soften the fur). For + *detail* work, prefer to buy VRAM back from the **temporal window** + (`window_size 16 → 8`) instead, so spatial resolution stays high. (The default + backoff-resolution behaviour favours temporal stability over sharpness — a + deliberate but detail-costing trade.) +3. **FocusCrop hero-box** (`KoolookMatteFocusCrop`) crops to the subject's union + bbox so `max_resolution` applies to just the subject. It never upscales and + keeps aspect (single scale factor); `max_long_side = 0` = native crop. + Run it on the **native-res plate before any downscale** — a crop downstream of + a shrink can't recover lost detail. +4. **Aspect-preserving `/64` alignment.** The sampler aligns W and H independently + to `/64` (SVD requirement); this nudges the aspect each step, and the alpha is + subtly stretched on stitch. Minimise resamples: crop native → process native + (or one clean cap) → stitch, with rescaling only when OOM forces it. + +--- + +## 4. Fast debugging loop + +Full clips are 150–220s and can OOM. To iterate on fidelity: + +1. **Run one frame.** Set the shot-length control to `1` (Custom). ~4s, ~9GB, no + OOM backoff — so you see the *intended* processing resolution, not a degraded + one. (A ComfyUI restart resets the shot-length widget; re-set it to 1.) +2. **Side-by-side the oracle.** Drop okdalto's `VideoMaMaPipelineLoader` + + `VideoMaMaSampler` into the graph, wire them to the **same** FocusCrop crop and + SAM3 mask, copy the identical widget values. Our node and the oracle then differ + only in `pipeline.py`, so any gap is ours. +3. **Rate 1–5 by eye** (1 = the known-blurry state, 5 = matches the reference). + Automated sharpness metrics on ComfyUI temp previews are fragile (temp files get + recycled); the human rating is the reliable signal. +4. Change one lever, re-run the single frame, re-rate. + +The vanilla nodes stay in the graph as a permanent control. + +--- + +## 5. dev-sync & checkpoints + +The SVD base model + VideoMaMa UNet (~10 GB) live **only in the live install** +under `matte/checkpoints/` — they're gitignored and downloaded at runtime, so the +repo has no copy. + +`scripts/sync_to_dev.py` mirrors each runtime path (`rmtree` + `copytree`). Left +naive, that **deletes `matte/checkpoints/` on every matte sync** and forces a +multi-GB re-download on the next node load (this manifested as ComfyUI "stuck at +8%"). The sync now protects it via `PRESERVE_IN_DEST = {"matte/checkpoints"}` +(an **exact path relative to the sync target**, not a bare name matched at any +depth) and a selective `_clean_dest` that keeps preserved subtrees while still +mirroring code. + +If you add another node pack with runtime-downloaded weights, add its +**target-relative path** (e.g. `"otherpack/checkpoints"`) to `PRESERVE_IN_DEST`. + +See also [`dev-iteration-loop.md`](dev-iteration-loop.md) for the general +push/publish gates. diff --git a/matte/NOTICE b/matte/NOTICE new file mode 100644 index 0000000..714d113 --- /dev/null +++ b/matte/NOTICE @@ -0,0 +1,35 @@ +Koolook Matte — NOTICE +====================== + +The CODE in this pack is licensed under GPL-3.0 (see the repository LICENSE) and +is an ORIGINAL, independent implementation — it does not incorporate third-party +wrapper code. + +Method credit +------------- +This pack implements the method described in: + "VideoMaMa: Mask-Guided Video Matting via Generative Prior" (CVPR 2026), + cvlab-kaist — https://github.com/cvlab-kaist/VideoMaMa +The method, the original reference code, and the trained weights are the work of +those authors. This project re-implements the (published) inference method using +the standard `diffusers` API and adds its own windowing / focus-crop tooling. + +Model weights (NOT bundled or redistributed — loaded at runtime) +---------------------------------------------------------------- +* VideoMaMa fine-tuned UNet — **CC BY-NC 4.0 (NON-COMMERCIAL)**. + Source: `SammyLim/VideoMaMa` (Hugging Face) / cvlab-kaist/VideoMaMa. + ==> Do NOT use these weights for commercial work without a separate license + from the model authors. +* Stable Video Diffusion img2vid-xt (its temporal VAE is used) — + **Stability AI Community License**. + Source: `stabilityai/stable-video-diffusion-img2vid-xt`. + +You (the user) are solely responsible for complying with the above model +licenses. This pack downloading/loading a model does not grant you any rights +to it. + +Technique note +-------------- +The temporal sliding-window batching and the subject bounding-box crop/paste +implemented here are standard, well-known techniques, written independently for +this pack. diff --git a/matte/__init__.py b/matte/__init__.py new file mode 100644 index 0000000..30da533 --- /dev/null +++ b/matte/__init__.py @@ -0,0 +1,12 @@ +"""Koolook Matte — mask-guided one-step video matting nodes. + +Part of the Koolook pack. Registered via the parent ``__init__``'s guarded +``_merge_node_group("Koolook Matte", ".matte")`` so a missing ``diffusers`` +disables only these nodes, not the rest of Koolook. + +The nodes LOAD (never bundle) the VideoMaMa UNet + SVD VAE weights at runtime; +those weights carry their own non-commercial licenses — see ``matte/NOTICE``. +""" +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/matte/nodes.py b/matte/nodes.py new file mode 100644 index 0000000..4696dd1 --- /dev/null +++ b/matte/nodes.py @@ -0,0 +1,737 @@ +""" +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", +} diff --git a/matte/pipeline.py b/matte/pipeline.py new file mode 100644 index 0000000..2000759 --- /dev/null +++ b/matte/pipeline.py @@ -0,0 +1,221 @@ +""" +Koolook Matte — independent inference pipeline for mask-guided one-step video matting. + +This is an ORIGINAL implementation of the *published* VideoMaMa method +("VideoMaMa: Mask-Guided Video Matting via Generative Prior", CVPR 2026), written +from the method description + the standard `diffusers` API — it does not incorporate +any third-party wrapper code. The method itself: + + * a Stable-Video-Diffusion UNet fine-tuned to take **12 latent channels** + = 4 (noise) + 4 (VAE-encoded RGB frame) + 4 (VAE-encoded guide mask), + * a **single** denoising step at timestep 1 with the CLIP cross-attention + context zeroed (the fine-tune ignores it), + * decode the predicted latents with SVD's temporal VAE and take luminance + as the alpha matte. + +The model WEIGHTS it loads (the VideoMaMa UNet, and SVD's VAE) carry their own +licenses (CC BY-NC 4.0 and the Stability Community License respectively); this +code does not redistribute them and the end user is responsible for compliance. + +Licensed under GPL-3.0 — see the repository LICENSE. +""" + +import numpy as np +import torch + +from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel + + +class KoolookMattePipeline: + """One-step mask-guided SVD matting. Reusable across many clips/windows.""" + + def __init__(self, base_model_path: str, unet_checkpoint_path: str, device: str = "cuda", + weight_dtype: torch.dtype = torch.float16, enable_model_cpu_offload: bool = False, + vae_encode_chunk_size: int = 1, attention_mode: str = "auto", + enable_vae_tiling: bool = False, enable_vae_slicing: bool = True): + self.device = torch.device(device) + self.weight_dtype = weight_dtype + self.enable_model_cpu_offload = enable_model_cpu_offload + self.vae_encode_chunk_size = max(1, int(vae_encode_chunk_size)) + + # SVD's temporal VAE (from the base model) + the VideoMaMa-fine-tuned UNet. + try: + self.vae = AutoencoderKLTemporalDecoder.from_pretrained(base_model_path, subfolder="vae", variant="fp16") + except Exception: + # some snapshots ship the vae without an fp16 variant + self.vae = AutoencoderKLTemporalDecoder.from_pretrained(base_model_path, subfolder="vae") + self.unet = UNetSpatioTemporalConditionModel.from_pretrained(unet_checkpoint_path, subfolder="unet") + + self.vae.eval() + self.unet.eval() + self._apply_attention_optimization(attention_mode) + + if enable_vae_slicing: + try: + self.vae.enable_slicing() + except (AttributeError, NotImplementedError): + pass + if enable_vae_tiling: + try: + self.vae.enable_tiling() + except (AttributeError, NotImplementedError): + pass + + if self.enable_model_cpu_offload: + self.vae.to("cpu", dtype=self.weight_dtype) + self.unet.to("cpu", dtype=self.weight_dtype) + else: + self.vae.to(self.device, dtype=self.weight_dtype) + self.unet.to(self.device, dtype=self.weight_dtype) + + # -- setup helpers ------------------------------------------------------ + + def _apply_attention_optimization(self, attention_mode: str): + if attention_mode == "none": + return + if attention_mode in ("auto", "xformers"): + try: + import xformers # noqa: F401 + self.unet.enable_xformers_memory_efficient_attention() + return + except Exception: + if attention_mode == "xformers": + print("KoolookMatte: xformers unavailable, falling back") + if attention_mode in ("auto", "sdpa"): + if hasattr(torch.nn.functional, "scaled_dot_product_attention"): + try: + from diffusers.models.attention_processor import AttnProcessor2_0 + self.unet.set_attn_processor(AttnProcessor2_0()) + return + except Exception: + pass + + @staticmethod + def _clear_cache(device): + if device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + elif device.type == "mps" and torch.backends.mps.is_available(): + torch.mps.empty_cache() + elif device.type == "xpu": + torch.xpu.empty_cache() + + # -- core inference ----------------------------------------------------- + + @torch.no_grad() + def run(self, cond_frames, mask_frames, seed=42, fps=7, motion_bucket_id=127, + noise_aug_strength=0.0, pbar=None): + """Run one-step matting on a list of PIL frames + PIL guide masks. + + Returns a list of grayscale PIL alpha frames (mode 'L'). + `pbar` is an optional ComfyUI ProgressBar (updated 4x per pass). + """ + def tick(): + if pbar is not None: + pbar.update(1) + + cond = self._pil_to_tensor(cond_frames).to(self.device) # [1,F,3,H,W] in [-1,1] + mask = self._pil_to_tensor(mask_frames, grayscale_ok=True).to(self.device) + if mask.shape[2] != 3: # replicate L -> 3ch for the VAE + mask = mask.repeat(1, 1, 3, 1, 1) + + hidden = self.unet.config.cross_attention_dim + encoder_hidden_states = torch.zeros((1, 1, hidden), dtype=self.weight_dtype, device=self.device) + + # --- encode RGB frame + guide mask to latents --- + if self.enable_model_cpu_offload: + self.vae.to(self.device) + cond_latents = self._encode_latents(cond.to(self.weight_dtype)) + mask_latents = self._encode_latents(mask.to(self.weight_dtype)) + tick() + del cond, mask + self._clear_cache(self.device) + if self.enable_model_cpu_offload: + self.vae.to("cpu") + self._clear_cache(self.device) + + # --- single UNet step over all frames at once --- + if self.enable_model_cpu_offload: + self.unet.to(self.device) + gen = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn(cond_latents.shape, generator=gen, device="cpu", + dtype=self.weight_dtype).to(self.device) + timesteps = torch.ones((1,), device=self.device, dtype=torch.long) + added_time_ids = self._add_time_ids(fps, motion_bucket_id, noise_aug_strength) + + unet_in = torch.cat([noise, cond_latents, mask_latents], dim=2) # 12 channels + del noise, cond_latents, mask_latents + self._clear_cache(self.device) + tick() + + pred = self.unet(unet_in, timesteps, encoder_hidden_states, added_time_ids=added_time_ids).sample + del unet_in + if self.enable_model_cpu_offload: + self.unet.to("cpu") + self._clear_cache(self.device) + tick() + + # --- decode latents -> luminance alpha --- + if self.enable_model_cpu_offload: + self.vae.to(self.device) + pred = pred.squeeze(0) / self.vae.config.scaling_factor # [F,4,h,w] + chunk = min(self.vae_encode_chunk_size, pred.shape[0]) + decoded = [] + for i in range(0, pred.shape[0], chunk): + part = pred[i:i + chunk] + decoded.append(self.vae.decode(part, num_frames=part.shape[0]).sample.cpu()) + del pred + if self.enable_model_cpu_offload: + self.vae.to("cpu") + self._clear_cache(self.device) + tick() + + video = torch.cat(decoded, dim=0) # [F,3,H,W] + del decoded + video = (video / 2.0 + 0.5).clamp(0, 1).mean(dim=1) # [F,H,W] luminance + out = [] + for frame in video.float().cpu().numpy(): + out.append(_to_pil_L(frame)) + return out + + # -- tensor plumbing ---------------------------------------------------- + + def _pil_to_tensor(self, frames, grayscale_ok=False): + arrs = [] + for f in frames: + a = np.asarray(f) + if a.ndim == 2: + a = np.stack([a] * 3, axis=-1) + t = torch.from_numpy(a.copy()).permute(2, 0, 1).float() / 255.0 + arrs.append(t) + video = torch.stack(arrs).unsqueeze(0) # [1,F,3,H,W] + return video * 2.0 - 1.0 + + def _encode_latents(self, t): + """Encode [1,F,3,H,W] -> raw VAE latents [1,F,4,h,w] (chunked over frames). + + The UNet is fed conditioning + mask latents at the RAW VAE scale (no + ``scaling_factor`` applied), matching okdalto's ``VideoInferencePipeline`` + (which encodes ``* scaling_factor`` then divides it back out, netting to + raw). Dividing here would over-scale the conditioning ~5.5x (1/0.18215), + pushing the fine-tuned UNet out of distribution and softening the matte. + """ + b, f = t.shape[0], t.shape[1] + t = t.reshape(b * f, *t.shape[2:]) + parts = [] + for i in range(0, t.shape[0], self.vae_encode_chunk_size): + parts.append(self.vae.encode(t[i:i + self.vae_encode_chunk_size]).latent_dist.sample()) + lat = torch.cat(parts, dim=0) + return lat.reshape(b, f, *lat.shape[1:]) + + def _add_time_ids(self, fps, motion_bucket_id, noise_aug_strength): + vals = [fps, motion_bucket_id, noise_aug_strength] + want = self.unet.config.addition_time_embed_dim * len(vals) + have = self.unet.add_embedding.linear_1.in_features + if want != have: + raise ValueError(f"added_time_ids dim mismatch: model expects {have}, built {want}") + return torch.tensor([vals], dtype=self.weight_dtype, device=self.device) + + +def _to_pil_L(arr01): + from PIL import Image + return Image.fromarray((arr01 * 255).clip(0, 255).astype(np.uint8), mode="L") diff --git a/scripts/sync_to_dev.py b/scripts/sync_to_dev.py index dfb4b11..e90cf73 100644 --- a/scripts/sync_to_dev.py +++ b/scripts/sync_to_dev.py @@ -176,6 +176,7 @@ RUNTIME_PATHS: tuple[str, ...] = ( "koolook_setups.py", "koolook_versioning.py", "forks", + "matte", "video_formats", "web", ) @@ -190,6 +191,15 @@ DIR_EXCLUDES: tuple[str, ...] = ( "README.md", ) +# Paths (POSIX, relative to the sync TARGET root) that live ONLY in the live +# install — gitignored, downloaded at runtime — and must survive a sync. The +# repo has no copy of these, so a blanket ``rmtree`` of the dest subtree would +# delete them and force a multi-GB re-download on the next node load. +# ``matte/checkpoints/`` holds the SVD base model + VideoMaMa UNet (~10GB). +# Matched by EXACT relative path (not by bare name at any depth), so an +# unrelated ``checkpoints/`` elsewhere in the tree is still mirrored normally. +PRESERVE_IN_DEST: frozenset[str] = frozenset({"matte/checkpoints"}) + def load_dotenv(env_path: Path) -> None: """Minimal `.env` loader. No dependency on python-dotenv.""" @@ -239,6 +249,42 @@ def _ignore(_dir: str, names: list[str]) -> list[str]: return [n for n in names if n in DIR_EXCLUDES] +def _is_preserved(rel: str) -> bool: + """True if ``rel`` (POSIX, relative to the target root) is exactly a + preserved path.""" + return rel in PRESERVE_IN_DEST + + +def _contains_preserved(rel: str) -> bool: + """True if a preserved path lives strictly BELOW ``rel`` (so we descend into + ``rel`` instead of deleting it wholesale).""" + prefix = rel + "/" + return any(p.startswith(prefix) for p in PRESERVE_IN_DEST) + + +def _clean_dest(dst: Path, base: Path) -> None: + """Recursively remove ``dst``'s contents so the copy mirrors the repo, but + keep any ``PRESERVE_IN_DEST`` path (matched by its POSIX path relative to + ``base``). Replaces a blanket ``shutil.rmtree(dst)`` that would delete the + gitignored, runtime-downloaded model weights and force a ~10GB re-download. + """ + for child in dst.iterdir(): + rel = child.relative_to(base).as_posix() + if _is_preserved(rel): + continue # keep runtime-downloaded data (model checkpoints) + # Symlinks first: a symlinked dir reports is_dir()==True, so recursing / + # rmtree'ing it could reach outside the tree. Just drop the link. + if child.is_symlink(): + child.unlink() + elif child.is_dir(): + if _contains_preserved(rel): + _clean_dest(child, base) # descend to spare a nested preserved subtree + else: + shutil.rmtree(child) + else: + child.unlink() + + def sync( target: Path, dry_run: bool, @@ -267,9 +313,23 @@ def sync( copied += 1 continue if src.is_dir(): - if dst.exists(): - shutil.rmtree(dst) - shutil.copytree(src, dst, ignore=_ignore) + # Clean the dest for a fresh mirror, but keep PRESERVE_IN_DEST + # subtrees (runtime-downloaded model checkpoints the repo doesn't + # carry). A blanket rmtree here would nuke matte/checkpoints/ and + # trigger a ~10GB re-download on the next load. + # A top-level symlinked dest (e.g. target/matte -> elsewhere) must + # never be cleaned/copied THROUGH: iterdir()/copytree would follow it + # and mutate files outside the target. Drop the link and own the path + # as a real dir. (is_symlink is checked before exists(), which follows + # symlinks.) + if dst.is_symlink(): + dst.unlink() + dst.mkdir(parents=True, exist_ok=True) + elif dst.exists(): + _clean_dest(dst, target) + else: + dst.mkdir(parents=True, exist_ok=True) + shutil.copytree(src, dst, ignore=_ignore, dirs_exist_ok=True) else: dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(src, dst)