XIAOTsune-xt-matte-toolbox-comfyui-node/nodes.py

"""ComfyUI custom nodes for XT Matte Toolbox core matting."""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Callable, Dict, Optional, Sequence, Tuple

import numpy as np
import torch
from PIL import Image

_RESAMPLING = getattr(Image, "Resampling", Image)


FALLBACK_MODEL_CHOICES: Tuple[str, ...] = (
    "General",
    "General-Lite",
    "General-Lite-2K",
    "Matting",
    "Portrait",
    "DIS",
    "HRSOD",
    "COD",
    "DIS-TR_TEs",
    "General-legacy",
)

_BOOTSTRAP_ERROR: Optional[str] = None
_PROJECT_ROOT: Optional[Path] = None
_MODEL_CHOICES: Tuple[str, ...] = FALLBACK_MODEL_CHOICES
_APPLY_BACKGROUND_REPLACEMENT: Optional[Callable] = None


def _find_project_root() -> Path:
    """Locate project root that contains src/models/loader.py."""
    env_root = os.environ.get("XT_MATTE_PROJECT_ROOT")
    if env_root:
        root = Path(env_root).expanduser().resolve()
        if (root / "src" / "models" / "loader.py").exists():
            return root
        raise RuntimeError(
            f"XT_MATTE_PROJECT_ROOT is set but invalid: {root}. "
            "Expected src/models/loader.py under this path."
        )

    here = Path(__file__).resolve()
    for parent in [here.parent] + list(here.parents):
        if (parent / "src" / "models" / "loader.py").exists():
            return parent

    raise RuntimeError(
        "Project root not found. Set env XT_MATTE_PROJECT_ROOT to your project path."
    )


def _bootstrap() -> None:
    """Wire this node to your existing matting core under src/."""
    global _BOOTSTRAP_ERROR
    global _PROJECT_ROOT
    global _MODEL_CHOICES
    global _APPLY_BACKGROUND_REPLACEMENT

    try:
        _PROJECT_ROOT = _find_project_root()

        # Keep core model/output directories in your project root, not ComfyUI cwd.
        os.environ.setdefault("XT_MATTE_BASE_DIR", str(_PROJECT_ROOT))

        root_str = str(_PROJECT_ROOT)
        if root_str not in sys.path:
            sys.path.insert(0, root_str)

        from src.config import USAGE_TO_WEIGHTS_FILE
        from src.processors.image_processor import apply_background_replacement

        _MODEL_CHOICES = tuple(USAGE_TO_WEIGHTS_FILE.keys()) or FALLBACK_MODEL_CHOICES
        _APPLY_BACKGROUND_REPLACEMENT = apply_background_replacement
        _BOOTSTRAP_ERROR = None
    except Exception as exc:  # noqa: BLE001
        _BOOTSTRAP_ERROR = str(exc)


_bootstrap()


def _ensure_bootstrap_ok() -> None:
    if _BOOTSTRAP_ERROR is not None or _APPLY_BACKGROUND_REPLACEMENT is None:
        raise RuntimeError(
            "XT Matte node bootstrap failed. "
            f"Details: {_BOOTSTRAP_ERROR or 'Unknown error'}"
        )


def _to_uint8_rgb(image_tensor: torch.Tensor) -> np.ndarray:
    img = image_tensor.detach().cpu().numpy()
    if img.ndim != 3 or img.shape[2] < 3:
        raise ValueError(f"Invalid IMAGE shape: {img.shape}")
    img = np.clip(img[..., :3], 0.0, 1.0)
    return (img * 255.0 + 0.5).astype(np.uint8)


def _pick_roi_mask(
    roi_mask: Optional[torch.Tensor], batch_index: int, target_hw: Tuple[int, int]
) -> Optional[np.ndarray]:
    if roi_mask is None:
        return None

    h, w = target_hw
    mask = roi_mask.detach().cpu().float()
    if mask.ndim == 2:
        selected = mask
    elif mask.ndim == 3:
        selected = mask[batch_index] if batch_index < mask.shape[0] else mask[0]
    else:
        raise ValueError(f"Invalid roi_mask shape: {tuple(mask.shape)}")

    m = torch.clamp(selected, 0.0, 1.0).numpy().astype(np.float32)
    if m.shape != (h, w):
        # Nearest keeps user strokes hard-edged.
        resized = Image.fromarray((m * 255.0 + 0.5).astype(np.uint8)).resize(
            (w, h), resample=_RESAMPLING.NEAREST
        )
        m = np.asarray(resized, dtype=np.float32) / 255.0
    return m


