ComfyUI-MatAnyone/mat_anyone2.py at main · siliconflow/ComfyUI-MatAnyone · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
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).cpu()

        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