402 lines
14 KiB
Python
402 lines
14 KiB
Python
"""
|
||
LTX First Last Frame Control with Middle Frame and Context Support
|
||
|
||
包含三个节点:
|
||
1. LTXVContext_TTP: 视频续接context节点
|
||
2. LTXVMiddleFrame_TTP: 可串联的中间帧节点
|
||
3. LTXVFirstLastFrameControl_TTP: 首尾帧控制节点
|
||
|
||
Author: TTP
|
||
"""
|
||
|
||
import torch
|
||
import comfy.utils
|
||
from comfy_api.latest import io
|
||
|
||
|
||
class LTXVContext_TTP:
|
||
"""
|
||
视频续接Context节点
|
||
|
||
功能:
|
||
- 从上一个视频提取结尾帧作为context
|
||
- 应用到新视频的开头
|
||
- 实现平滑续接
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"previous_video": ("IMAGE",), # 上一个视频（IMAGE类型，batch维度是帧）
|
||
"vae": ("VAE",),
|
||
"latent": ("LATENT",), # 新视频的空白latent
|
||
"context_latent_frames": ("INT", {
|
||
"default": 6,
|
||
"min": 2,
|
||
"max": 20,
|
||
"step": 1,
|
||
"tooltip": "Context的latent帧数（推荐6）\n"
|
||
"注：latent帧0=1原始帧，其余=8原始帧\n"
|
||
"6帧 = 1+40 = 41原始帧 ≈ 1.64秒@25fps"
|
||
}),
|
||
},
|
||
"optional": {
|
||
"context_strength": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.05,
|
||
"tooltip": "Context固定强度（1.0=完全固定）"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
RETURN_NAMES = ("latent",)
|
||
FUNCTION = "apply_context"
|
||
CATEGORY = "conditioning/video_models"
|
||
|
||
def apply_context(self, previous_video, vae, latent, context_latent_frames, context_strength=1.0):
|
||
"""
|
||
从上一个视频提取context并应用到新latent
|
||
|
||
Args:
|
||
previous_video: 上一个视频的帧（IMAGE格式）
|
||
vae: VAE模型
|
||
latent: 新视频的空白latent
|
||
context_latent_frames: 要使用的context latent帧数
|
||
context_strength: context强度
|
||
|
||
Returns:
|
||
包含context的latent
|
||
"""
|
||
samples = latent["samples"].clone()
|
||
batch, channels, latent_frames, latent_height, latent_width = samples.shape
|
||
|
||
# 获取VAE缩放因子
|
||
_, height_scale_factor, width_scale_factor = vae.downscale_index_formula
|
||
target_width = latent_width * width_scale_factor
|
||
target_height = latent_height * height_scale_factor
|
||
|
||
# 计算需要多少原始帧来生成context_latent_frames个latent帧
|
||
# LTX的公式: latent_frames = ((original_frames - 1) // 8) + 1
|
||
# 反推: original_frames = (latent_frames - 1) * 8 + 1
|
||
# 8N+1结构：第0帧独立，其余每8帧压缩为1个latent帧
|
||
required_frames = (context_latent_frames - 1) * 8 + 1
|
||
|
||
# 从previous_video提取最后N帧
|
||
total_video_frames = previous_video.shape[0]
|
||
start_idx = max(0, total_video_frames - required_frames)
|
||
context_frames = previous_video[start_idx:]
|
||
|
||
# 确保帧数符合要求（可能不足required_frames）
|
||
actual_frames = context_frames.shape[0]
|
||
|
||
# 调整图像尺寸
|
||
if context_frames.shape[1] != target_height or context_frames.shape[2] != target_width:
|
||
pixels = comfy.utils.common_upscale(
|
||
context_frames.movedim(-1, 1),
|
||
target_width,
|
||
target_height,
|
||
"bilinear",
|
||
"center"
|
||
).movedim(1, -1)
|
||
else:
|
||
pixels = context_frames
|
||
|
||
# 只取RGB通道
|
||
encode_pixels = pixels[:, :, :, :3]
|
||
|
||
# VAE编码
|
||
context_latent = vae.encode(encode_pixels)
|
||
actual_latent_frames = context_latent.shape[2]
|
||
|
||
# 将context嵌入到新latent的开头
|
||
embed_frames = min(actual_latent_frames, latent_frames)
|
||
samples[:, :, :embed_frames] = context_latent[:, :, :embed_frames]
|
||
|
||
# 设置noise_mask
|
||
noise_mask = torch.ones(
|
||
(batch, 1, latent_frames, 1, 1),
|
||
dtype=torch.float32,
|
||
device=samples.device,
|
||
)
|
||
|
||
# Context帧固定
|
||
noise_mask[:, :, :embed_frames] = 1.0 - context_strength
|
||
|
||
return ({"samples": samples, "noise_mask": noise_mask},)
|
||
|
||
|
||
class LTXVMiddleFrame_TTP:
|
||
"""
|
||
中间帧节点 - 可串联（传统ComfyUI格式）
|
||
|
||
使用传统的INPUT_TYPES和RETURN_TYPES定义，支持anytype"*"
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"position": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
||
},
|
||
"optional": {
|
||
"middle_frames": ("*",), # anytype - 接受任何类型
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("*",) # anytype - 返回任何类型
|
||
RETURN_NAMES = ("middle_frames",)
|
||
FUNCTION = "execute"
|
||
CATEGORY = "conditioning/video_models"
|
||
|
||
def execute(self, image, position, strength, middle_frames=None):
|
||
"""
|
||
执行中间帧处理
|
||
|
||
Args:
|
||
image: 中间帧图像
|
||
position: 位置(0.0-1.0)
|
||
strength: 强度
|
||
middle_frames: 上一个节点的中间帧数据（可选）
|
||
|
||
Returns:
|
||
累积的中间帧数据
|
||
"""
|
||
# 初始化或累积中间帧列表
|
||
if middle_frames is None:
|
||
frames_list = []
|
||
else:
|
||
# 复制现有列表
|
||
frames_list = list(middle_frames.get("frames", []))
|
||
|
||
# 添加新的中间帧信息
|
||
frames_list.append({
|
||
"image": image,
|
||
"position": position,
|
||
"strength": strength,
|
||
})
|
||
|
||
# 返回累积的数据
|
||
return ({"frames": frames_list},)
|
||
|
||
|
||
class LTXVFirstLastFrameControl_TTP(io.ComfyNode):
|
||
"""
|
||
首尾帧控制节点（io.Schema格式）
|
||
|
||
功能:
|
||
- 控制首帧和尾帧
|
||
- 接受中间帧数据并融合
|
||
- 输出包含所有关键帧的latent
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
return io.Schema(
|
||
node_id="LTXVFirstLastFrameControl_TTP",
|
||
display_name="LTX First Last Frame Control (TTP)",
|
||
category="conditioning/video_models",
|
||
description="控制LTX视频生成的首帧和尾帧，支持中间帧融合",
|
||
inputs=[
|
||
io.Vae.Input("vae"),
|
||
io.Latent.Input("latent"),
|
||
io.Image.Input("first_image", optional=True, tooltip="首帧图像(可选)"),
|
||
io.Image.Input("last_image", optional=True, tooltip="尾帧图像(可选)"),
|
||
io.Float.Input(
|
||
"first_strength",
|
||
default=1.0,
|
||
min=0.0,
|
||
max=1.0,
|
||
step=0.05,
|
||
tooltip="首帧嵌入强度 (1.0=完全替换, 0.0=不嵌入)"
|
||
),
|
||
io.Float.Input(
|
||
"last_strength",
|
||
default=1.0,
|
||
min=0.0,
|
||
max=1.0,
|
||
step=0.05,
|
||
tooltip="尾帧嵌入强度 (1.0=完全替换, 0.0=不嵌入)"
|
||
),
|
||
],
|
||
# 不定义middle_frames输入，改用可选参数在execute中接收
|
||
outputs=[
|
||
io.Latent.Output(display_name="latent"),
|
||
],
|
||
)
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
"""添加传统格式的INPUT_TYPES以支持middle_frames anytype"""
|
||
base_inputs = {
|
||
"required": {
|
||
"vae": ("VAE",),
|
||
"latent": ("LATENT",),
|
||
"first_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
||
"last_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
||
},
|
||
"optional": {
|
||
"first_image": ("IMAGE",),
|
||
"last_image": ("IMAGE",),
|
||
"middle_frames": ("*",), # anytype
|
||
}
|
||
}
|
||
return base_inputs
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
FUNCTION = "execute"
|
||
CATEGORY = "conditioning/video_models"
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
vae,
|
||
latent,
|
||
first_strength=1.0,
|
||
last_strength=1.0,
|
||
first_image=None,
|
||
last_image=None,
|
||
middle_frames=None
|
||
) -> io.NodeOutput:
|
||
"""
|
||
执行首尾帧控制和中间帧融合
|
||
|
||
Args:
|
||
vae: VAE模型
|
||
latent: 输入的视频latent
|
||
first_strength: 首帧强度
|
||
last_strength: 尾帧强度
|
||
first_image: 首帧图像(可选)
|
||
last_image: 尾帧图像(可选)
|
||
middle_frames: 中间帧数据(可选)
|
||
|
||
Returns:
|
||
包含处理后latent和noise_mask的字典
|
||
"""
|
||
# 如果没有提供任何图像和中间帧，直接返回
|
||
has_middle = middle_frames is not None and len(middle_frames.get("frames", [])) > 0
|
||
if first_image is None and last_image is None and not has_middle:
|
||
return (latent,)
|
||
|
||
# 获取latent信息
|
||
samples = latent["samples"]
|
||
batch, _, latent_frames, latent_height, latent_width = samples.shape
|
||
|
||
# 获取VAE的缩放因子
|
||
_, height_scale_factor, width_scale_factor = vae.downscale_index_formula
|
||
width = latent_width * width_scale_factor
|
||
height = latent_height * height_scale_factor
|
||
|
||
# 初始化noise_mask (全1表示不mask)
|
||
noise_mask = torch.ones(
|
||
(batch, 1, latent_frames, 1, 1),
|
||
dtype=torch.float32,
|
||
device=samples.device,
|
||
)
|
||
|
||
# 处理首帧
|
||
if first_image is not None and first_strength > 0.0:
|
||
first_latent = cls._encode_image(vae, first_image, height, width)
|
||
first_latent_frames = first_latent.shape[2]
|
||
|
||
# 嵌入首帧
|
||
samples[:, :, :first_latent_frames] = first_latent
|
||
|
||
# 设置首帧的noise_mask
|
||
noise_mask[:, :, :first_latent_frames] = 1.0 - first_strength
|
||
|
||
# 处理尾帧
|
||
if last_image is not None and last_strength > 0.0:
|
||
last_latent = cls._encode_image(vae, last_image, height, width)
|
||
last_latent_frames = last_latent.shape[2]
|
||
|
||
# 计算尾帧起始位置
|
||
last_start_idx = latent_frames - last_latent_frames
|
||
if last_start_idx < 0:
|
||
last_latent = last_latent[:, :, :latent_frames]
|
||
last_start_idx = 0
|
||
last_latent_frames = latent_frames
|
||
|
||
# 嵌入尾帧
|
||
samples[:, :, last_start_idx:] = last_latent
|
||
|
||
# 设置尾帧的noise_mask
|
||
noise_mask[:, :, last_start_idx:] = 1.0 - last_strength
|
||
|
||
# 处理中间帧（如果有）
|
||
if has_middle:
|
||
frames_list = middle_frames["frames"]
|
||
for frame_data in frames_list:
|
||
image = frame_data["image"]
|
||
position = frame_data["position"]
|
||
strength = frame_data["strength"]
|
||
|
||
if strength <= 0.0:
|
||
continue
|
||
|
||
# 编码中间帧图像
|
||
middle_latent = cls._encode_image(vae, image, height, width)
|
||
middle_latent_frames = middle_latent.shape[2]
|
||
|
||
# 根据百分比计算中间帧位置
|
||
middle_frame_idx = round(position * (latent_frames - 1))
|
||
|
||
# 确保不超出范围
|
||
middle_frame_idx = max(0, min(middle_frame_idx, latent_frames - middle_latent_frames))
|
||
|
||
# 嵌入中间帧
|
||
samples[:, :, middle_frame_idx:middle_frame_idx+middle_latent_frames] = middle_latent
|
||
|
||
# 设置中间帧的noise_mask（使用min确保不覆盖更强的约束）
|
||
current_mask = noise_mask[:, :, middle_frame_idx:middle_frame_idx+middle_latent_frames, :, :]
|
||
new_mask = 1.0 - strength
|
||
# 保留更强的约束（更小的mask值）
|
||
noise_mask[:, :, middle_frame_idx:middle_frame_idx+middle_latent_frames, :, :] = torch.minimum(
|
||
current_mask,
|
||
torch.full_like(current_mask, new_mask)
|
||
)
|
||
|
||
return ({"samples": samples, "noise_mask": noise_mask},)
|
||
|
||
@staticmethod
|
||
def _encode_image(vae, image, target_height, target_width):
|
||
"""
|
||
编码图像到latent空间
|
||
|
||
Args:
|
||
vae: VAE模型
|
||
image: 输入图像张量
|
||
target_height: 目标高度
|
||
target_width: 目标宽度
|
||
|
||
Returns:
|
||
编码后的latent张量
|
||
"""
|
||
# 调整图像尺寸
|
||
if image.shape[1] != target_height or image.shape[2] != target_width:
|
||
pixels = comfy.utils.common_upscale(
|
||
image.movedim(-1, 1),
|
||
target_width,
|
||
target_height,
|
||
"bilinear",
|
||
"center"
|
||
).movedim(1, -1)
|
||
else:
|
||
pixels = image
|
||
|
||
# 只取RGB通道(去掉alpha)
|
||
encode_pixels = pixels[:, :, :, :3]
|
||
|
||
# VAE编码
|
||
latent = vae.encode(encode_pixels)
|
||
|
||
return latent
|
||
|
||
# 兼容旧API
|
||
generate = execute
|