Files

689 lines
24 KiB
Python

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