BiyikgokhanandClaude Opus 4.6
bf31dde091
Single ComfyUI node that auto-scales Wan2.1 InfiniTalk video generation to any audio duration. No manual segment chaining needed. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
300 lines
12 KiB
Python
300 lines
12 KiB
Python
"""
|
|
InfiniTalk AutoScale — single ComfyUI node that auto-scales to any audio duration.
|
|
|
|
Internally loops N segments based on audio length, handling:
|
|
- Audio embedding windowing per segment
|
|
- Motion frame overlap between segments
|
|
- Latent sampling + VAE decode per segment
|
|
- Frame accumulation with overlap trimming
|
|
|
|
Canvas usage:
|
|
[LoadImage] ─────────┐
|
|
[LoadAudio] ─► [AudioEncoderEncode] ─┐
|
|
[UNETLoader] ─► [LoRA] ─┐ │
|
|
[ModelPatchLoader] ──────┤ │
|
|
[CLIPTextEncode] ────────┤ │
|
|
[ConditioningZeroOut] ───┤ │
|
|
[VAELoader] ─────────────┤ │
|
|
▼ ▼
|
|
[InfiniTalkAutoScale]
|
|
│
|
|
▼ IMAGE
|
|
[CreateVideo] ─► [SaveVideo]
|
|
"""
|
|
|
|
import logging
|
|
|
|
import torch
|
|
import comfy.model_management
|
|
import comfy.patcher_extension
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import latent_preview
|
|
import node_helpers
|
|
|
|
from comfy.ldm.wan.model_multitalk import (
|
|
InfiniteTalkOuterSampleWrapper,
|
|
MultiTalkCrossAttnPatch,
|
|
MultiTalkGetAttnMapPatch,
|
|
project_audio_features,
|
|
)
|
|
from comfy_extras.nodes_wan import linear_interpolation
|
|
|
|
logger = logging.getLogger("InfiniTalkAutoScale")
|
|
|
|
|
|
class InfiniTalkAutoScale:
|
|
"""Single node that generates a full InfiniTalk video, auto-scaled to audio duration."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"model_patch": ("MODEL_PATCH",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"vae": ("VAE",),
|
|
"audio_encoder_output_1": ("AUDIO_ENCODER_OUTPUT",),
|
|
"start_image": ("IMAGE",),
|
|
"width": ("INT", {"default": 832, "min": 16, "max": 16384, "step": 16}),
|
|
"height": ("INT", {"default": 480, "min": 16, "max": 16384, "step": 16}),
|
|
"frame_window": ("INT", {
|
|
"default": 81, "min": 17, "max": 257, "step": 4,
|
|
"tooltip": "Frames per segment (must be 4n+1). 81 = ~3.2s per segment.",
|
|
}),
|
|
"motion_frame_count": ("INT", {
|
|
"default": 9, "min": 1, "max": 33, "step": 1,
|
|
"tooltip": "Overlap frames between segments for temporal continuity.",
|
|
}),
|
|
"max_segments": ("INT", {
|
|
"default": 0, "min": 0, "max": 100, "step": 1,
|
|
"tooltip": "Max segments (0 = auto from audio length, no limit).",
|
|
}),
|
|
"steps": ("INT", {"default": 4, "min": 1, "max": 100}),
|
|
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"sampler_name": (comfy.samplers.SAMPLER_NAMES,),
|
|
"scheduler": (comfy.samplers.SCHEDULER_NAMES,),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"audio_scale": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("video_frames",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "video/infinitalk"
|
|
OUTPUT_NODE = False
|
|
DESCRIPTION = (
|
|
"Auto-scaling InfiniTalk video generation. "
|
|
"Determines segment count from audio duration and loops internally."
|
|
)
|
|
|
|
# ── helpers ────────────────────────────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _process_audio_embeddings(audio_encoder_output):
|
|
"""Convert raw audio encoder output to frame-level embeddings. Returns (encoded_audio, total_frames)."""
|
|
all_layers = audio_encoder_output["encoded_audio_all_layers"]
|
|
encoded_audio = torch.stack(all_layers, dim=0).squeeze(1)[1:] # [num_layers, T, 512]
|
|
encoded_audio = linear_interpolation(encoded_audio, input_fps=50, output_fps=25) # [num_layers, T', 512]
|
|
encoded_audio = encoded_audio.movedim(0, 1) # [T', num_layers, 512]
|
|
total_frames = encoded_audio.shape[0]
|
|
return encoded_audio, total_frames
|
|
|
|
@staticmethod
|
|
def _calc_segments(total_audio_frames, frame_window, motion_frame_count, max_segments):
|
|
"""Calculate how many segments are needed to cover the audio."""
|
|
new_frames_per_seg = frame_window - motion_frame_count
|
|
if total_audio_frames <= frame_window:
|
|
n = 1
|
|
else:
|
|
n = 1 + max(1, -(-((total_audio_frames - frame_window)) // new_frames_per_seg)) # ceil div
|
|
if max_segments > 0:
|
|
n = min(n, max_segments)
|
|
return n
|
|
|
|
def _prepare_segment(self, model, model_patch, positive, negative, vae,
|
|
encoded_audio_list, start_image, previous_frames,
|
|
width, height, length, motion_frame_count,
|
|
audio_start, audio_end, audio_scale):
|
|
"""Replicates WanInfiniteTalkToVideo.execute() logic for one segment."""
|
|
|
|
if previous_frames is not None and previous_frames.shape[0] < motion_frame_count:
|
|
raise ValueError("Not enough previous frames for motion context.")
|
|
|
|
latent = torch.zeros(
|
|
[1, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
|
|
device=comfy.model_management.intermediate_device(),
|
|
)
|
|
|
|
if start_image is not None:
|
|
si = comfy.utils.common_upscale(
|
|
start_image[:length].movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
image = torch.ones(
|
|
(length, height, width, si.shape[-1]), device=si.device, dtype=si.dtype
|
|
) * 0.5
|
|
image[:si.shape[0]] = si
|
|
concat_latent_image = vae.encode(image[:, :, :, :3])
|
|
concat_mask = torch.ones(
|
|
(1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]),
|
|
device=si.device, dtype=si.dtype,
|
|
)
|
|
concat_mask[:, :, :((si.shape[0] - 1) // 4) + 1] = 0.0
|
|
positive = node_helpers.conditioning_set_values(
|
|
positive, {"concat_latent_image": concat_latent_image, "concat_mask": concat_mask}
|
|
)
|
|
negative = node_helpers.conditioning_set_values(
|
|
negative, {"concat_latent_image": concat_latent_image, "concat_mask": concat_mask}
|
|
)
|
|
|
|
model_patched = model.clone()
|
|
|
|
if previous_frames is not None:
|
|
motion_frames = comfy.utils.common_upscale(
|
|
previous_frames[-motion_frame_count:].movedim(-1, 1), width, height, "bilinear", "center"
|
|
).movedim(1, -1)
|
|
motion_frames_latent = vae.encode(motion_frames[:, :, :, :3])
|
|
trim_image = motion_frame_count
|
|
else:
|
|
motion_frames_latent = concat_latent_image[:, :, :1]
|
|
trim_image = 0
|
|
|
|
audio_embed = project_audio_features(
|
|
model_patch.model.audio_proj, encoded_audio_list, audio_start, audio_end
|
|
).to(model_patched.model_dtype())
|
|
model_patched.model_options["transformer_options"]["audio_embeds"] = audio_embed
|
|
|
|
model_patched.add_wrapper_with_key(
|
|
comfy.patcher_extension.WrappersMP.OUTER_SAMPLE,
|
|
"infinite_talk_outer_sample",
|
|
InfiniteTalkOuterSampleWrapper(
|
|
motion_frames_latent,
|
|
model_patch,
|
|
is_extend=previous_frames is not None,
|
|
),
|
|
)
|
|
|
|
model_patched.set_model_patch(MultiTalkCrossAttnPatch(model_patch, audio_scale), "attn2_patch")
|
|
|
|
out_latent = {"samples": latent}
|
|
return model_patched, positive, negative, out_latent, trim_image
|
|
|
|
def _sample_segment(self, model_patched, positive, negative, latent_dict,
|
|
sampler_name, scheduler, steps, cfg, seed):
|
|
"""Run CFGGuider + SamplerCustomAdvanced logic for one segment."""
|
|
guider = comfy.samplers.CFGGuider(model_patched)
|
|
guider.set_conds(positive, negative)
|
|
guider.set_cfg(cfg)
|
|
|
|
sampler = comfy.samplers.sampler_object(sampler_name)
|
|
|
|
sigmas = comfy.samplers.calculate_sigmas(
|
|
model_patched.get_model_object("model_sampling"), scheduler, steps
|
|
).cpu()
|
|
|
|
latent = latent_dict.copy()
|
|
latent_image = latent["samples"]
|
|
latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image, None)
|
|
latent["samples"] = latent_image
|
|
|
|
noise = comfy.sample.prepare_noise(latent_image, seed)
|
|
|
|
callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1)
|
|
|
|
samples = guider.sample(
|
|
noise, latent_image, sampler, sigmas,
|
|
denoise_mask=None, callback=callback,
|
|
disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED,
|
|
seed=seed,
|
|
)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
return {"samples": samples}
|
|
|
|
# ── main entry ─────────────────────────────────────────────────────────
|
|
|
|
def generate(self, model, model_patch, positive, negative, vae,
|
|
audio_encoder_output_1, start_image, width, height,
|
|
frame_window, motion_frame_count, max_segments,
|
|
steps, cfg, sampler_name, scheduler, seed, audio_scale):
|
|
|
|
encoded_audio, total_audio_frames = self._process_audio_embeddings(audio_encoder_output_1)
|
|
encoded_audio_list = [encoded_audio]
|
|
num_segments = self._calc_segments(total_audio_frames, frame_window, motion_frame_count, max_segments)
|
|
|
|
total_video_frames = frame_window + (num_segments - 1) * (frame_window - motion_frame_count)
|
|
audio_duration = total_audio_frames / 25.0
|
|
video_duration = total_video_frames / 25.0
|
|
|
|
logger.info(
|
|
f"InfiniTalkAutoScale: audio={audio_duration:.1f}s ({total_audio_frames} frames) "
|
|
f"-> {num_segments} segments -> ~{video_duration:.1f}s video"
|
|
)
|
|
|
|
pbar = comfy.utils.ProgressBar(num_segments)
|
|
all_frames = []
|
|
previous_frames = None
|
|
audio_offset = 0
|
|
|
|
for seg_idx in range(num_segments):
|
|
seg_seed = seed + seg_idx
|
|
|
|
audio_start = audio_offset
|
|
audio_end = audio_start + frame_window
|
|
audio_end = min(audio_end, total_audio_frames + frame_window)
|
|
|
|
logger.info(f" Segment {seg_idx + 1}/{num_segments}: audio frames {audio_start}-{audio_end}, seed={seg_seed}")
|
|
|
|
model_patched, pos, neg, latent_dict, trim_image = self._prepare_segment(
|
|
model=model,
|
|
model_patch=model_patch,
|
|
positive=positive,
|
|
negative=negative,
|
|
vae=vae,
|
|
encoded_audio_list=encoded_audio_list,
|
|
start_image=start_image,
|
|
previous_frames=previous_frames,
|
|
width=width, height=height,
|
|
length=frame_window,
|
|
motion_frame_count=motion_frame_count,
|
|
audio_start=audio_start,
|
|
audio_end=audio_end,
|
|
audio_scale=audio_scale,
|
|
)
|
|
|
|
sampled = self._sample_segment(
|
|
model_patched, pos, neg, latent_dict,
|
|
sampler_name, scheduler, steps, cfg, seg_seed,
|
|
)
|
|
|
|
decoded = vae.decode(sampled["samples"])
|
|
if len(decoded.shape) == 5:
|
|
decoded = decoded.reshape(-1, decoded.shape[-3], decoded.shape[-2], decoded.shape[-1])
|
|
|
|
if seg_idx == 0:
|
|
all_frames.append(decoded)
|
|
else:
|
|
all_frames.append(decoded[trim_image:])
|
|
|
|
previous_frames = decoded
|
|
|
|
audio_offset += (frame_window - motion_frame_count)
|
|
|
|
pbar.update(1)
|
|
|
|
final_frames = torch.cat(all_frames, dim=0)
|
|
|
|
logger.info(f"InfiniTalkAutoScale: Done! {final_frames.shape[0]} total frames")
|
|
return (final_frames,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"InfiniTalkAutoScale": InfiniTalkAutoScale,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"InfiniTalkAutoScale": "InfiniTalk AutoScale",
|
|
}
|