""" LTXVSceneExtender: All-in-one node for extending video scenes with synchronized audio. Combines LTXVLoopingSampler + LTXVNormalizingSampler + image guides into a single node that processes timestamped scene scripts. """ import copy from typing import Optional import torch from comfy_api.latest import io # Import ComfyUI components try: import nodes import comfy.model_management as mm import comfy.utils from comfy.nested_tensor import NestedTensor COMFY_AVAILABLE = True except ImportError: COMFY_AVAILABLE = False NestedTensor = None # Import LTXV nodes/components try: from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced # Try importing LTXV-specific utilities from the custom node package # This tries standard installation paths try: from custom_nodes.ComfyUI_LTXVideo.easy_samplers import LinearOverlapLatentTransition from custom_nodes.ComfyUI_LTXVideo.latents import LTXVAddLatentGuide, LTXVSelectLatents LTXV_UTILS_AVAILABLE = True except ImportError: # Check if we are inside the package (relative import) or handle import differently LTXV_UTILS_AVAILABLE = False except ImportError: LTXV_UTILS_AVAILABLE = False # Local imports from .time_manager import TimeManager from .script_parser import ( parse_scene_script, resolve_image_refs, get_chunk_guide_images, SceneChunk, ) from .audio_blender import AudioOverlapBlender # Global Cache to avoid Locked Class issues _CACHE = {} def get_noise_mask(latent): """Helper to extract noise mask, handling nested tensors.""" if "noise_mask" in latent: nm = latent["noise_mask"] if isinstance(nm, NestedTensor): nm = nm.tensors[0] # Get Video Mask return nm return None class LTXVSceneExtender(io.ComfyNode): """ All-in-one node for extending video scenes with synchronized audio. Combines temporal tiling, audio normalization, and image guides into a single node that processes timestamped scene scripts. Features: - Timestamped prompts with audio specs (silent, ambient, dialogue) - Image guides at first/end/specific frames - Smooth audio blending at chunk transitions - All time inputs in SECONDS (internal frame math abstracted) """ @classmethod def define_schema(cls): return io.Schema( node_id="LTXVSceneExtender", display_name="💜 LTXV Scene Extender ErosDiffusion", category="ErosDiffusion/ltxv", description="Extend video with synchronized audio, image guides, and timestamped prompts", inputs=[ # === Model Inputs === io.Model.Input( "model", tooltip="Diffusion model (LTXAVModel for audio-video, or video-only model)" ), io.Vae.Input("video_vae", tooltip="Video VAE for encoding/decoding"), io.Vae.Input( "audio_vae", optional=True, tooltip="Audio VAE (required for audio generation)" ), io.Sampler.Input("sampler", tooltip="Sampler to use"), io.Sigmas.Input("sigmas", tooltip="Sigma schedule"), io.Noise.Input("noise", tooltip="Noise source"), io.Guider.Input( "guider", tooltip="Guider (e.g., STGGuiderAdvanced)" ), io.Clip.Input("clip", tooltip="CLIP model for encoding prompts"), # === Existing Content === io.Latent.Input( "latent", tooltip="Existing video/AV latent to extend (optional for new generation)", optional=True ), # === Extension Configuration === io.Float.Input( "overlap_duration", default=1.0, min=0.5, max=3.0, step=0.1, tooltip="Overlap between chunks for smooth transitions" ), io.Float.Input( "video_fps", default=25.0, min=1.0, max=60.0, step=1.0, tooltip="Video frame rate" ), io.Int.Input( "width", default=960, min=64, max=2048, step=32, tooltip="Output video width" ), io.Int.Input( "height", default=512, min=64, max=2048, step=32, tooltip="Output video height" ), # === Scene Script === io.String.Input( "scene_script", default="""[00:00-00:01] A woman with short hair speaks to camera | audio:"Hello" | first:$0 | end:$0 [00:01-00:03] A woman with short hair turns around and speaks to camera , camera orbits around | audio:"I like this camera" | first:$0 | end:$1 [00:03-00:06] A woman with short hair opens her arms and tilts head backwards happy | audio: "I feel free" | first:$1 | end:$2 [00:06-00:08] A woman with short hair puts her hands in her hair | first:$2""", multiline=True, tooltip="""Timestamped scene script. Format: [MM:SS-MM:SS] Scene description | audio:spec | first:$0 | MM:SS:$1 | end:$2 Audio specs: audio:silent, audio:ambient, audio:"dialogue text" Guide refs: $0, $1, etc. reference guide_images batch by index""" ), # === Image Guides === io.Image.Input( "guide_images", optional=True, tooltip="Batch of guide images (referenced as $0, $1, etc.)" ), io.Float.Input( "guide_strength", default=1.0, min=0.0, max=1.0, step=0.05, tooltip="Strength of image guides" ), # === Audio Controls === io.Float.Input( "audio_overlap_duration", default=0.5, min=0.1, max=2.0, step=0.1, tooltip="Audio overlap for smooth blending at transitions" ), io.Int.Input( "audio_slope_frames", default=5, min=1, max=20, step=1, tooltip="Crossfade slope length for seamless audio" ), io.String.Input( "audio_normalization", default="1,1,0.25,1,1,0.25,1,1", tooltip="Per-step audio normalization factors" ), # === Advanced === io.Int.Input("video_overlap_frames", default=8, min=0, max=64), io.Int.Input("audio_overlap_frames", default=32, min=0, max=256), io.Float.Input( "temporal_cond_strength", default=1.0, min=0.0, max=1.0, step=0.05, tooltip="Conditioning strength from previous tile overlap" ), io.Float.Input( "adain_factor", default=0.1, min=0.0, max=1.0, step=0.05, tooltip="AdaIN factor to prevent oversaturation" ), io.Boolean.Input( "enable_cache", default=False, tooltip="Enable chunk caching for faster re-runs (can use extra memory)" ), ], outputs=[ io.Latent.Output(display_name="latent"), io.Latent.Output(display_name="video_latent"), io.Latent.Output(display_name="audio_latent"), io.Conditioning.Output(display_name="positive"), io.Conditioning.Output(display_name="negative"), ], ) @classmethod def execute( cls, model, video_vae, sampler, sigmas, noise, guider, clip, overlap_duration: float, video_fps: float, width: int, height: int, scene_script: str, guide_strength: float, audio_overlap_duration: float, audio_slope_frames: int, audio_normalization: str, video_overlap_frames: int, audio_overlap_frames: int, temporal_cond_strength: float, adain_factor: float, enable_cache: bool, audio_vae=None, latent=None, guide_images=None, ) -> io.NodeOutput: batch_size = 1 # Hardcoded """Execute the scene extension.""" # Initialize Cache global _CACHE # Clear cache if script fundamentally changed structure? # But maybe limit cache size? if len(_CACHE) > 50: # Arbitrary limit _CACHE.clear() # Check if we have an audio-video model is_av_model = cls._is_av_model(model) # Initialize TimeManager time_mgr = TimeManager( video_fps=video_fps, audio_sample_rate=16000 if audio_vae is None else getattr( audio_vae.autoencoder, 'sampling_rate', 16000 ), mel_hop_length=160 if audio_vae is None else getattr( audio_vae.autoencoder, 'mel_hop_length', 160 ), ) # Parse scene script if scene_script.strip(): chunks = parse_scene_script(scene_script) else: chunks = cls._generate_default_chunks( extension_duration, tile_duration, overlap_duration ) print(f"LTXVSceneExtender: Processing {len(chunks)} chunks.") # --- OPTIMIZATION PHASE: PRE-ENCODING --- print("Optimization: Pre-encoding prompts and guides...") chunk_data, resolved_refs = cls._pre_encode_assets( chunks, clip, video_vae, guide_images, width, height, time_mgr ) print("Optimization: Encoding complete. Starting Generation Loop.") # Trigger soft cache cleanup to free VAE/CLIP memory if possible mm.soft_empty_cache() # --- GENERATION LOOP --- full_video = None full_audio = None # Initialize audio blender if we have audio audio_blender = None if is_av_model and audio_vae is not None: audio_blender = AudioOverlapBlender( overlap_frames=int( audio_overlap_duration * time_mgr.config.audio_latents_per_second ), slope_len=audio_slope_frames, ) # Get dimensions time_scale_factor, width_scale_factor, height_scale_factor = ( video_vae.downscale_index_formula ) latent_height = height // height_scale_factor latent_width = width // width_scale_factor scale_factors = video_vae.downscale_index_formula # Keep track of previous latent for extension prev_latent = latent final_video_list = [] positive = None negative = None # Store last used for return for i, (chunk, c_data) in enumerate(zip(chunks, chunk_data)): chunk_duration = chunk.end_sec - chunk.start_sec latent_length = max(1, time_mgr.calculate_video_latent_count(chunk_duration)) # Ensure at least 1 print(f"Processing chunk {i+1}/{len(chunks)}: [{chunk.start_sec:.1f}s - {chunk.end_sec:.1f}s] ({latent_length} latents / ~{latent_length*8} frames)") # === CACHING CHECK === prev_context_hash = 0.0 if prev_latent is not None: if isinstance(prev_latent["samples"], NestedTensor): v_tens = prev_latent["samples"].tensors[0] else: v_tens = prev_latent["samples"] prev_context_hash = float(v_tens.mean().cpu()) # Simple hash cache_key = ( i, chunk.prompt, chunk.duration, chunk.transition_type, prev_context_hash, width, height ) if enable_cache and cache_key in _CACHE: print(f"Using Cached Chunk {i}") chunk_res = _CACHE[cache_key] final_video_list.append(chunk_res) prev_latent = chunk_res continue # --- Allocation & Extension Logic --- is_extension = (prev_latent is not None) and (chunk.transition_type != "cut") if not is_extension: # NEW GENERATION # Basic Mask video_mask = torch.ones( (batch_size, 1, latent_length, 1, 1), device=mm.intermediate_device() ) video_latent = torch.zeros( [batch_size, 128, latent_length, latent_height, latent_width], device=mm.intermediate_device() if COMFY_AVAILABLE else "cpu" ) # Handle AV Model if is_av_model and NestedTensor is not None and audio_vae is not None: audio_len = time_mgr.calculate_audio_latent_count(chunk_duration) if audio_len < 1: audio_len = 1 print(f" [DEBUG] NEW GEN: audio_len forced to 1") # Get correct dimensions from VAE (defaults match typical LTXAV values) a_ch = getattr(audio_vae, "latent_channels", 8) a_freq = getattr(audio_vae, "latent_frequency_bins", 16) audio_latent = torch.zeros( [batch_size, a_ch, audio_len, a_freq], device=mm.intermediate_device() if COMFY_AVAILABLE else "cpu" ) nt = NestedTensor((video_latent, audio_latent)) input_latent = {"samples": nt} # Mask handling for AV # Video part: video_mask # Audio part: Ones audio_mask = torch.ones((batch_size, 1, audio_len, 1), device=mm.intermediate_device()) nt_mask = NestedTensor((video_mask, audio_mask)) input_latent["noise_mask"] = nt_mask else: input_latent = {"samples": video_latent, "noise_mask": video_mask} # If audio_vae missing, we skip audio latent creation else: # EXTENSION # Extract Previous Tail prev_samples = prev_latent["samples"] if is_av_model and NestedTensor is not None and isinstance(prev_samples, NestedTensor): prev_video = prev_samples.tensors[0] # [B, 128, T, H, W] prev_audio_src = prev_samples.tensors[1] else: prev_video = prev_samples prev_audio_src = None # Copy Overlap (Robust) video_overlap = video_overlap_frames if video_overlap > prev_video.shape[2]: video_overlap = prev_video.shape[2] if video_overlap > latent_length: # Clamp to new chunk size video_overlap = latent_length src_v = prev_video[..., -video_overlap:, :, :] # Allocate New current_v = torch.zeros( [batch_size, 128, latent_length, latent_height, latent_width], device=mm.intermediate_device() if COMFY_AVAILABLE else "cpu" ) current_v[:, :, :video_overlap, :, :] = src_v # AV Logic if is_av_model and NestedTensor is not None and audio_vae is not None: # Calculate audio length with defensive minimum audio_len = time_mgr.calculate_audio_latent_count(chunk_duration) if audio_len < 1: audio_len = 1 print(f" [DEBUG] audio_len forced to 1 (was {time_mgr.calculate_audio_latent_count(chunk_duration)})") a_ch = getattr(audio_vae, "latent_channels", 8) # Default 8, not 128 a_freq = getattr(audio_vae, "latent_frequency_bins", 16) # Default 16 current_a = torch.zeros( [batch_size, a_ch, audio_len, a_freq], device=mm.intermediate_device() if COMFY_AVAILABLE else "cpu" ) audio_overlap = audio_overlap_frames if prev_audio_src is not None: # Ensure prev_audio_src has batch dimension [B, C, T, F] if prev_audio_src.ndim == 3: prev_audio_src = prev_audio_src.unsqueeze(0) # Safe clamp of copy length a_copy_len = min(audio_overlap, prev_audio_src.shape[2], current_a.shape[2]) if a_copy_len > 0: src_a = prev_audio_src[:, :, -a_copy_len:, :] current_a[:, :, :a_copy_len, :] = src_a else: print(f" [DEBUG] Skipping audio overlap copy (a_copy_len={a_copy_len})") nt = NestedTensor((current_v, current_a)) input_latent = {"samples": nt} # Create Masks (Mask out overlap regions) v_mask = torch.ones((batch_size, 1, latent_length, 1, 1), device=mm.intermediate_device()) v_mask[:, :, :video_overlap] = 1.0 - temporal_cond_strength a_mask = torch.ones((batch_size, 1, audio_len, 1), device=mm.intermediate_device()) a_mask[:, :, :audio_overlap] = 1.0 - temporal_cond_strength input_latent["noise_mask"] = NestedTensor((v_mask, a_mask)) else: input_latent = {"samples": current_v} v_mask = torch.ones((batch_size, 1, latent_length, 1, 1), device=mm.intermediate_device()) v_mask[:, :, :video_overlap] = 1.0 - temporal_cond_strength input_latent["noise_mask"] = v_mask # --- SETUP CONDITIONING --- chunk_guider = copy.copy(guider) # Using pre-encoded prompts c_pos = c_data["cond_pos"] c_pos = c_data["cond_pos"] # Use Global Negative from Guider (since we didn't pre-encode negative) _, global_neg = cls._get_conds_from_guider(guider) c_neg = global_neg if global_neg is not None else [] # Sanitize Negative (KJ Fix) if c_neg and len(c_neg) > 0 and isinstance(c_neg[0], dict): c_neg = [] # --- APPLY PRE-ENCODED GUIDES --- guides = c_data["guides"] if guides: print(f" Applying {len(guides)} image guides (Pre-encoded)") # Extract working Tensors if isinstance(input_latent["samples"], NestedTensor): working_v = input_latent["samples"].tensors[0] # We handle audio parts later else: working_v = input_latent["samples"] if isinstance(input_latent["noise_mask"], NestedTensor): working_mask = input_latent["noise_mask"].tensors[0] else: working_mask = input_latent["noise_mask"] for g in guides: g_latent = g["latent"] frame_offset = g["frame_idx"] # From script ref g_strength = g["strength"] # Adjust offset for Extension chunks? # Chunk logic: if extension, new latent contains OVERLAP. # Script 'first' ($0) usually maps to 0. # If latent starts with overlap, index 0 IS the overlap start. # So offset logic should be consistent with how user expects it. # If user says "first", they mean Frame 0 of this chunk. # Which is Frame 0 of latent. # So no adjustment needed? # BUT 'resolve_image_refs' logic maps 'first'->0. # IF is_extension, we might want to offset by overlap? # No, let's assume raw frame index into the latent currently being generated. # 1. Update Conds (Keyframe Index) # LTXVAddGuide method c_pos = LTXVAddGuide.add_keyframe_index(c_pos, frame_offset, g_latent, scale_factors) c_neg = LTXVAddGuide.add_keyframe_index(c_neg, frame_offset, g_latent, scale_factors) # 2. Update Latent/Mask time_scale = scale_factors[0] # Use Floor division for safer mapping of Frame -> Latent l_idx = frame_offset // time_scale # Clamp to bounds to prevent crash if 'end' lands on boundary cond_len = g_latent.shape[2] max_idx = working_v.shape[2] - cond_len l_idx = max(0, min(l_idx, max_idx)) working_v, working_mask = LTXVAddGuide.replace_latent_frames( working_v, working_mask, g_latent, l_idx, g_strength ) # Pack updated Tensors back if isinstance(input_latent["samples"], NestedTensor): orig_samples = input_latent["samples"] new_samples = NestedTensor((working_v, orig_samples.tensors[1])) orig_mask = input_latent["noise_mask"] new_mask = NestedTensor((working_mask, orig_mask.tensors[1])) input_latent["samples"] = new_samples input_latent["noise_mask"] = new_mask else: input_latent["samples"] = working_v input_latent["noise_mask"] = working_mask # Set Conditioning on Guider # Manual set to bypass validation if custom types try: chunk_guider.conds = {"positive": c_pos, "negative": c_neg} except: try: chunk_guider.set_conds(c_pos, c_neg) except: pass # Hope it worked # --- SAMPLING --- print(" Sampling...") _, denoised = SamplerCustomAdvanced().sample( noise, chunk_guider, sampler, sigmas, input_latent ) final_video_list.append(denoised) prev_latent = denoised # Cache it (if caching is enabled) if enable_cache: _CACHE[cache_key] = denoised # Keep refs for return positive = c_pos negative = c_neg # --- HANDLE EMPTY CHUNKS --- if len(final_video_list) == 0: print("WARNING: No chunks were processed. Returning empty latents.") # Create minimal placeholder latents empty_video = torch.zeros([1, 128, 1, latent_height, latent_width], device=mm.intermediate_device()) empty_audio = torch.zeros([1, 8, 1, 16], device=mm.intermediate_device()) if audio_vae is not None else None video_out = {"samples": empty_video} audio_out = {"samples": empty_audio} if empty_audio is not None else {"samples": torch.zeros([1,8,1,16])} if is_av_model and NestedTensor is not None and empty_audio is not None: combined = {"samples": NestedTensor((empty_video, empty_audio))} else: combined = video_out return io.NodeOutput(combined, video_out, audio_out, [], []) # --- BLENDING --- print("Generation complete. Blending chunks...") full_video = None full_audio = None video_overlap_frames = time_mgr.calculate_video_latent_count(overlap_duration) audio_overlap_frames = time_mgr.calculate_audio_latent_count(overlap_duration) for i, chunk_res in enumerate(final_video_list): chunk = chunks[i] is_cut = chunk.transition_type == "cut" if is_av_model and NestedTensor is not None and isinstance(chunk_res["samples"], NestedTensor): v_part = chunk_res["samples"].tensors[0] a_part = chunk_res["samples"].tensors[1] else: v_part = chunk_res["samples"] a_part = None if i == 0: full_video = v_part full_audio = a_part else: if is_cut: print(f" Chunk {i+1}: Hard Cut") overlap_use = 0 else: overlap_use = video_overlap_frames full_video = cls._blend_latents(full_video, v_part, overlap_use) if full_audio is not None and a_part is not None: if is_cut: # Hard cut - just concatenate without overlap full_audio = torch.cat([full_audio, a_part], dim=2) elif audio_blender is not None: # Use specialized audio blender for proper crossfade full_audio = audio_blender.blend_chunks(full_audio, a_part) else: # Fallback to generic blending a_ov_use = audio_overlap_frames full_audio = cls._blend_latents(full_audio, a_part, a_ov_use) # Output video_out = {"samples": full_video} # Use correct audio latent shape: [B, 8, T, 16] where 8 is channels, 16 is freq bins audio_out = {"samples": full_audio} if full_audio is not None else {"samples": torch.zeros([1,8,1,16], device=mm.intermediate_device())} if is_av_model and NestedTensor is not None and full_audio is not None: combined = {"samples": NestedTensor((full_video, full_audio))} else: combined = video_out return io.NodeOutput( combined, video_out, audio_out, # Fixed Typo positive, negative ) @classmethod def _pre_encode_assets(cls, chunks, clip, video_vae, guide_images, width, height, time_mgr): """Helper to pre-encode prompts and guide images.""" chunk_data = [] # Resolve all references first (to get paths) resolved_refs = resolve_image_refs(chunks, guide_images) latent_width = width // 32 latent_height = height // 32 scale_factors = video_vae.downscale_index_formula for i, chunk in enumerate(chunks): # Text cond_pos = cls._encode_prompt(clip, chunk.prompt) # Use empty string for neg as simple default or passed param? # We implemented global negative in execute, but let's use empty here # and let execute merge global? # User wants optimization. We can encode global NEG once outside loop? # Wait, signature of execute has `negative_prompt`? No, it has `guider`. # We extract from Guider. # So pre-encoding NEG is done via Guider extraction in Loop. # We only pre-encode POS specific to chunk. # Guides processed_guides = [] chunk_guides, chunk_indices = get_chunk_guide_images(chunk, resolved_refs, time_mgr) if chunk_guides is not None: # chunk_guides is [N, H, W, C] # Indices string indices_list = [int(x) for x in chunk_indices.split(",")] # Resize resized_guides = comfy.utils.common_upscale( chunk_guides.movedim(-1, 1), width, height, "bilinear", "center" ).movedim(1, -1) # Encode for img, idx in zip(resized_guides, indices_list): try: # Use LTXVAddGuide.encode logic # (It expects batch of images but we do one by one for simplicity/safety) # Actually LTXVAddGuide.encode handles batch. # But we need granular control per index. _, g_latent = LTXVAddGuide.encode( video_vae, latent_width, latent_height, img.unsqueeze(0), scale_factors ) processed_guides.append({ "latent": g_latent, "frame_idx": idx, "strength": 1.0 # TODO: Get from chunk ref? Parser supports it, but here we used batch strength? # Input guide_strength is global. # Script parser has strength in $0:1.0? # resolve_image_refs logic... }) except Exception as e: print(f"Warning: Guide encode failed: {e}") # Basic Negative placeholder (will be replaced by Global Neg in Loop) chunk_data.append({ "cond_pos": cond_pos, "cond_neg": None, # Will use Global "guides": processed_guides }) return chunk_data, resolved_refs @classmethod def _blend_latents(cls, prev: torch.Tensor, next_t: torch.Tensor, overlap: int) -> torch.Tensor: """Blend two latents with linear crossfade on dim 2 (time).""" if overlap <= 0: return torch.cat([prev, next_t], dim=2) overlap = min(overlap, prev.shape[2], next_t.shape[2]) prev_cut = prev[:, :, :-overlap] prev_tail = prev[:, :, -overlap:] next_head = next_t[:, :, :overlap] next_cut = next_t[:, :, overlap:] alpha = torch.linspace(0, 1, overlap, device=prev.device, dtype=prev.dtype) # Dynamic Reshape for Broadcasting (Video 5D or Audio 4D) shape = [1, 1, -1] + [1] * (prev.ndim - 3) alpha = alpha.view(*shape) blended = prev_tail * (1.0 - alpha) + next_head * alpha return torch.cat([prev_cut, blended, next_cut], dim=2) @classmethod def _is_av_model(cls, model) -> bool: try: return model.model.diffusion_model.__class__.__name__ == "LTXAVModel" except AttributeError: return False @classmethod def _get_conds_from_guider(cls, guider): conds = None if hasattr(guider, "conds"): conds = guider.conds elif hasattr(guider, "raw_conds"): conds = guider.raw_conds elif hasattr(guider, "original_conds"): conds = guider.original_conds if conds is not None: if isinstance(conds, dict): return conds.get("positive"), conds.get("negative") return conds return None, None @classmethod def _encode_prompt(cls, clip, prompt: str): tokens = clip.tokenize(prompt) try: cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) return [[cond, {"pooled_output": pooled}]] except: conds = clip.encode_from_tokens_scheduled(tokens) if conds and isinstance(conds, list) and len(conds) > 0 and isinstance(conds[0], dict): new_conds = [] for c in conds: tensor = None if "cross_attn" in c: tensor = c["cross_attn"] elif "pooled_output" in c: tensor = c["pooled_output"] if tensor is not None: new_conds.append([tensor, c]) if new_conds: return new_conds return conds @classmethod def _generate_default_chunks( cls, duration: float, tile_duration: float, overlap_duration: float ) -> list[SceneChunk]: chunks = [] effective_tile = tile_duration - overlap_duration current_start = 0.0 while current_start < duration: chunk_end = min(current_start + tile_duration, duration) chunks.append(SceneChunk( start_sec=current_start, end_sec=chunk_end, prompt="", audio_spec="silent", guides=[], )) current_start += effective_tile if chunk_end >= duration: break return chunks