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, )