latents.py

import comfy.utils
import comfy_extras.nodes_lt as nodes_lt
import numpy as np
import torch
from comfy.ldm.lightricks.vae.audio_vae import LATENT_DOWNSAMPLE_FACTOR
from comfy.nested_tensor import NestedTensor

from .iclora_attention import set_last_guide_attention_strength
from .nodes_registry import comfy_node


@comfy_node(name="LTXVSelectLatents")
class LTXVSelectLatents:
    """
    Selects a range of frames from a video latent.

    Features:
    - Supports positive and negative indexing
    - Preserves batch processing capabilities
    - Handles noise masks if present
    - Maintains 5D tensor format
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "samples": ("LATENT",),
                "start_index": (
                    "INT",
                    {"default": 0, "min": -9999, "max": 9999, "step": 1},
                ),
                "end_index": (
                    "INT",
                    {"default": -1, "min": -9999, "max": 9999, "step": 1},
                ),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "select_latents"
    CATEGORY = "latent/video"
    DESCRIPTION = (
        "Selects a range of frames from the video latent. "
        "start_index and end_index define a closed interval (inclusive of both endpoints)."
    )

    def select_latents(self, samples: dict, start_index: int, end_index: int) -> tuple:
        """
        Selects a range of frames from the video latent.

        Args:
            samples (dict): Video latent dictionary
            start_index (int): Starting frame index (supports negative indexing)
            end_index (int): Ending frame index (supports negative indexing)

        Returns:
            tuple: Contains modified latent dictionary with selected frames

        Raises:
            ValueError: If indices are invalid
        """
        try:
            s = samples.copy()
            video_latent = s["samples"]
            batch, channels, frames, height, width = video_latent.shape

            # Handle negative indices
            start_idx = frames + start_index if start_index < 0 else start_index
            end_idx = frames + end_index if end_index < 0 else end_index

            # Validate and clamp indices
            start_idx = max(0, min(start_idx, frames - 1))
            end_idx = max(0, min(end_idx, frames - 1))
            if start_idx > end_idx:
                start_idx = min(start_idx, end_idx)

            # Select frames while maintaining 5D format
            s["samples"] = video_latent[:, :, start_idx : end_idx + 1, :, :]

            # Handle noise mask if present
            if "noise_mask" in s and s["noise_mask"] is not None:
                s["noise_mask"] = s["noise_mask"][:, :, start_idx : end_idx + 1, :, :]

            return (s,)

        except Exception as e:
            print(f"[LTXVSelectLatents] Error: {str(e)}")
            raise


@comfy_node(name="LTXVSetVideoLatentNoiseMasks")
class LTXVSetVideoLatentNoiseMasks:
    """
    Applies multiple masks to a video latent.

    Features:
    - Supports multiple input mask formats (2D, 3D, 4D)
    - Automatically handles fewer masks than frames by reusing the last mask
    - Resizes masks to match latent dimensions
    - Preserves batch processing capabilities

    Input Formats:
    - 2D mask: Single mask [H, W]
    - 3D mask: Multiple masks [M, H, W]
    - 4D mask: Multiple masks with channels [M, C, H, W]
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "samples": ("LATENT",),
                "masks": ("MASK",),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"
    CATEGORY = "latent/video"
    DESCRIPTION = (
        "Applies multiple masks to a video latent. "
        "masks can be 2D, 3D, or 4D tensors. "
        "If there are fewer masks than frames, the last mask will be reused."
    )

    def set_mask(self, samples: dict, masks: torch.Tensor) -> tuple:
        """
        Applies masks to video latent frames.

        Args:
            samples (dict): Video latent dictionary containing 'samples' tensor
            masks (torch.Tensor): Mask tensor in various possible formats
                - 2D: [H, W] single mask
                - 3D: [M, H, W] multiple masks
                - 4D: [M, C, H, W] multiple masks with channels

        Returns:
            tuple: Contains modified latent dictionary with applied masks

        Raises:
            ValueError: If mask dimensions are unsupported
            RuntimeError: If tensor operations fail
        """
        try:
            s = samples.copy()
            video_latent = s["samples"]
            batch_size, channels, num_frames, height, width = video_latent.shape

            # Initialize noise_mask if not present or resize if dimensions don't match
            if "noise_mask" not in s:
                s["noise_mask"] = torch.zeros(
                    (batch_size, 1, num_frames, height, width),
                    dtype=video_latent.dtype,
                    device=video_latent.device,
                )
            else:
                existing_shape = s["noise_mask"].shape
                # Check if noise_mask has the right number of frames
                if existing_shape[2] != num_frames:
                    s["noise_mask"] = torch.zeros(
                        (batch_size, 1, num_frames, height, width),
                        dtype=video_latent.dtype,
                        device=video_latent.device,
                    )

            # Process masks
            masks_reshaped = self._reshape_masks(masks)
            M = masks_reshaped.shape[0]
            resized_masks = self._resize_masks(masks_reshaped, height, width)

            # Apply masks efficiently
            self._apply_masks(s["noise_mask"], resized_masks, num_frames, M)
            return (s,)

        except Exception as e:
            print(f"[LTXVSetVideoLatentNoiseMasks] Error: {str(e)}")
            raise

    def _reshape_masks(self, masks: torch.Tensor) -> torch.Tensor:
        """Reshapes input masks to consistent 4D format."""
        original_shape = tuple(masks.shape)
        ndims = masks.ndim

        if ndims == 2:
            return masks.unsqueeze(0).unsqueeze(0)
        elif ndims == 3:
            return masks.reshape(masks.shape[0], 1, masks.shape[1], masks.shape[2])
        elif ndims == 4:
            return masks.reshape(masks.shape[0], 1, masks.shape[2], masks.shape[3])
        else:
            raise ValueError(
                f"Unsupported 'masks' dimension: {original_shape}. "
                "Must be 2D (H,W), 3D (M,H,W), or 4D (M,C,H,W)."
            )

    def _resize_masks(
        self, masks: torch.Tensor, height: int, width: int
    ) -> torch.Tensor:
        """Resizes all masks to match latent dimensions."""
        return torch.nn.functional.interpolate(
            masks, size=(height, width), mode="bilinear", align_corners=False
        )

    def _apply_masks(
        self,
        noise_mask: torch.Tensor,
        resized_masks: torch.Tensor,
        num_frames: int,
        M: int,
    ) -> None:
        """Applies resized masks to all frames."""
        for f in range(num_frames):
            mask_idx = min(f, M - 1)  # Reuse last mask if we run out
            noise_mask[:, :, f] = resized_masks[mask_idx]