class XTMatteCutoutNode:
    """
    Core matting node:
    - Input IMAGE
    - Output foreground IMAGE + alpha MASK + mask preview IMAGE
    """

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Dict[str, Sequence]]:
        return {
            "required": {
                "image": ("IMAGE",),
                "model_name": (_MODEL_CHOICES,),
                "resolution": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}),
                "semi_transparent": ("BOOLEAN", {"default": False}),
                "semi_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "defringe": ("BOOLEAN", {"default": False}),
                "defringe_strength": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}),
                "roi_crop_before": ("BOOLEAN", {"default": True}),
                "roi_pad_px": ("INT", {"default": 16, "min": 0, "max": 256, "step": 1}),
            },
            "optional": {
                "roi_mask": ("MASK",),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
    RETURN_NAMES = ("foreground", "alpha_mask", "mask_preview")
    FUNCTION = "matte"
    CATEGORY = "XT/Matting"

    def matte(
        self,
        image: torch.Tensor,
        model_name: str,
        resolution: int,
        semi_transparent: bool,
        semi_strength: float,
        defringe: bool,
        defringe_strength: float,
        roi_crop_before: bool,
        roi_pad_px: int,
        roi_mask: Optional[torch.Tensor] = None,
    ):
        _ensure_bootstrap_ok()
        assert _APPLY_BACKGROUND_REPLACEMENT is not None

        if image.ndim != 4:
            raise ValueError(f"Expected IMAGE shape [B,H,W,C], got {tuple(image.shape)}")

        batch = image.shape[0]
        fg_list = []
        alpha_list = []
        preview_list = []

        for i in range(batch):
            rgb_u8 = _to_uint8_rgb(image[i])
            h, w = rgb_u8.shape[:2]
            roi = _pick_roi_mask(roi_mask, i, (h, w))

            result, mask_preview, err = _APPLY_BACKGROUND_REPLACEMENT(
                image=rgb_u8,
                model_name=model_name,
                input_size=(int(resolution), int(resolution)),
                semi_transparent=bool(semi_transparent),
                semi_strength=float(semi_strength),
                semi_mode="auto",
                remove_white_halo=bool(defringe),
                defringe_strength=float(defringe_strength),
                roi_mask_fullres=roi,
                roi_crop_before=bool(roi_crop_before),
                roi_pad_px=int(roi_pad_px),
            )

            if result is None:
                raise RuntimeError(f"Matting failed on batch index {i}: {err}")

            rgba = np.asarray(result.convert("RGBA"), dtype=np.uint8)
            alpha = (rgba[:, :, 3].astype(np.float32) / 255.0).copy()
            fg = (rgba[:, :, :3].astype(np.float32) / 255.0).copy()

            preview_img = (
                mask_preview.convert("RGB")
                if isinstance(mask_preview, Image.Image)
                else Image.fromarray((alpha * 255.0 + 0.5).astype(np.uint8)).convert("RGB")
            )
            preview = (np.asarray(preview_img, dtype=np.float32) / 255.0).copy()

            fg_list.append(torch.from_numpy(fg))
            alpha_list.append(torch.from_numpy(alpha))
            preview_list.append(torch.from_numpy(preview))

        foreground = torch.stack(fg_list, dim=0)
        alpha_mask = torch.stack(alpha_list, dim=0)
        mask_preview_tensor = torch.stack(preview_list, dim=0)
        return (foreground, alpha_mask, mask_preview_tensor)


NODE_CLASS_MAPPINGS = {
    "XTMatteCutout": XTMatteCutoutNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "XTMatteCutout": "XT Matte Cutout",
}