831 lines
28 KiB
Python
831 lines
28 KiB
Python
from typing import Optional
|
|
|
|
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 .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="LTXVAddLatents")
|
|
class LTXVAddLatents:
|
|
"""
|
|
Concatenates two video latents along the frames dimension.
|
|
|
|
Features:
|
|
- Validates dimension compatibility
|
|
- Handles device placement
|
|
- Preserves noise masks with proper handling
|
|
- Supports batch processing
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latents1": ("LATENT",),
|
|
"latents2": ("LATENT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "add_latents"
|
|
CATEGORY = "latent/video"
|
|
DESCRIPTION = (
|
|
"Concatenates two video latents along the frames dimension. "
|
|
"latents1 and latents2 must have the same dimensions except for the frames dimension."
|
|
)
|
|
|
|
def add_latents(
|
|
self, latents1: torch.Tensor, latents2: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Concatenates two video latents along the frames dimension.
|
|
|
|
Args:
|
|
latents1 (dict): First video latent dictionary
|
|
latents2 (dict): Second video latent dictionary
|
|
|
|
Returns:
|
|
tuple: Contains concatenated latent dictionary
|
|
|
|
Raises:
|
|
ValueError: If latent dimensions don't match
|
|
RuntimeError: If tensor operations fail
|
|
"""
|
|
try:
|
|
s = latents1.copy()
|
|
video_latent1 = latents1["samples"]
|
|
video_latent2 = latents2["samples"]
|
|
|
|
# Ensure tensors are on the same device
|
|
target_device = video_latent1.device
|
|
video_latent2 = video_latent2.to(target_device)
|
|
|
|
# Validate dimensions
|
|
self._validate_dimensions(video_latent1, video_latent2)
|
|
|
|
# Concatenate along frames dimension
|
|
s["samples"] = torch.cat([video_latent1, video_latent2], dim=2)
|
|
|
|
# Handle noise masks
|
|
s["noise_mask"] = self._merge_noise_masks(
|
|
latents1, latents2, video_latent1.shape[2], video_latent2.shape[2]
|
|
)
|
|
|
|
return (s,)
|
|
|
|
except Exception as e:
|
|
print(f"[LTXVAddLatents] Error: {str(e)}")
|
|
raise
|
|
|
|
def _validate_dimensions(self, latent1: torch.Tensor, latent2: torch.Tensor):
|
|
"""Validates that latent dimensions match except for frames."""
|
|
# video latents
|
|
if latent1.ndim == 5 and latent2.ndim == 5:
|
|
b1, c1, f1, h1, w1 = latent1.shape
|
|
b2, c2, f2, h2, w2 = latent2.shape
|
|
|
|
if not (b1 == b2 and c1 == c2 and h1 == h2 and w1 == w2):
|
|
raise ValueError(
|
|
f"Latent dimensions must match (except frames dimension).\n"
|
|
f"Got shapes {latent1.shape} and {latent2.shape}"
|
|
)
|
|
# audio latents
|
|
elif latent1.ndim == 4 and latent2.ndim == 4:
|
|
b1, c1, f1, s1 = latent1.shape
|
|
b2, c2, f2, s2 = latent2.shape
|
|
|
|
if not (b1 == b2 and c1 == c2 and s1 == s2):
|
|
raise ValueError(
|
|
f"Latent dimensions must match (except frames dimension).\n"
|
|
f"Got shapes {latent1.shape} and {latent2.shape}"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Latent dimensions must be 4 (audio) or 5 (video) for both inputs.\n"
|
|
f"Got shapes {latent1.shape} and {latent2.shape}"
|
|
)
|
|
|
|
def _merge_noise_masks(
|
|
self, latents1: torch.Tensor, latents2: torch.Tensor, frames1: int, frames2: int
|
|
) -> Optional[torch.Tensor]:
|
|
"""Merges noise masks from both latents with proper handling."""
|
|
if "noise_mask" in latents1 and "noise_mask" in latents2:
|
|
return torch.cat([latents1["noise_mask"], latents2["noise_mask"]], dim=2)
|
|
elif "noise_mask" in latents1 and latents1["noise_mask"] is not None:
|
|
zeros = torch.zeros_like(latents1["noise_mask"][:, :, :frames2, :, :])
|
|
return torch.cat([latents1["noise_mask"], zeros], dim=2)
|
|
elif "noise_mask" in latents2 and latents2["noise_mask"] is not None:
|
|
zeros = torch.zeros_like(latents2["noise_mask"][:, :, :frames1, :, :])
|
|
return torch.cat([zeros, latents2["noise_mask"]], dim=2)
|
|
return None
|
|
|
|
|
|
@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,)
|
|
|
|
|
|
@comfy_node(name="LTXVAddLatentGuide")
|
|
class LTXVAddLatentGuide:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent": ("LATENT",),
|
|
"guiding_latent": ("LATENT",),
|
|
"latent_idx": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": -9999,
|
|
"max": 9999,
|
|
"step": 1,
|
|
"tooltip": "Latent index to start the conditioning at. Can be negative to"
|
|
"indicate that the conditioning is on the frames before the latent.",
|
|
},
|
|
),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive", "negative", "latent")
|
|
|
|
CATEGORY = "ltxtricks"
|
|
FUNCTION = "generate"
|
|
|
|
DESCRIPTION = "Adds a keyframe or a video segment at a specific frame index."
|
|
|
|
def generate(
|
|
self, vae, positive, negative, latent, guiding_latent, latent_idx, strength
|
|
):
|
|
noise_mask = nodes_lt.get_noise_mask(latent)
|
|
latent = latent["samples"]
|
|
guide = guiding_latent["samples"]
|
|
|
|
assert (
|
|
latent.shape[4] % guide.shape[4] == 0
|
|
and latent.shape[3] % guide.shape[3] == 0
|
|
), "The ratio of the height and width of the latents and optional_guiding_latents must be an integer"
|
|
|
|
guiding_latent = LTXVDilateLatent().dilate_latent(
|
|
guiding_latent,
|
|
horizontal_scale=latent.shape[4] // guide.shape[4],
|
|
vertical_scale=latent.shape[3] // guide.shape[3],
|
|
)[0]
|
|
|
|
guide = guiding_latent["samples"]
|
|
guide_mask = guiding_latent.get("noise_mask", None)
|
|
|
|
scale_factors = vae.downscale_index_formula
|
|
|
|
if latent_idx <= 0:
|
|
frame_idx = latent_idx * scale_factors[0]
|
|
else:
|
|
frame_idx = 1 + (latent_idx - 1) * scale_factors[0]
|
|
|
|
positive, negative, latent, noise_mask = nodes_lt.LTXVAddGuide.append_keyframe(
|
|
positive=positive,
|
|
negative=negative,
|
|
frame_idx=frame_idx,
|
|
latent_image=latent,
|
|
noise_mask=noise_mask,
|
|
guiding_latent=guide,
|
|
strength=strength,
|
|
scale_factors=scale_factors,
|
|
guide_mask=guide_mask,
|
|
)
|
|
|
|
return (
|
|
positive,
|
|
negative,
|
|
{"samples": latent, "noise_mask": noise_mask},
|
|
)
|
|
|
|
|
|
@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
|
|
|
|
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,
|
|
)
|