@comfy_node(name="LTXVDilateLatent")
class LTXVDilateLatent:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "latent": ("LATENT",),
                "horizontal_scale": (
                    "INT",
                    {"default": 1, "min": 1, "max": 100, "step": 1},
                ),
                "vertical_scale": (
                    "INT",
                    {"default": 1, "min": 1, "max": 100, "step": 1},
                ),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "dilate_latent"
    CATEGORY = "latent/video"
    DESCRIPTION = "Dilates a latent by a grid size."

    def dilate_latent(
        self, latent: dict, horizontal_scale: int, vertical_scale: int
    ) -> tuple:
        if horizontal_scale == 1 and vertical_scale == 1:
            return (latent,)

        samples = latent["samples"]
        mask = latent.get("noise_mask", None)
        dilated_shape = samples.shape[:3] + (
            samples.shape[3] * vertical_scale,
            samples.shape[4] * horizontal_scale,
        )

        dilated_samples = torch.zeros(
            dilated_shape,
            device=samples.device,
            dtype=samples.dtype,
            requires_grad=False,
        )
        dilated_samples[..., ::vertical_scale, ::horizontal_scale] = samples

        dilated_mask_shape = (
            dilated_samples.shape[0],
            1,
            dilated_samples.shape[2],
            dilated_samples.shape[3],
            dilated_samples.shape[4],
        )
        dilated_mask = torch.full(
            dilated_mask_shape,
            -1.0,
            device=samples.device,
            dtype=samples.dtype,
            requires_grad=False,
        )
        dilated_mask[..., ::vertical_scale, ::horizontal_scale] = (
            mask if mask is not None else 1.0
        )
        latent = {"samples": dilated_samples, "noise_mask": dilated_mask}

        return (latent,)


