""" 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", }