nodes/FuouM-ComfyUI-MatAnyone
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

memory optimization and missing mat anyone 2 py

Browse Source
This commit is contained in:
FuouM
2026-03-07 15:27:19 +07:00
parent cbfd7abd31
commit c4e8f67608
9 changed files with 1685 additions and 2013 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+10 -10
View File
@@ -18,7 +18,9 @@ checkpoint/
## Workflow
[workflow/workflow_mat_anyone.json](workflow/workflow_mat_anyone.json)
The extension now includes the `MatAnyone2Video` node, which runs the improved MatAnyoneV2 model for higher-quality and more robust video matting.
**Workflow**: [workflow/workflow_mat_anyone.json](workflow/workflow_mat_anyone.json)
![matanyone_example](workflow/mat_anyone_preview.png)
(Not a workflow-embedded image)
@@ -31,18 +33,16 @@ Inputs:
- `mask_frame`: The input mask's index (defaults to 0). Support first (0), last and middle frame.
- `n_warmup`: Number of iterations to warm up the model. Defaults to 10.
**Memory Management Inputs**:
- `max_internal_size` (optional): Resizes the internal processing resolution to save memory (e.g., 360, 480). Default is -1 (full resolution).
- `max_mem_frames` (optional): The number of key frames kept in high-resolution working memory. Default is 5.
- `use_long_term` (optional): Limits active memory scaling by compressing older frames into long-term prototype memories, preventing OOM over time. Default is False. (Useful for long videos)
Your input mask won't actually be in the final matte. Instead, the warmup process generate a new input mask, which is then propagated throughout the video.
### MatAnyone 2 Support
The extension now includes the `MatAnyone2Video` node, which runs the improved MatAnyoneV2 model for higher-quality and more robust video matting.
**Workflow V2**: [workflow/workflow_mat_anyone_v2.json](workflow/workflow_mat_anyone_v2.json)
![matanyone_example](workflow/mat_anyone_preview_v2.png)
(Not a workflow-embedded image)
Additional Inputs for V2:
- `r_erode` (optional): The radius for morphological erosion applied to the `foreground_mask` before processing (defaults to 0). Useful for refining rough masks.
- `r_dilate` (optional): The radius for morphological dilation applied to the `foreground_mask` before processing (defaults to 0).
mat_anyone.py
+19 -16
View File
@@ -66,7 +66,6 @@ def inference_matanyone(
n_warmup=10,
):
mask = preprocess_mask(mask)
vframes = vframes.to(device)
mask = mask.to(device)
length = vframes.shape[0]
phas = [None] * length # initialize output list
@@ -80,13 +79,13 @@ def inference_matanyone(
repeated_frames = repeated_frames.to(device)
output_prob = warming_up(mask, processor, n_warmup, repeated_frames)
phas[0] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
phas[0] = processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
pbar = ProgressBar(length - 1)
for ti in tqdm(range(1, length), desc="Forward Propagation"):
image = vframes[ti]
image = vframes[ti].to(device)
output_prob = processor.step(image)
phas[ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
phas[ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
pbar.update_absolute(ti - 1, length - 1)
elif frame_index == length - 1:
@@ -98,15 +97,17 @@ def inference_matanyone(
reversed_mask = mask # mask for the last frame
output_prob = warming_up(reversed_mask, processor, n_warmup, repeated_frames)
phas[frame_index] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
phas[frame_index] = (
processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
)
pbar = ProgressBar(length - 1)
for ti in tqdm(range(1, length), desc="Backward Propagation"):
image = reversed_vframes[ti]
image = reversed_vframes[ti].to(device)
output_prob = processor.step(image)
phas[length - 1 - ti] = processor.output_prob_to_mask(
output_prob
).unsqueeze(0) # reverse index
phas[length - 1 - ti] = (
processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
) # reverse index
pbar.update_absolute(ti - 1, length - 1)
elif 0 < frame_index < length - 1:
@@ -117,14 +118,16 @@ def inference_matanyone(
# Warm up at frame_index
output_prob = warming_up(mask, processor, n_warmup, repeated_frames)
phas[frame_index] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
phas[frame_index] = (
processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
)
# Forward Propagation (from frame_index + 1 to end)
pbar_forward = ProgressBar(length - 1 - frame_index)
for ti in tqdm(range(frame_index + 1, length), desc="Forward Propagation"):
image = vframes[ti]
image = vframes[ti].to(device)
output_prob = processor.step(image)
phas[ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0)
phas[ti] = processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
pbar_forward.update_absolute(
ti - (frame_index + 1), length - 1 - frame_index
)
@@ -134,11 +137,11 @@ def inference_matanyone(
reversed_vframes_backward = torch.flip(vframes[:frame_index], dims=[0])
pbar_backward = ProgressBar(frame_index)
for ti in tqdm(range(frame_index), desc="Backward Propagation"):
image = reversed_vframes_backward[ti]
image = reversed_vframes_backward[ti].to(device)
output_prob = processor.step(image)
phas[frame_index - 1 - ti] = processor.output_prob_to_mask(
output_prob
).unsqueeze(0) # reverse index
phas[frame_index - 1 - ti] = (
processor.output_prob_to_mask(output_prob).unsqueeze(0).cpu()
) # reverse index
pbar_backward.update_absolute(ti, frame_index)
return phas
mat_anyone2.py
+182
View File
@@ -0,0 +1,182 @@
"""
MatAnyone2 inference wrapper for ComfyUI.
"""
from pathlib import Path
import numpy as np
import torch
from comfy.utils import ProgressBar
from .constants import ckpt_path_matanyone2
from .src.matanyone2.inference.inference_core import InferenceCore
from .src.matanyone2.utils.get_default_model import get_matanyone2_model
from .src.matanyone2.utils.inference_utils import gen_dilate, gen_erosion
base_dir = Path(__file__).resolve().parent
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def get_matanyone2_model_cached():
"""Loads MatAnyone2 model from checkpoint (user must place matanyone2.pth in checkpoint/)."""
ckpt_path = base_dir / ckpt_path_matanyone2
if not ckpt_path.exists():
raise FileNotFoundError(
"MatAnyone2 checkpoint not found at %s. "
"Download matanyone2.pth and place it in the checkpoint folder." % ckpt_path
)
return get_matanyone2_model(str(ckpt_path), device)
def warming_up2(
mask_t, processor, n_warmup, repeated_frames, pbar=None, pbar_start=0, total_len=0
):
output_prob = None
objects = [1]
for ti in range(n_warmup):
image = repeated_frames[ti]
if ti == 0:
output_prob = processor.step(image, mask_t, objects=objects)
output_prob = processor.step(image, first_frame_pred=True)
else:
output_prob = processor.step(image, first_frame_pred=True)
if pbar is not None:
pbar.update_absolute(pbar_start + ti, total_len)
return output_prob
def inference_matanyone2(
vframes: torch.Tensor,
mask: torch.Tensor,
processor: InferenceCore,
mask_frame: int = 0,
n_warmup: int = 10,
r_erode: int = 0,
r_dilate: int = 0,
):
"""
Runs MatAnyone2 inference on video frames.
Args:
vframes: (T, C, H, W) RGB, float [0, 1]
mask: (H, W) or (1, H, W) float [0, 1] or int
processor: InferenceCore instance
mask_frame: Frame index for mask
n_warmup: Warmup iterations
r_erode: Erosion kernel size
r_dilate: Dilation kernel size
Returns:
List of alpha tensors (1, 1, H, W) per frame
"""
mask = mask.squeeze().float()
if mask.max() <= 1.0:
mask = (mask * 255).clamp(0, 255)
mask_np = mask.cpu().numpy().astype(np.uint8)
if r_dilate > 0:
mask_np = gen_dilate(mask_np, r_dilate, r_dilate)
if r_erode > 0:
mask_np = gen_erosion(mask_np, r_erode, r_erode)
mask_t = torch.from_numpy(mask_np).float().to(device)
length = vframes.shape[0]
if mask_frame >= length:
raise ValueError(f"Expected 0 <= x < {length}, got {mask_frame}")
phas = [None] * length
if mask_frame == 0:
total_len = n_warmup + length
pbar = ProgressBar(total_len)
repeated_frames = vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1).to(device)
output_prob = warming_up2(
mask_t, processor, n_warmup, repeated_frames, pbar, 0, total_len
)
mask_out = processor.output_prob_to_mask(output_prob)
phas[0] = mask_out.unsqueeze(0).unsqueeze(0)
for ti in range(1, length):
image = vframes[ti].to(device)
output_prob = processor.step(image)
mask_out = processor.output_prob_to_mask(output_prob)
phas[ti] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
pbar.update_absolute(n_warmup + ti, total_len)
elif mask_frame == length - 1:
total_len = n_warmup + length
pbar = ProgressBar(total_len)
reversed_vframes = torch.flip(vframes, dims=[0])
repeated_frames = (
reversed_vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1).to(device)
)
output_prob = warming_up2(
mask_t, processor, n_warmup, repeated_frames, pbar, 0, total_len
)
mask_out = processor.output_prob_to_mask(output_prob)
phas[mask_frame] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
for ti in range(1, length):
image = reversed_vframes[ti].to(device)
output_prob = processor.step(image)
mask_out = processor.output_prob_to_mask(output_prob)
phas[length - 1 - ti] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
pbar.update_absolute(n_warmup + ti, total_len)
elif 0 < mask_frame < length - 1:
# Two passes: forward then backward
total_len = n_warmup + (length - mask_frame) + n_warmup + mask_frame
pbar = ProgressBar(total_len)
current_step = 0
# Pass 1: Forward
repeated_frames = (
vframes[mask_frame].unsqueeze(0).repeat(n_warmup, 1, 1, 1).to(device)
)
output_prob = warming_up2(
mask_t, processor, n_warmup, repeated_frames, pbar, current_step, total_len
)
current_step += n_warmup
mask_out = processor.output_prob_to_mask(output_prob)
phas[mask_frame] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
for ti in range(mask_frame + 1, length):
image = vframes[ti].to(device)
output_prob = processor.step(image)
mask_out = processor.output_prob_to_mask(output_prob)
phas[ti] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
pbar.update_absolute(current_step, total_len)
current_step += 1
# Pass 2: Backward
processor = type(processor)(processor.network, cfg=processor.network.cfg)
repeated_frames_back = (
vframes[mask_frame].unsqueeze(0).repeat(n_warmup, 1, 1, 1).to(device)
)
output_prob = warming_up2(
mask_t,
processor,
n_warmup,
repeated_frames_back,
pbar,
current_step,
total_len,
)
current_step += n_warmup
reversed_vframes_backward = torch.flip(vframes[:mask_frame], dims=[0])
for ti in range(mask_frame):
image = reversed_vframes_backward[ti].to(device)
output_prob = processor.step(image)
mask_out = processor.output_prob_to_mask(output_prob)
phas[mask_frame - 1 - ti] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
pbar.update_absolute(current_step, total_len)
current_step += 1
return phas
run.py
+38
View File
@@ -45,6 +45,18 @@ class MatAnyoneVideo:
"foreground_mask": ("IMAGE",),
"foreground_MASK": ("MASK",),
"solid_color": ("IMAGE",),
"max_internal_size": (
"INT",
{"default": -1, "min": -1, "step": 1},
),
"max_mem_frames": (
"INT",
{"default": 5, "min": 1, "step": 1},
),
"use_long_term": (
"BOOLEAN",
{"default": False},
),
},
}
@@ -67,12 +79,19 @@ class MatAnyoneVideo:
foreground_mask: torch.Tensor | None = None,
foreground_MASK: torch.Tensor | None = None,
solid_color: torch.Tensor | None = None,
max_internal_size: int = -1,
max_mem_frames: int = 5,
use_long_term: bool = False,
):
mask = get_mask(foreground_mask, foreground_MASK)
src_video = src_video.permute(0, 3, 1, 2) # T CHW RGB
# load MatAnyone model
matanyone = get_matanyone_model(f"{base_dir}/{ckpt_path}")
matanyone.cfg.max_internal_size = max_internal_size
matanyone.cfg.max_mem_frames = max_mem_frames
matanyone.cfg.use_long_term = use_long_term
processor = InferenceCore(matanyone, cfg=matanyone.cfg)
phas = inference_matanyone(src_video, mask, processor, mask_frame, n_warmup)
out_mask = torch.cat(phas).unsqueeze(1).permute(0, 2, 3, 1)
@@ -123,6 +142,18 @@ class MatAnyone2Video:
"INT",
{"default": 0, "min": 0, "step": 1},
),
"max_internal_size": (
"INT",
{"default": -1, "min": -1, "step": 1},
),
"max_mem_frames": (
"INT",
{"default": 5, "min": 1, "step": 1},
),
"use_long_term": (
"BOOLEAN",
{"default": False},
),
},
}
@@ -147,11 +178,18 @@ class MatAnyone2Video:
solid_color: torch.Tensor | None = None,
r_erode: int = 0,
r_dilate: int = 0,
max_internal_size: int = -1,
max_mem_frames: int = 5,
use_long_term: bool = False,
):
mask = get_mask(foreground_mask, foreground_MASK)
src_video = src_video.permute(0, 3, 1, 2) # T C H W RGB, [0, 1]
matanyone2 = get_matanyone2_model_cached()
matanyone2.cfg.max_internal_size = max_internal_size
matanyone2.cfg.max_mem_frames = max_mem_frames
matanyone2.cfg.use_long_term = use_long_term
processor = InferenceCore2(matanyone2, cfg=matanyone2.cfg)
phas = inference_matanyone2(
src_video,
src/matanyone2/inference/inference_core.py
+14 -6
View File
@@ -234,12 +234,20 @@ class InferenceCore:
.long()
)
else:
mask = F.interpolate(
mask.unsqueeze(0),
size=(new_h, new_w),
mode="bilinear",
align_corners=False,
)[0]
if mask.dim() == 2:
mask = F.interpolate(
mask.unsqueeze(0).unsqueeze(0).float(),
size=(new_h, new_w),
mode="bilinear",
align_corners=False,
)[0, 0]
else:
mask = F.interpolate(
mask.unsqueeze(0).float(),
size=(new_h, new_w),
mode="bilinear",
align_corners=False,
)[0]
self.curr_ti += 1
workflow/mat_anyone_preview.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 332 KiB

After

Width:  |  Height:  |  Size: 260 KiB

workflow/mat_anyone_preview_v2.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 269 KiB

workflow/workflow_mat_anyone.json
+1422 -617
View File
File diff suppressed because it is too large Load Diff
workflow/workflow_mat_anyone_v2.json
-1364
View File
File diff suppressed because it is too large Load Diff