masks.py

from typing import Any

import torch
from comfy.model_management import is_device_cpu
from comfy_extras.nodes_mask import GrowMask

from .nodes_registry import comfy_node


@comfy_node(name="LTXVPreprocessMasks")
class LTXVPreprocessMasks:
    """
    Preprocesses masks for use with the LTXVideo model's latent masking.

    This node handles temporal mask processing by:
    1. Validating mask dimensions against VAE downscaling factors
    2. Optionally inverting masks
    3. Handling the first frame mask separately (can be ignored)
    4. Combining temporal masks using max pooling
    5. Growing/shrinking masks with morphological operations
    6. Clamping mask values to ensure proper opacity

    The output is a set of masks ready for latent-space masking operations.
    """

    @classmethod
    def INPUT_TYPES(cls: "LTXVPreprocessMasks") -> dict[str, Any]:  # noqa: N802
        return {
            "required": {
                "masks": ("MASK",),
                "vae": ("VAE",),
                "invert_input_masks": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "Invert the input masks before processing."
                        "Useful for masking vs unmasking operations.",
                    },
                ),
                "ignore_first_mask": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Zero out the first mask, typically used when"
                        "it corresponds to the conditioning frame.",
                    },
                ),
                "pooling_method": (
                    ["max", "mean", "min"],
                    {
                        "default": "max",
                        "tooltip": "Method to combine temporal masks. Max preserves strongest values, "
                        "mean averages them, min takes weakest values.",
                    },
                ),
                "grow_mask": (
                    "INT",
                    {
                        "default": 0,
                        "min": -16384,
                        "max": 16384,
                        "step": 1,
                        "tooltip": "Pixels to grow (positive) or shrink (negative) the mask."
                        "Uses morphological operations.",
                    },
                ),
                "tapered_corners": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Use tapered corners in morphological operations for smoother mask edges.",
                    },
                ),
                "clamp_min": (
                    "FLOAT",
                    {
                        "default": 0.5,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Minimum value to clamp masks to, ensuring minimum opacity.",
                    },
                ),
                "clamp_max": (
                    "FLOAT",
                    {
                        "default": 1.0,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Maximum value to clamp masks to.",
                    },
                ),
            },
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "preprocess_masks"
    CATEGORY = "Lightricks/mask_operations"
    DESCRIPTION = (
        "Preprocess masks to be used for masking latents in the LTXVideo model."
    )

    def preprocess_masks(
        self,
        masks: torch.Tensor,
        vae,  # noqa: ANN001
        pooling_method: str,
        grow_mask: int,
        tapered_corners: bool,
        clamp_min: float,
        clamp_max: float,
        ignore_first_mask: bool,
        invert_input_masks: bool,
    ) -> tuple[torch.Tensor]:
        """
        Preprocess masks for temporal latent masking.

        Args:
            masks: Input masks tensor of shape (batch_size, H, W) where batch_size = 1 + (num_frames - 1)
                   The first mask is for the conditioning frame, remaining masks are for generated frames
            vae: VAE model object containing downscale_index_formula attribute
            pooling_method: Method to combine temporal masks ('max', 'mean', 'min')
            grow_mask: Pixels to grow (positive) or shrink (negative) the mask
            tapered_corners: Whether to use tapered corners in morphological operations
            clamp_min: Minimum value to clamp masks to (ensures minimum opacity)
            clamp_max: Maximum value to clamp masks to
            ignore_first_mask: Whether to zero out the first mask (often used for conditioning)
            invert_input_masks: Whether to invert all input masks (1 - mask)

        Returns:
            Tuple containing the processed masks tensor

        Raises:
            ValueError: If mask dimensions are not compatible with VAE downscaling factors
        """

        if masks.ndim != 3:
            raise ValueError("Masks must be of shape (batch_size, H, W).")

        batch_size, H, W = masks.shape

        # Get VAE downscaling factors for temporal and spatial dimensions
        time_scale_factor, width_scale_factor, height_scale_factor = (
            vae.downscale_index_formula
        )

        # Validate that mask dimensions are compatible with VAE downscaling
        # The batch size minus 1 (for the first frame) must be divisible by time_scale_factor
        if (batch_size - 1) % time_scale_factor != 0:
            raise ValueError(
                f"Masks batch size must have a multiple of {time_scale_factor} masks + 1."
            )

        if H % height_scale_factor != 0:
            raise ValueError(
                f"Masks height must have a multiple of {height_scale_factor}."
            )

        if W % width_scale_factor != 0:
            raise ValueError(
                f"Masks width must have a multiple of {width_scale_factor}."
            )

        # Optionally invert all masks (useful for masking vs unmasking operations)
        if invert_input_masks:
            masks = 1.0 - masks

        # Handle first mask separately - often used as conditioning frame
        if ignore_first_mask:
            first_mask = masks[0:1, :, :] * 0  # Zero out the first mask
        else:
            first_mask = masks[0:1, :, :]

        # Process remaining masks by grouping them according to time_scale_factor
        # and taking the maximum value across each temporal group
        remaining = masks[1:].view(-1, time_scale_factor, H, W)

        # Apply selected pooling method
        pooling_methods = {
            "max": lambda x: x.max(dim=1)[0],
            "mean": lambda x: x.mean(dim=1),
            "min": lambda x: x.min(dim=1)[0],
        }

        combined = pooling_methods[pooling_method](remaining)

        # Apply morphological operations to grow or shrink masks
        if grow_mask != 0:
            grow_mask_node = GrowMask()

            # GrowMask node is not supported on GPU, so we need to move the combined tensor to CPU
            # if it is not already on CPU
            input_device = combined.device
            if not is_device_cpu(input_device):
                combined = combined.cpu()

            (combined,) = grow_mask_node.expand_mask(
                mask=combined, expand=grow_mask, tapered_corners=tapered_corners
            )

            # Move the combined tensor back to the original device
            if not is_device_cpu(input_device):
                combined = combined.to(input_device)

        # Combine first mask with processed remaining masks and clamp values
        # Clamping ensures masks have proper opacity values
        processed_masks = torch.cat(
            [first_mask, torch.clamp(combined, clamp_min, clamp_max)], dim=0
        )

        return (processed_masks,)