def add_latent_guide(
    *,
    vae,
    positive,
    negative,
    latent,
    guiding_latent,
    latent_idx,
    strength,
    attention_mask=None,
):
    """Pin an already-encoded latent as a guide, via ComfyUI's native node.

    ComfyUI ships LTXVAddLatentGuide itself (Comfy-Org/ComfyUI#16176), so the
    implementation lives there rather than being duplicated here; this is only the
    call-shape adapter the pack's samplers want -- keyword arguments and a plain
    tuple instead of a NodeOutput.

    Deliberately does not forward pixel_frame_offset: that input exists only in
    Lightricks' ComfyUI fork, and the pack has to keep working against a stock
    install.
    """
    positive, negative, latent = nodes_lt.LTXVAddLatentGuide.execute(
        positive=positive,
        negative=negative,
        vae=vae,
        latent=latent,
        guiding_latent=guiding_latent,
        latent_idx=latent_idx,
        strength=strength,
        attention_mask=attention_mask,
    )
    if attention_mask is None:
        # The core node records `strength` as the guide's attention strength too, so a
        # weakly pinned guide would also become weakly attended. The implementation this
        # replaced always left attention at full strength, and the samplers calling this
        # pass strengths well below 1.0 (0.15 in the tiled sampler, 0.5 when extending),
        # so keep the two decoupled rather than silently reweighting their attention.
        positive = set_last_guide_attention_strength(positive, 1.0)
        negative = set_last_guide_attention_strength(negative, 1.0)
    return positive, negative, latent


