ComfyUI/comfy_extras/nodes_scail2.py

"""SCAIL-2 preprocessing nodes that turn SAM3 video tracks into the conditioning
bundle the SCAIL-2 model consumes."""

from typing_extensions import override

import torch
import torch.nn.functional as F

import comfy.model_management
import comfy.utils
from comfy_api.latest import ComfyExtension, io


SAM3TrackData = io.Custom("SAM3_TRACK_DATA")


# Model was trained on these exact colors; deviating degrades multi-identity quality.
DEFAULT_PALETTE = [
    (0.0, 0.0, 1.0),  # Blue
    (1.0, 0.0, 0.0),  # Red
    (0.0, 1.0, 0.0),  # Green
    (1.0, 0.0, 1.0),  # Magenta
    (0.0, 1.0, 1.0),  # Cyan
    (1.0, 1.0, 0.0),  # Yellow
]


def _unpack(track_data):
    from comfy.ldm.sam3.tracker import unpack_masks
    packed = track_data["packed_masks"]
    if packed is None or packed.shape[1] == 0:
        return None
    return unpack_masks(packed)


def _first_frame_cx_area(masks_bool):
    first = masks_bool[0].float()
    H, W = first.shape[-2], first.shape[-1]
    n_pixels = H * W
    grid_x = torch.arange(W, device=first.device, dtype=first.dtype).view(1, W)
    area = first.sum(dim=(-1, -2)).clamp_(min=1)
    cx = (first * grid_x).sum(dim=(-1, -2)) / area
    return (cx / W).tolist(), (area / n_pixels).tolist()


def _sort_tracks(track_data, sort_by):
    masks_bool = _unpack(track_data)
    if masks_bool is None:
        return []
    cx, area = _first_frame_cx_area(masks_bool)
    if sort_by == "x":
        return sorted(range(len(cx)), key=lambda i: cx[i])
    return sorted(range(len(area)), key=lambda i: -area[i])  # "area"


def _subset_track_data(track_data, obj_indices):
    out = dict(track_data)
    packed = track_data["packed_masks"]
    if packed is None or not obj_indices:
        out["packed_masks"] = None
        if "scores" in out:
            out["scores"] = []
        return out
    out["packed_masks"] = packed[:, obj_indices].contiguous()
    scores = track_data.get("scores")
    if scores is not None:
        out["scores"] = [scores[i] for i in obj_indices if i < len(scores)]
    return out


def _bg_to_rgb(background):
    if background.startswith("white"):
        return (1.0, 1.0, 1.0)
    return (0.0, 0.0, 0.0)


def _render_colored_masks(track_data, background="black"):
    from comfy.ldm.sam3.tracker import unpack_masks
    packed = track_data["packed_masks"]
    H, W = track_data["orig_size"]
    device = comfy.model_management.intermediate_device()
    bg_rgb = _bg_to_rgb(background)
    if packed is None or packed.shape[1] == 0:
        T = track_data.get("n_frames", 1) if packed is None else packed.shape[0]
        out = torch.empty(T, H, W, 3, device=device)
        out[..., 0], out[..., 1], out[..., 2] = bg_rgb[0], bg_rgb[1], bg_rgb[2]
        return out
    T, N_obj = packed.shape[0], packed.shape[1]
    colors = torch.tensor(
        [DEFAULT_PALETTE[i % len(DEFAULT_PALETTE)] for i in range(N_obj)],
        device=device, dtype=torch.float32,
    )
    masks_full = unpack_masks(packed.to(device)).float()
    Hm, Wm = masks_full.shape[-2], masks_full.shape[-1]
    masks_full = F.interpolate(
        masks_full.view(T * N_obj, 1, Hm, Wm), size=(H, W), mode="nearest"
    ).view(T, N_obj, H, W) > 0.5
    any_mask = masks_full.any(dim=1)
    obj_idx_map = masks_full.to(torch.uint8).argmax(dim=1)
    color_overlay = colors[obj_idx_map]
    bg_tensor = torch.tensor(bg_rgb, device=device, dtype=color_overlay.dtype).view(1, 1, 1, 3)
    return torch.where(any_mask.unsqueeze(-1), color_overlay, bg_tensor.expand_as(color_overlay))


class SCAIL2ColoredMask(io.ComfyNode):
    """Render SAM3 tracks for the driving video and (optionally) the reference
    image into the two colored masks WanSCAILToVideo consumes. Shared `sort_by`
    across both outputs guarantees identity K maps to the same color on both
    sides, so multi-person workflows stay consistent without a separate
    alignment node. ref_mask is always rendered black-bg (model convention);
    mask_video bg follows the mode you'll set on WanSCAILToVideo."""

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SCAIL2ColoredMask",
            display_name="SCAIL-2 Colored Mask",
            category="conditioning/video_models/scail",
            inputs=[
                SAM3TrackData.Input("driving_track_data"),
                SAM3TrackData.Input("ref_track_data", optional=True,
                                    tooltip="SAM3 track of the reference image. Optional — wire it for the ref_mask_image output."),
                io.String.Input("object_indices", default="",
                                tooltip="Comma-separated object indices to include (e.g. '0,2,3'). Applied to both sides. Empty = all."),
                io.Combo.Input("sort_by", options=["none", "x", "area"],
                               tooltip="Applied to both sides identically so index K = same logical slot. x = left-to-right by first-frame centroid; area = descending mask area; none = SAM3's order."),
                io.Boolean.Input("replacement_mode", default=False,
                                 tooltip="False = mask_video has black bg (Animation Mode). True = white bg (Replacement Mode). WanSCAILToVideo auto-detects mode from the wired mask_video's bg color, so this is the single source of truth. ref_mask_image is always black-bg regardless."),
            ],
            outputs=[
                io.Image.Output("driving_mask_video"),
                io.Image.Output("ref_mask_image"),
            ],
            is_experimental=True,
        )

    @classmethod
    def execute(cls, driving_track_data, object_indices, sort_by, replacement_mode, ref_track_data=None):
        def _prep(td):
            if sort_by != "none":
                td = _subset_track_data(td, _sort_tracks(td, sort_by))
            if object_indices.strip():
                indices = [int(i.strip()) for i in object_indices.split(",") if i.strip().isdigit()]
                packed = td.get("packed_masks")
                n_obj = packed.shape[1] if packed is not None else 0
                indices = [i for i in indices if 0 <= i < n_obj]
                td = _subset_track_data(td, indices)
            return td

        drv = _prep(driving_track_data)
        mask_video = _render_colored_masks(drv, "white" if replacement_mode else "black")

        if ref_track_data is not None:
            ref = _prep(ref_track_data)
            ref_mask_image = _render_colored_masks(ref, "black")
        else:
            H, W = drv["orig_size"]
            ref_mask_image = torch.zeros(1, H, W, 3, device=comfy.model_management.intermediate_device())

        return io.NodeOutput(mask_video, ref_mask_image)


class SCAIL2Extension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            SCAIL2ColoredMask,
        ]


async def comfy_entrypoint() -> SCAIL2Extension:
    return SCAIL2Extension()