@comfy_node(name="LTXVImgToVideoConditionOnly")
class LTXVImgToVideoConditionOnly:
    """
    Applies image conditioning to the first frames of a video latent.

    Features:
    - Takes existing latent and applies image conditioning
    - Automatically resizes image to match latent dimensions
    - Creates noise mask for strength control
    - Supports bypass mode
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "vae": ("VAE",),
                "image": ("IMAGE",),
                "latent": ("LATENT",),
                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
            },
            "optional": {
                "bypass": (
                    "BOOLEAN",
                    {"default": False, "tooltip": "Bypass the conditioning."},
                ),
            },
        }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("latent",)
    CATEGORY = "conditioning/video_models"
    FUNCTION = "generate"
    DESCRIPTION = (
        "Applies image conditioning to the first frames of an existing latent. "
        "Creates a noise mask to control conditioning strength."
    )

    def encode_image(self, image, shape, vae):

        time_scale_factor, height_scale_factor, width_scale_factor = (
            vae.downscale_index_formula
        )
        batch, channels, frames, height, width = shape
        width *= width_scale_factor
        height *= height_scale_factor

        # Resize image to match latent dimensions
        if image.shape[1] != height or image.shape[2] != width:
            pixels = comfy.utils.common_upscale(
                image.movedim(-1, 1), width, height, "bilinear", "center"
            ).movedim(1, -1)
        else:
            pixels = image

        # Encode image (only RGB channels)
        encode_pixels = pixels[:, :, :, :3]
        t = vae.encode(encode_pixels)
        return t

    def generate(self, image, vae, latent, strength, bypass=False):
        if bypass:
            return (latent,)

        samples = latent["samples"]
        t = self.encode_image(image, samples.shape, vae)
        # Apply encoded image to first frames
        samples[:, :, : t.shape[2]] = t

        # Create noise mask for conditioning
        conditioning_latent_frames_mask = torch.ones(
            (1, 1, samples.shape[2], 1, 1),
            dtype=torch.float32,
            device=samples.device,
        )
        conditioning_latent_frames_mask[:, :, : t.shape[2]] = 1.0 - strength

        return ({"samples": samples, "noise_mask": conditioning_latent_frames_mask},)


def get_video_latent_blend_coefficients(
    video_frame_index_start,
    video_frame_index_end,
    video_frame_count,
    slope_len=3,
):
    """
    Returns a blend coefficient list shaped such that:
    - It is 0.0 outside the range [video_frame_index_start, video_frame_index_end]
    - It ramps up from 0.0 to 1.0 starting at video_frame_index_start-slope_len to video_frame_index_start
    - It stays at 1.0 during [video_frame_index_start, video_frame_index_end]
    - It ramps down from 1.0 to 0.0 starting at video_frame_index_end to video_frame_index_end+slope_len
    - The 'slope_len' parameter controls the sharpness (in frames) of the slope.
    """

    coeffs = [0.0] * video_frame_count

    # Clamp arguments to safe range
    video_frame_index_start = max(
        0, min(video_frame_count - 1, video_frame_index_start)
    )
    video_frame_index_end = max(
        video_frame_index_start, min(video_frame_count - 1, video_frame_index_end)
    )
    slope_len = max(1, slope_len)

    # Ramp up
    rampl_start = max(0, video_frame_index_start - slope_len)
    for i in range(rampl_start, video_frame_index_start):
        # Linear ramp, but could be smoother (sigmoid/cosine)
        coeffs[i] = (i - rampl_start + 1) / slope_len

    # Plateau
    for i in range(video_frame_index_start, video_frame_index_end + 1):
        coeffs[i] = 1.0

    # Ramp down
    rampr_end = min(video_frame_count, video_frame_index_end + slope_len + 1)
    for i in range(video_frame_index_end + 1, rampr_end):
        coeffs[i] = 1.0 - ((i - (video_frame_index_end + 1) + 1) / slope_len)
        coeffs[i] = max(0.0, coeffs[i])

    import numpy as np

    num_coeffs = len(coeffs)
    pixel_frame_length = (num_coeffs - 1) * 8 + 1

    xp = np.array([0] + list(range(1, pixel_frame_length, 8)))
    fp = np.array(coeffs)

    # Calculate the upsampled coefficients using np.interp
    pixel_frame_positions = np.arange(pixel_frame_length)
    pixel_frame_coefficients = np.interp(pixel_frame_positions, xp, fp).tolist()

    return coeffs, pixel_frame_coefficients


@comfy_node(description="LTXV Set Audio Video Mask By Time")
class LTXVSetAudioVideoMaskByTime:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "av_latent": ("LATENT",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "model": ("MODEL",),
                "vae": ("VAE",),
                "audio_vae": ("VAE",),
                "start_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2000.0}),
                "end_time": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 2000.0}),
                "video_fps": ("FLOAT", {"default": 24.0, "min": 0.0, "max": 500.0}),
                "mask_video": ("BOOLEAN", {"default": True}),
                "mask_audio": ("BOOLEAN", {"default": True}),
                "mask_init_value_video": (
                    "FLOAT",
                    {"default": 0.0, "min": 0.0, "max": 1.0},
                ),
                "mask_init_value_audio": (
                    "FLOAT",
                    {"default": 0.0, "min": 0.0, "max": 1.0},
                ),
                "slope_len": (
                    "INT",
                    {"default": 3, "min": 1, "max": 100, "step": 1},
                ),
            },
            "optional": {
                "spatial_mask": (
                    "MASK",
                    {"default": None, "tooltip": "Spatial mask."},
                ),
            },
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "FLOAT", "FLOAT")
    RETURN_NAMES = (
        "positive",
        "negative",
        "av_latent",
        "video_latent_blend_coefficients",
        "video_pixel_blend_coefficients",
    )

    FUNCTION = "run"
    CATEGORY = "utility"
    DESCRIPTION = "Sets the audio and video mask by time."

    def run(
        self,
        av_latent,
        positive,
        negative,
        model,
        vae,
        audio_vae,
        start_time,
        end_time,
        video_fps,
        mask_video,
        mask_audio,
        mask_init_value_video,
        mask_init_value_audio,
        slope_len,
        spatial_mask=None,
    ):
        from comfy.ldm.lightricks.av_model import LTXAVModel

        if model.model.diffusion_model.__class__.__name__ != "LTXAVModel":
            raise ValueError()

        ltxav: LTXAVModel = model.model.diffusion_model

        # Extract configuration from the audio VAE
        sampling_rate = audio_vae.autoencoder.sampling_rate
        mel_hop_length = audio_vae.autoencoder.mel_hop_length
        audio_latents_per_second = (
            sampling_rate / mel_hop_length / LATENT_DOWNSAMPLE_FACTOR
        )

        time_scale_factor = vae.downscale_index_formula[0]
        video_latents_per_second = video_fps / time_scale_factor

        if not isinstance(av_latent["samples"], NestedTensor):
            raise ValueError("av_latent must be a NestedTensor")

        video_samples, audio_samples = ltxav.separate_audio_and_video_latents(
            av_latent["samples"].tensors,
            None,
        )
        video_mask = torch.full(
            (
                video_samples.shape[0],
                video_samples.shape[1],
                video_samples.shape[2],
                video_samples.shape[3],
                video_samples.shape[4],
            ),
            fill_value=mask_init_value_video,
        )
        audio_mask = torch.full(
            (
                audio_samples.shape[0],
                audio_samples.shape[1],
                audio_samples.shape[2],
                audio_samples.shape[3],
            ),
            fill_value=mask_init_value_audio,
        )
        if spatial_mask is not None:
            if spatial_mask.ndim == 3:
                spatial_mask = spatial_mask.unsqueeze(0)
            if spatial_mask.ndim == 2:
                spatial_mask = spatial_mask.unsqueeze(0).unsqueeze(0)
            spatial_mask = torch.nn.functional.interpolate(
                spatial_mask,
                size=(video_samples.shape[3], video_samples.shape[4]),
                mode="bilinear",
                align_corners=False,
            )
        video_latent_frame_count = video_samples.shape[2]
        audio_latent_frame_count = audio_samples.shape[2]
        video_pixel_frame_count = (video_latent_frame_count - 1) * time_scale_factor + 1
        xp = np.array(
            [0]
            + list(
                range(1, video_pixel_frame_count + time_scale_factor, time_scale_factor)
            )
        )
        video_pixel_frame_start_raw = int(round(start_time * video_fps))
        # video_frame_index_start = index of the value in xp rounding up
        video_latent_frame_index_start = np.searchsorted(
            xp, video_pixel_frame_start_raw, side="left"
        )
        video_pixel_frame_end_raw = int(round(end_time * video_fps))
        # video_frame_index_end = index of the value in xp rounding down
        video_latent_frame_index_end = (
            np.searchsorted(xp, video_pixel_frame_end_raw, side="right") - 1
        )
        audio_latent_frame_index_start = int(
            round(start_time * audio_latents_per_second)
        )
        audio_latent_frame_index_end = (
            int(round(end_time * audio_latents_per_second)) + 1
        )
        # clamping
        video_latent_frame_index_start = max(0, video_latent_frame_index_start)
        video_latent_frame_index_end = min(
            video_latent_frame_index_end, video_latent_frame_count
        )
        audio_latent_frame_index_start = max(0, audio_latent_frame_index_start)
        audio_latent_frame_index_end = min(
            audio_latent_frame_index_end, audio_latent_frame_count
        )
        print(
            "noise mask start and end indices: video (%d %d), audio (%d %d), video fps: %f, video_latents_per_second: %f, audio_latents_per_second: %f, "
            "video_latent_frame_count: %d, video_pixel_frame_count: %d, video_pixel_frame_start_raw: %d, video_pixel_frame_end_raw: %d, start_time: %f, end_time: %f"
            % (
                video_latent_frame_index_start,
                video_latent_frame_index_end,
                audio_latent_frame_index_start,
                audio_latent_frame_index_end,
                video_fps,
                video_latents_per_second,
                audio_latents_per_second,
                video_latent_frame_count,
                video_pixel_frame_count,
                video_pixel_frame_start_raw,
                video_pixel_frame_end_raw,
                start_time,
                end_time,
            )
        )
        if mask_video:
            if spatial_mask is not None:
                video_mask[
                    :,
                    :,
                    video_latent_frame_index_start:video_latent_frame_index_end,
                    :,
                    :,
                ] = spatial_mask
            else:
                video_mask[
                    :, :, video_latent_frame_index_start:video_latent_frame_index_end
                ] = 1.0
        if mask_audio:
            audio_mask[
                :, :, audio_latent_frame_index_start:audio_latent_frame_index_end
            ] = 1.0

        if "noise_mask" in av_latent:
            base_mask = av_latent["noise_mask"].tensors[0].clone()
            if (
                base_mask.shape[0]
                == base_mask.shape[1]
                == 1
                == base_mask.shape[3]
                == base_mask.shape[4]
            ):
                for frame in range(base_mask.shape[2]):
                    video_mask[:, :, frame, :, :] *= base_mask[0, 0, frame, 0, 0]

        av_latent["noise_mask"] = NestedTensor(
            ltxav.recombine_audio_and_video_latents(video_mask, audio_mask)
        )

        video_latent_blend_coefficients, video_pixel_blend_coefficients = (
            get_video_latent_blend_coefficients(
                video_latent_frame_index_start,
                video_latent_frame_index_end,
                video_latent_frame_count,
                slope_len=slope_len,
            )
        )
        return (
            positive,
            negative,
            av_latent,
            video_latent_blend_coefficients,
            video_pixel_blend_coefficients,
        )