From 7bb60125b9d8ee63af0a08ea88c4a713f9649edf Mon Sep 17 00:00:00 2001 From: Greg Date: Mon, 23 Feb 2026 10:21:15 -0800 Subject: [PATCH] =?UTF-8?q?feat:=20Initial=20release=20=E2=80=94=209=20nod?= =?UTF-8?q?es=20for=20chunked=20long=20video=20generation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Core nodes (7): BlendVideoChunks, BlendLatentChunks, ConcatVideoChunks, ExtractVideoChunk, GetFrameByIndex, GetFrameRange, VideoChunkPlanner Wan-specific nodes (2): WanChunkedI2VSampler (all-in-one with FLF/multi-keyframe), WanVideoChunkCalculator Rolling-reference approach solves context window reversal in I2V models. --- .gitignore | 19 ++ LICENSE | 21 ++ README.md | 212 ++++++++++++ __init__.py | 16 + nodes.py | 635 ++++++++++++++++++++++++++++++++++ pyproject.toml | 9 + wan_chunked_sampler.py | 753 +++++++++++++++++++++++++++++++++++++++++ 7 files changed, 1665 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 pyproject.toml create mode 100644 wan_chunked_sampler.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..845b64b --- /dev/null +++ b/.gitignore @@ -0,0 +1,19 @@ +__pycache__/ +*.py[cod] +*$py.class +*.so +*.egg-info/ +dist/ +build/ +.eggs/ +*.egg +.env +.venv +venv/ +.idea/ +.vscode/ +*.swp +*.swo +*~ +.DS_Store +Thumbs.db diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..b536033 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Greg Tee + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..0e675bc --- /dev/null +++ b/README.md @@ -0,0 +1,212 @@ +# ComfyUI VideoChunkTools + +**9 utility nodes for generating long videos in ComfyUI by splitting them into overlapping chunks with rolling reference frames.** + +Designed to solve the **"context window reversal" problem** — where video generation models (Wan 2.1/2.2, FantasyPortrait, etc.) revert to the starting state after ~135 frames because the reference image embedding pulls the generation back. + +![ComfyUI](https://img.shields.io/badge/ComfyUI-Custom_Nodes-blue) +![License](https://img.shields.io/badge/license-MIT-green) + +--- + +## The Problem + +Most image-to-video models have a limited context window (~81–137 frames). When you try to generate longer videos, the model "forgets" the accumulated motion and snaps back to looking like frame 1. The result is an obvious loop or identity shift. + +## The Solution + +**Rolling Reference**: Generate video in chunks, using the **last frame** of each chunk as the **reference image** for the next chunk. This keeps identity consistent while allowing pose/action to progress naturally. + +``` +Chunk 0: [ref=original portrait] → 81 frames → last frame becomes next ref +Chunk 1: [ref=chunk 0 last frame] → 81 frames → last frame becomes next ref +Chunk 2: [ref=chunk 1 last frame] → 81 frames → ... +Final: Concatenate all chunks → 240+ frame seamless video +``` + +--- + +## Nodes + +### Core Utility Nodes (7 nodes — no dependencies beyond ComfyUI) + +| Node | Purpose | +|------|---------| +| **Extract Video Chunk** | Pull chunk N from a longer driving video, with configurable overlap | +| **Blend Video Chunks (Crossfade)** | Crossfade-blend two overlapping pixel-space chunks into one seamless video | +| **Blend Latent Chunks (Pre-Decode)** | Join two latent chunks before VAE decode — supports slerp, hard_cut, crossfade | +| **Concat Video Chunks** | Simple concatenation with optional first-frame trim (for rolling reference) | +| **Get Frame By Index** | Extract a single frame by index (-1 = last frame = next reference) | +| **Get Frame Range** | Extract a range of frames with negative indexing support | +| **Video Chunk Planner** | Calculate chunking strategy — shows chunk count, frame ranges, workflow steps | + +### Wan-Specific Nodes (2 nodes — require [ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper)) + +| Node | Purpose | +|------|---------| +| **Wan Chunked I2V Sampler ♾️** | All-in-one node: encode→sample→decode→extract ref→repeat. Handles everything internally. | +| **Wan Chunk Calculator 🧮** | Calculate exact total_frames for N chunks with 4n+1 normalization | + +> The Wan nodes gracefully degrade — if WanVideoWrapper isn't installed, the 7 core nodes still load and work fine. + +--- + +## Installation + +### Option 1: ComfyUI Manager (Recommended) +Search for `VideoChunkTools` in ComfyUI Manager and click Install. + +### Option 2: Git Clone +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/gregtee2/ComfyUI_VideoChunkTools.git +``` + +### Option 3: Manual Download +Download the ZIP from GitHub, extract to `ComfyUI/custom_nodes/ComfyUI_VideoChunkTools/`. + +**Restart ComfyUI after installation.** + +No pip dependencies required — uses only PyTorch (already in ComfyUI). + +--- + +## Workflows + +### Rolling Reference (Manual — Any I2V Model) + +Use the core utility nodes to build a rolling-reference pipeline with any image-to-video model: + +1. **Video Chunk Planner** → See how many chunks you need +2. **Extract Video Chunk** (index=0) → Get first driving chunk +3. Run your I2V model with your **original reference image** +4. **Get Frame By Index** (index=-1) → Extract last frame as new reference +5. **Extract Video Chunk** (index=1) → Get next driving chunk +6. Run your I2V model with the **new reference** +7. **Blend Video Chunks** or **Concat Video Chunks** → Join the results +8. Repeat for each chunk + +### All-In-One (Wan Models) + +The **Wan Chunked I2V Sampler** handles everything in a single node: + +1. Connect your Wan model, VAE, and start image +2. Set `total_frames` (e.g., 241 for ~15 seconds at 16fps) +3. Set `chunk_frames` (e.g., 81) +4. Hit Queue — the node generates all chunks automatically + +**Features:** +- **Single-pass** or **Two-pass** sampling (connect `model_b` for split denoising) +- **FLF (First-Last-Frame)** — connect an `end_image` to guide the final frame +- **Multi-keyframe FLF** — provide a batch of end images to distribute across chunks +- **Crossfade overlap** — set `end_blend_chunks` for smooth FLF transitions +- **Auto 4n+1 normalization** — chunk sizes are automatically adjusted for Wan's requirements + +--- + +## Node Details + +### Extract Video Chunk + +Divides a video into chunks with overlap. Adjacent chunks share `overlap_frames` at their boundary. + +| Input | Type | Default | Description | +|-------|------|---------|-------------| +| images | IMAGE | — | Full driving video | +| chunk_index | INT | 0 | Which chunk to extract (0-based) | +| chunk_frames | INT | 81 | Frames per chunk | +| overlap_frames | INT | 16 | Frames shared between adjacent chunks | + +| Output | Type | Description | +|--------|------|-------------| +| chunk | IMAGE | Extracted chunk | +| total_chunks | INT | How many chunks cover the full video | +| chunk_index | INT | Pass-through for chaining | +| is_last_chunk | BOOLEAN | True if this is the final chunk | + +### Blend Video Chunks (Crossfade) + +Crossfade two overlapping pixel-space chunks. The last N frames of chunk_a smooth-transition into the first N frames of chunk_b. + +| Input | Type | Default | Description | +|-------|------|---------|-------------| +| chunk_a | IMAGE | — | First (earlier) chunk | +| chunk_b | IMAGE | — | Second (later) chunk | +| overlap_frames | INT | 16 | Frames to crossfade | +| blend_curve | ENUM | ease_in_out | linear, ease_in_out, sigmoid | + +### Blend Latent Chunks (Pre-Decode) + +Join two 5D video latents along the temporal dimension **before** VAE decode. Operates in latent space — overlap is in latent frames (Wan: pixel_overlap / 4). + +| Input | Type | Default | Description | +|-------|------|---------|-------------| +| latent_a | LATENT | — | First latent chunk | +| latent_b | LATENT | — | Second latent chunk | +| overlap_frames | INT | 4 | Overlap in latent temporal frames | +| blend_curve | ENUM | hard_cut | hard_cut, slerp, linear, ease_in_out, sigmoid | + +**hard_cut** (recommended for rolling reference): Clean cut at the overlap midpoint — no dissolve artifacts. + +**slerp**: Spherical linear interpolation — preserves latent vector magnitude. Standard technique for diffusion model interpolation. + +### Concat Video Chunks + +Simple concatenation. For rolling-reference workflows where chunk B's first frame naturally matches chunk A's last frame. + +| Input | Type | Default | Description | +|-------|------|---------|-------------| +| chunk_a | IMAGE | — | First chunk | +| chunk_b | IMAGE | — | Second chunk | +| trim_b_start | INT | 1 | Frames to trim from B's start (1 = drop duplicate ref frame) | + +### Get Frame By Index + +Extract a single frame. Use `-1` for the last frame (rolling reference). + +### Get Frame Range + +Extract a range of frames. Supports negative indexing. `end=0` means "to the end". + +### Video Chunk Planner + +Outputs the total number of chunks needed and a detailed text plan showing frame ranges and workflow steps. + +### Wan Chunked I2V Sampler ♾️ + +All-in-one node for Wan I2V models. See the [Workflows](#all-in-one-wan-models) section above. + +### Wan Chunk Calculator 🧮 + +Simple math: calculates `total_frames = chunk_frames + (num_chunks - 1) * (chunk_frames - 1)` with 4n+1 normalization. + +--- + +## Tips + +- **Overlap of 16 frames** works well for most cases (about 1 second at 16fps) +- **ease_in_out blend curve** gives the smoothest pixel-space transitions +- **hard_cut in latent space** is usually best — rolling reference already makes the overlap zone match +- **slerp** is the gold standard for latent interpolation if you need blending +- For Wan models, **chunk_frames must be 4n+1** (5, 9, 13, ..., 77, 81, 85, ...). The nodes auto-normalize this. +- Use **Video Chunk Planner** first to understand how your video will be divided + +--- + +## Requirements + +- **ComfyUI** (any recent version) +- **PyTorch** (included with ComfyUI) +- **[ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper)** — only needed for the 2 Wan-specific nodes. The 7 core nodes work without it. + +--- + +## License + +MIT — see [LICENSE](LICENSE) + +--- + +## Credits + +Built by [Greg Tee](https://github.com/gregtee2) for the ComfyUI community. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..181894a --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +# MIT License — Copyright (c) 2026 Greg Tee — see LICENSE file +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# WanChunkedI2VSampler — requires ComfyUI-WanVideoWrapper at runtime (not import time). +# If the import fails for any reason, the base VideoChunkTools nodes still work fine. +try: + from .wan_chunked_sampler import ( + NODE_CLASS_MAPPINGS as _WAN_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as _WAN_NAMES, + ) + NODE_CLASS_MAPPINGS.update(_WAN_MAPPINGS) + NODE_DISPLAY_NAME_MAPPINGS.update(_WAN_NAMES) +except Exception: + pass + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..03aeacd --- /dev/null +++ b/nodes.py @@ -0,0 +1,635 @@ +# MIT License — Copyright (c) 2026 Greg Tee — see LICENSE file +""" +ComfyUI_VideoChunkTools — Utility nodes for generating long videos +by splitting into overlapping chunks with rolling reference frames. + +Designed to solve the "context window reversal" problem where video +generation models (like Wan/FantasyPortrait) revert to the starting +state after ~135 frames because the reference image embedding pulls +the generation back. + +WORKFLOW (Rolling Reference for FantasyPortrait): +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +1. Load your full driving video (e.g., 300 frames) +2. ExtractVideoChunk(driving, chunk_index=0, chunk_frames=81, overlap=16) + → driving_chunk_0 (frames 0-80) +3. Run FantasyPortrait pipeline with ORIGINAL portrait as reference + → generated_chunk_0 +4. GetFrameByIndex(generated_chunk_0, index=-1) → last_frame +5. ExtractVideoChunk(driving, chunk_index=1, ...) → driving_chunk_1 +6. Run FantasyPortrait pipeline with last_frame as NEW reference + → generated_chunk_1 +7. BlendVideoChunks(generated_chunk_0, generated_chunk_1, overlap=16) + → seamless output +8. Repeat steps 4-7 for additional chunks + +The key insight: by using the LAST generated frame as the reference +for the next chunk, the model's identity stays consistent but the +pose/action can progress forward naturally without reversal. +""" + +import math +import torch + + +class BlendVideoChunks: + """ + Crossfade-blend two video chunks that share overlapping frames + at the boundary into a single continuous video. + + chunk_a ends where chunk_b begins — they share 'overlap_frames' + frames at the junction. This node smoothly transitions between them. + + Example: + chunk_a = 81 frames (frames 0-80 of final video) + chunk_b = 81 frames (frames 65-145 of final video) + overlap = 16 frames + Result = 146 frames: A[0:65] + crossfade[65:81] + B[16:81] + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "chunk_a": ("IMAGE", { + "tooltip": "First video chunk (the earlier segment)" + }), + "chunk_b": ("IMAGE", { + "tooltip": "Second video chunk (continues from chunk_a). " + "The first 'overlap_frames' of chunk_b overlap " + "with the last 'overlap_frames' of chunk_a." + }), + "overlap_frames": ("INT", { + "default": 16, "min": 1, "max": 200, "step": 1, + "tooltip": "Number of frames shared between the two chunks. " + "These frames will be crossfaded." + }), + "blend_curve": (["linear", "ease_in_out", "sigmoid"], { + "default": "ease_in_out", + "tooltip": "Blend curve shape: linear = straight ramp, " + "ease_in_out = smooth S-curve (3t²-2t³, best for most cases), " + "sigmoid = steep S-curve (snappier transition)" + }), + }, + } + + RETURN_TYPES = ("IMAGE", "INT",) + RETURN_NAMES = ("video", "total_frames",) + FUNCTION = "blend" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ("Crossfade-blend two overlapping video chunks into a single " + "seamless video. Use with ExtractVideoChunk and GetFrameByIndex " + "for rolling-reference long video generation.") + + def blend(self, chunk_a, chunk_b, overlap_frames, blend_curve): + a_len = chunk_a.shape[0] + b_len = chunk_b.shape[0] + + if overlap_frames >= a_len: + raise ValueError( + f"overlap_frames ({overlap_frames}) must be less than " + f"chunk_a length ({a_len})" + ) + if overlap_frames >= b_len: + raise ValueError( + f"overlap_frames ({overlap_frames}) must be less than " + f"chunk_b length ({b_len})" + ) + + # Split into non-overlapping + overlapping regions + a_unique = chunk_a[:a_len - overlap_frames] # frames only in A + a_overlap = chunk_a[a_len - overlap_frames:] # last N frames of A + b_overlap = chunk_b[:overlap_frames] # first N frames of B + b_unique = chunk_b[overlap_frames:] # frames only in B + + # Build blend weights (0→1 = A→B transition) + t = torch.linspace(0.0, 1.0, overlap_frames, device=chunk_a.device, + dtype=chunk_a.dtype) + + if blend_curve == "ease_in_out": + # Hermite smoothstep: 3t² - 2t³ (smooth acceleration/deceleration) + t = t * t * (3.0 - 2.0 * t) + elif blend_curve == "sigmoid": + # Steep logistic S-curve (sharper midpoint transition) + t = torch.sigmoid((t - 0.5) * 12.0) + t = (t - t[0]) / (t[-1] - t[0]) # renormalize to exact 0-1 + + # Reshape for broadcasting: (overlap_frames, 1, 1, 1) over (H, W, C) + weights = t.view(-1, 1, 1, 1) + + # Crossfade: A*(1-w) + B*w + blended = a_overlap * (1.0 - weights) + b_overlap * weights + + # Concatenate: [A unique] + [crossfade zone] + [B unique] + result = torch.cat([a_unique, blended, b_unique], dim=0) + + return (result, result.shape[0],) + + +class ExtractVideoChunk: + """ + Extract a specific chunk from a longer video, with overlap support. + + Divides a video into chunks of 'chunk_frames' length, where adjacent + chunks overlap by 'overlap_frames'. Use chunk_index to get each chunk. + + Chunk layout (e.g., 300 frames, chunk=81, overlap=16): + Chunk 0: frames 0-80 (81 frames) + Chunk 1: frames 65-145 (81 frames, starts 16 before chunk 0 ends) + Chunk 2: frames 130-210 (81 frames) + Chunk 3: frames 195-275 (81 frames) + Chunk 4: frames 219-299 (last chunk, may start closer if near end) + + For Wan models: chunk_frames should follow ((n-1)//4)*4+1 rule + (i.e., 5, 9, 13, ..., 77, 81, 85, ..., 129, 133, 137) + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", { + "tooltip": "Full driving video / image sequence to chunk" + }), + "chunk_index": ("INT", { + "default": 0, "min": 0, "max": 1000, + "tooltip": "Which chunk to extract (0-based)" + }), + "chunk_frames": ("INT", { + "default": 81, "min": 2, "max": 2000, "step": 1, + "tooltip": "Number of frames per chunk. " + "For Wan models, use values like 81, 97, 113, 129 " + "(must be 1 mod 4: ((n-1)//4)*4+1)" + }), + "overlap_frames": ("INT", { + "default": 16, "min": 0, "max": 200, "step": 1, + "tooltip": "Number of frames that overlap between adjacent chunks. " + "These will be crossfaded by BlendVideoChunks." + }), + }, + } + + RETURN_TYPES = ("IMAGE", "INT", "INT", "BOOLEAN",) + RETURN_NAMES = ("chunk", "total_chunks", "chunk_index", "is_last_chunk",) + FUNCTION = "extract" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ("Extract chunk N from a video for rolling-reference generation. " + "Chunks overlap by overlap_frames so they can be blended later.") + + def extract(self, images, chunk_index, chunk_frames, overlap_frames): + total_frames = images.shape[0] + + if overlap_frames >= chunk_frames: + raise ValueError( + f"overlap_frames ({overlap_frames}) must be less than " + f"chunk_frames ({chunk_frames})" + ) + + stride = chunk_frames - overlap_frames + + # Calculate total number of chunks needed to cover all frames + if total_frames <= chunk_frames: + total_chunks = 1 + else: + total_chunks = 1 + math.ceil((total_frames - chunk_frames) / stride) + + # Calculate start index for this chunk + start = chunk_index * stride + end = start + chunk_frames + + # Clamp to video bounds — last chunk may start earlier to stay full-size + if end > total_frames: + end = total_frames + start = max(0, end - chunk_frames) + + is_last = chunk_index >= (total_chunks - 1) + + chunk = images[start:end] + + return (chunk, total_chunks, chunk_index, is_last,) + + +class GetFrameByIndex: + """ + Extract a single frame from a video / image batch by index. + + Supports negative indexing: + -1 = last frame (default — use for rolling reference) + -2 = second-to-last + 0 = first frame + 5 = sixth frame + + Primary use case: Get the last frame of a generated chunk + to use as the reference image for the next chunk. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", { + "tooltip": "Image batch / video frames" + }), + "index": ("INT", { + "default": -1, "min": -10000, "max": 10000, + "tooltip": "Frame index to extract. Use -1 for last frame " + "(ideal for rolling reference), -2 for second-to-last, " + "0 for first frame, etc." + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("frame",) + FUNCTION = "get_frame" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ("Get a single frame from a video batch by index. " + "Use -1 (last frame) as the reference for the next chunk " + "in a rolling-reference workflow.") + + def get_frame(self, images, index): + total = images.shape[0] + + # Handle negative indexing (Python-style) + if index < 0: + index = total + index + + # Clamp to valid range + index = max(0, min(index, total - 1)) + + # Return as single-frame batch (1, H, W, C) + return (images[index:index + 1],) + + +class GetFrameRange: + """ + Extract a range of frames from a video / image batch. + + Supports negative indexing for both start and end. + The end index is EXCLUSIVE (Python slice convention). + + Examples: + start=0, end=16 → first 16 frames + start=-16, end=0 → last 16 frames (end=0 means "to the end") + start=10, end=30 → frames 10-29 + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", { + "tooltip": "Image batch / video frames" + }), + "start": ("INT", { + "default": 0, "min": -10000, "max": 10000, + "tooltip": "Start index (inclusive). Negative = from end." + }), + "end": ("INT", { + "default": 0, "min": -10000, "max": 10000, + "tooltip": "End index (exclusive). 0 = to the end of the video. " + "Negative = from end." + }), + }, + } + + RETURN_TYPES = ("IMAGE", "INT",) + RETURN_NAMES = ("frames", "frame_count",) + FUNCTION = "get_range" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ("Extract a range of frames from a video batch. " + "Supports negative indexing. End=0 means 'to the end'.") + + def get_range(self, images, start, end): + total = images.shape[0] + + # Handle negative indexing + if start < 0: + start = total + start + start = max(0, min(start, total - 1)) + + # end=0 means "to the end" + if end == 0: + end = total + elif end < 0: + end = total + end + end = max(start + 1, min(end, total)) + + result = images[start:end] + return (result, result.shape[0],) + + +class VideoChunkPlanner: + """ + Calculate the chunking plan for a video — outputs how many chunks + are needed and the frame ranges, given chunk_frames and overlap. + + Useful for planning how many times to run the pipeline and + understanding the layout before starting generation. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "total_frames": ("INT", { + "default": 300, "min": 1, "max": 100000, + "tooltip": "Total number of frames in your driving video" + }), + "chunk_frames": ("INT", { + "default": 81, "min": 2, "max": 2000, + "tooltip": "Frames per chunk (e.g., 81 for Wan models)" + }), + "overlap_frames": ("INT", { + "default": 16, "min": 0, "max": 200, + "tooltip": "Frames of overlap between adjacent chunks" + }), + }, + } + + RETURN_TYPES = ("INT", "STRING",) + RETURN_NAMES = ("total_chunks", "plan_info",) + FUNCTION = "plan" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ("Plan the chunking strategy for long video generation. " + "Shows total chunks needed and frame ranges for each.") + + def plan(self, total_frames, chunk_frames, overlap_frames): + if overlap_frames >= chunk_frames: + raise ValueError( + f"overlap_frames ({overlap_frames}) must be less than " + f"chunk_frames ({chunk_frames})" + ) + + stride = chunk_frames - overlap_frames + + if total_frames <= chunk_frames: + total_chunks = 1 + final_frames = total_frames + else: + total_chunks = 1 + math.ceil((total_frames - chunk_frames) / stride) + # output video length = total_frames without overlap duplication + final_frames = chunk_frames + (total_chunks - 1) * stride + + lines = [] + lines.append(f"=== Video Chunk Plan ===") + lines.append(f"Source video: {total_frames} frames") + lines.append(f"Chunk size: {chunk_frames} frames") + lines.append(f"Overlap: {overlap_frames} frames") + lines.append(f"Stride: {stride} frames") + lines.append(f"Total chunks needed: {total_chunks}") + lines.append(f"Output video length: ~{total_frames} frames") + lines.append(f"") + lines.append(f"--- Chunk Layout ---") + + for i in range(total_chunks): + start = i * stride + end = start + chunk_frames + if end > total_frames: + end = total_frames + start = max(0, end - chunk_frames) + + ref = "original reference" if i == 0 else f"last frame of chunk {i-1}" + lines.append( + f" Chunk {i}: frames {start:>4d}-{end-1:>4d} " + f"({end - start} frames) | ref: {ref}" + ) + + lines.append(f"") + lines.append(f"--- Workflow Steps ---") + lines.append(f"1. Load driving video ({total_frames} frames)") + for i in range(total_chunks): + start = i * stride + end = min(start + chunk_frames, total_frames) + if end > total_frames: + start = max(0, total_frames - chunk_frames) + if i == 0: + lines.append(f"2. ExtractVideoChunk(index=0) → Generate with original reference") + else: + lines.append( + f"{2+i}. GetFrameByIndex(chunk_{i-1}, -1) → New reference → " + f"ExtractVideoChunk(index={i}) → Generate" + ) + lines.append(f"{2+total_chunks}. Chain BlendVideoChunks for each pair") + + plan_text = "\n".join(lines) + return (total_chunks, plan_text,) + + +class BlendLatentChunks: + """ + Join two video chunk LATENTS along the temporal dimension BEFORE decoding. + + The overlap_frames parameter is in LATENT temporal space (not pixel frames). + For Wan models with 4x temporal compression: + latent_overlap = pixel_overlap / 4 + e.g. 16 pixel frames → 4 latent frames + + Blend modes: + hard_cut — No blending. Cuts at the midpoint of the overlap region. + Best for rolling-reference workflows where both chunks + generated similar content in the overlap zone. + slerp — Spherical linear interpolation. Preserves latent vector + magnitude (standard for diffusion model interpolation). + Avoids the norm-collapse that causes washed-out frames. + linear — Simple linear crossfade in latent space. + ease_in_out — Hermite smooth-step crossfade. + sigmoid — Steep S-curve crossfade. + + Wire BOTH sampler latent outputs into this node, then feed the result + into a SINGLE VAE decode for the final seamless video. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "latent_a": ("LATENT",), + "latent_b": ("LATENT",), + "overlap_frames": ("INT", { + "default": 4, + "min": 1, + "max": 64, + "step": 1, + "tooltip": "Overlap in LATENT temporal frames. For Wan: pixel_overlap / 4 (e.g. 16px = 4 latent)" + }), + "blend_curve": (["hard_cut", "slerp", "linear", "ease_in_out", "sigmoid"], { + "default": "hard_cut", + "tooltip": "hard_cut = clean cut at midpoint (best for rolling reference). " + "slerp = spherical interpolation (preserves latent magnitude). " + "linear/ease_in_out/sigmoid = alpha-blend crossfades." + }), + } + } + + RETURN_TYPES = ("LATENT", "INT",) + RETURN_NAMES = ("latent", "total_latent_frames",) + FUNCTION = "blend" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ( + "Join two video chunk latents along temporal dimension before decoding. " + "hard_cut (recommended): clean cut at overlap midpoint — no dissolve. " + "slerp: spherical interpolation preserving latent magnitude. " + "overlap_frames is in LATENT space (Wan: pixel_overlap / 4)." + ) + + @staticmethod + def _slerp_frame(t, a, b): + """Spherical linear interpolation between two latent frames. + Preserves the norm of the vectors — standard for diffusion latent spaces. + a, b: [B, C, H, W] tensors. t: scalar 0-1.""" + a_f = a.flatten().float() + b_f = b.flatten().float() + a_norm = torch.nn.functional.normalize(a_f, dim=0) + b_norm = torch.nn.functional.normalize(b_f, dim=0) + dot = (a_norm * b_norm).sum().clamp(-1.0, 1.0) + omega = torch.acos(dot) + # Fall back to lerp for nearly-parallel vectors + if omega.abs() < 1e-5: + return (a.float() * (1.0 - t) + b.float() * t).to(a.dtype) + sin_omega = torch.sin(omega) + wa = torch.sin((1.0 - t) * omega) / sin_omega + wb = torch.sin(t * omega) / sin_omega + return (a.float() * wa + b.float() * wb).to(a.dtype) + + def blend(self, latent_a, latent_b, overlap_frames, blend_curve): + a = latent_a["samples"] + b = latent_b["samples"] + + if a.dim() != 5 or b.dim() != 5: + raise ValueError( + f"Expected 5D video latents [B,C,T,H,W], " + f"got {a.dim()}D and {b.dim()}D. " + f"This node is for VIDEO latents, not image latents." + ) + + T_a = a.shape[2] + T_b = b.shape[2] + overlap = min(overlap_frames, T_a, T_b) + + # Slice the latent tensors + a_keep = a[:, :, :T_a - overlap, :, :] + a_over = a[:, :, T_a - overlap:, :, :] + b_over = b[:, :, :overlap, :, :] + b_keep = b[:, :, overlap:, :, :] + + if blend_curve == "hard_cut": + # Clean cut at the midpoint — no interpolation at all. + # With rolling reference, both chunks have similar content here, + # so a clean cut is virtually invisible. + mid = overlap // 2 + blended = torch.cat([a_over[:, :, :mid, :, :], + b_over[:, :, mid:, :, :]], dim=2) + + elif blend_curve == "slerp": + # Spherical linear interpolation — preserves latent vector norms. + # Standard technique for interpolating in diffusion latent spaces. + # Avoids the washed-out frames caused by lerp norm-collapse. + frames = [] + for i in range(overlap): + t = i / max(overlap - 1, 1) + frame = self._slerp_frame(t, a_over[:, :, i, :, :], + b_over[:, :, i, :, :]) + frames.append(frame.unsqueeze(2)) + blended = torch.cat(frames, dim=2) + + else: + # Alpha-blend crossfades (linear, ease_in_out, sigmoid) + weights = torch.linspace(0.0, 1.0, overlap, + device=a.device, dtype=a.dtype) + if blend_curve == "ease_in_out": + weights = weights * weights * (3.0 - 2.0 * weights) + elif blend_curve == "sigmoid": + weights = torch.sigmoid((weights - 0.5) * 10.0) + weights = (weights - weights[0]) / (weights[-1] - weights[0]) + weights = weights.reshape(1, 1, -1, 1, 1) + blended = a_over * (1.0 - weights) + b_over * weights + + # Concatenate: a_unique + overlap_region + b_unique + result = torch.cat([a_keep, blended, b_keep], dim=2) + + # Preserve any extra keys from latent_a (noise_mask, batch_index, etc.) + out = {k: v for k, v in latent_a.items() if k != "samples"} + out["samples"] = result + + return (out, int(result.shape[2]),) + + +class ConcatVideoChunks: + """ + Concatenate two decoded video chunks into a single continuous video. + + For rolling-reference workflows where chunk B was generated using chunk A's + last frame as the CLIP Vision reference: + - Chunk B's FIRST frame closely matches chunk A's LAST frame + - Set trim_b_start=1 to drop chunk B's duplicate first frame + - Result: A[all frames] + B[frame 1 onward] — clean seamless join + + No blending, no overlap, no dissolve — just a clean concatenation. + The Wan I2V model naturally makes chunk B's first frame match the reference, + so the boundary is virtually invisible. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "chunk_a": ("IMAGE", { + "tooltip": "First video chunk (decoded pixels)" + }), + "chunk_b": ("IMAGE", { + "tooltip": "Second video chunk (decoded pixels). " + "Its first frame should match chunk_a's last frame " + "if using rolling reference." + }), + "trim_b_start": ("INT", { + "default": 1, "min": 0, "max": 100, "step": 1, + "tooltip": "Number of frames to trim from the START of chunk_b " + "before concatenating. Default 1 = drop B's first frame " + "(which duplicates A's last frame in rolling-reference workflows)." + }), + }, + } + + RETURN_TYPES = ("IMAGE", "INT",) + RETURN_NAMES = ("video", "total_frames",) + FUNCTION = "concat" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ( + "Concatenate two decoded video chunks. Optionally trim B's first N frames " + "to remove the duplicate reference frame. Simplest join method — no blending needed " + "when using rolling reference (chunk B's first frame already matches A's last)." + ) + + def concat(self, chunk_a, chunk_b, trim_b_start): + if trim_b_start > 0: + chunk_b = chunk_b[trim_b_start:] + + if chunk_b.shape[0] == 0: + return (chunk_a, chunk_a.shape[0],) + + result = torch.cat([chunk_a, chunk_b], dim=0) + return (result, result.shape[0],) + + +# ---------- Registration ---------- + +NODE_CLASS_MAPPINGS = { + "BlendVideoChunks": BlendVideoChunks, + "BlendLatentChunks": BlendLatentChunks, + "ConcatVideoChunks": ConcatVideoChunks, + "ExtractVideoChunk": ExtractVideoChunk, + "GetFrameByIndex": GetFrameByIndex, + "GetFrameRange": GetFrameRange, + "VideoChunkPlanner": VideoChunkPlanner, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "BlendVideoChunks": "Blend Video Chunks (Crossfade)", + "BlendLatentChunks": "Blend Latent Chunks (Pre-Decode)", + "ConcatVideoChunks": "Concat Video Chunks", + "ExtractVideoChunk": "Extract Video Chunk", + "GetFrameByIndex": "Get Frame By Index", + "GetFrameRange": "Get Frame Range", + "VideoChunkPlanner": "Video Chunk Planner", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..46b84be --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,9 @@ +[project] +name = "ComfyUI_VideoChunkTools" +version = "1.0.0" +description = "ComfyUI nodes for generating long videos by splitting into overlapping chunks with rolling reference frames." +license = {text = "MIT"} +requires-python = ">=3.10" + +[project.urls] +Repository = "https://github.com/gregtee2/ComfyUI_VideoChunkTools" diff --git a/wan_chunked_sampler.py b/wan_chunked_sampler.py new file mode 100644 index 0000000..47c4269 --- /dev/null +++ b/wan_chunked_sampler.py @@ -0,0 +1,753 @@ +# MIT License — Copyright (c) 2026 Greg Tee — see LICENSE file +""" +WanChunkedI2VSampler — All-in-one rolling-reference chunked video generation. + +Internally handles: encode → sample (two-pass high/low noise) → decode → +extract reference frame → repeat... for as many chunks as needed to produce +`total_frames` of seamless video from a single node. + +Replicates the proven v04 concat approach: + - Each chunk's first frame naturally matches the previous chunk's last frame + (due to I2V conditioning on the reference image). + - We drop the overlapping first frame and concatenate — zero artifacts. + +Requires: ComfyUI-WanVideoWrapper (kijai) +""" + +import math +import importlib +import torch +import gc + +from comfy.utils import ProgressBar + +# --------------------------------------------------------------------------- +# Lazy import of WanVideoWrapper node classes from ComfyUI's global registry. +# We don't import at module level so that: +# 1. VideoChunkTools can load even if WanVideoWrapper is missing. +# 2. By the time process() runs, all custom nodes are fully registered. +# --------------------------------------------------------------------------- +_wan_cache = {} + +def _wan_class(name): + """Get a WanVideoWrapper node class from ComfyUI's global NODE_CLASS_MAPPINGS.""" + if name in _wan_cache: + return _wan_cache[name] + try: + comfy_nodes = importlib.import_module("nodes") # ComfyUI root nodes.py + cls = comfy_nodes.NODE_CLASS_MAPPINGS[name] + _wan_cache[name] = cls + return cls + except (KeyError, AttributeError, ImportError) as exc: + raise ImportError( + f"WanChunkedI2VSampler requires ComfyUI-WanVideoWrapper. " + f"Node class '{name}' not found. Error: {exc}" + ) + +# --------------------------------------------------------------------------- +# Scheduler list — dynamically read from WanVideoSampler if available, +# otherwise fall back to a hardcoded subset. +# --------------------------------------------------------------------------- +_FALLBACK_SCHEDULERS = [ + "euler", "euler/beta", + "unipc", "unipc/beta", + "dpm++", "dpm++/beta", + "dpm++_sde", "dpm++_sde/beta", + "deis", + "lcm", "lcm/beta", + "longcat_distill_euler", + "flowmatch_distill", + "flowmatch_causvid", + "res_multistep", +] + +_scheduler_list_cache = None + +def _get_scheduler_list(): + """Return the scheduler list, preferring WanVideoWrapper's live list.""" + global _scheduler_list_cache + if _scheduler_list_cache is not None: + return _scheduler_list_cache + try: + comfy_nodes = importlib.import_module("nodes") + sampler_cls = comfy_nodes.NODE_CLASS_MAPPINGS.get("WanVideoSampler") + if sampler_cls: + inputs = sampler_cls.INPUT_TYPES() + _scheduler_list_cache = list(inputs["required"]["scheduler"][0]) + return _scheduler_list_cache + except Exception: + pass + _scheduler_list_cache = list(_FALLBACK_SCHEDULERS) + return _scheduler_list_cache + +# --------------------------------------------------------------------------- +# Image Resizing Helper +# --------------------------------------------------------------------------- +def resize_image_to_target(image, target_width, target_height, keep_proportion, crop_position): + """ + Resizes an image tensor (B, H, W, C) to target_width and target_height. + Supports stretch, crop, and pad modes. + """ + import torch.nn.functional as F + from comfy.utils import common_upscale + + B, H, W, C = image.shape + + if W == target_width and H == target_height: + return image + + if keep_proportion == "stretch": + image = image.movedim(-1, 1) + image = common_upscale(image, target_width, target_height, "bicubic", "disabled") + image = image.movedim(1, -1) + return image + + elif keep_proportion == "crop": + old_aspect = W / H + new_aspect = target_width / target_height + if old_aspect > new_aspect: + crop_w = round(H * new_aspect) + crop_h = H + else: + crop_w = W + crop_h = round(W / new_aspect) + + if crop_position == "center": + x = (W - crop_w) // 2 + y = (H - crop_h) // 2 + elif crop_position == "top": + x = (W - crop_w) // 2 + y = 0 + elif crop_position == "bottom": + x = (W - crop_w) // 2 + y = H - crop_h + elif crop_position == "left": + x = 0 + y = (H - crop_h) // 2 + elif crop_position == "right": + x = W - crop_w + y = (H - crop_h) // 2 + + image = image[:, y:y+crop_h, x:x+crop_w, :] + image = image.movedim(-1, 1) + image = common_upscale(image, target_width, target_height, "bicubic", "disabled") + image = image.movedim(1, -1) + return image + + elif keep_proportion == "pad": + ratio = min(target_width / W, target_height / H) + new_width = round(W * ratio) + new_height = round(H * ratio) + + image = image.movedim(-1, 1) + image = common_upscale(image, new_width, new_height, "bicubic", "disabled") + image = image.movedim(1, -1) + + pad_left = pad_right = pad_top = pad_bottom = 0 + if crop_position == "center": + pad_left = (target_width - new_width) // 2 + pad_right = target_width - new_width - pad_left + pad_top = (target_height - new_height) // 2 + pad_bottom = target_height - new_height - pad_top + elif crop_position == "top": + pad_left = (target_width - new_width) // 2 + pad_right = target_width - new_width - pad_left + pad_top = 0 + pad_bottom = target_height - new_height + elif crop_position == "bottom": + pad_left = (target_width - new_width) // 2 + pad_right = target_width - new_width - pad_left + pad_top = target_height - new_height + pad_bottom = 0 + elif crop_position == "left": + pad_left = 0 + pad_right = target_width - new_width + pad_top = (target_height - new_height) // 2 + pad_bottom = target_height - new_height - pad_top + elif crop_position == "right": + pad_left = target_width - new_width + pad_right = 0 + pad_top = (target_height - new_height) // 2 + pad_bottom = target_height - new_height - pad_top + + if pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0: + image = image.movedim(-1, 1) + image = F.pad(image, (pad_left, pad_right, pad_top, pad_bottom), mode='constant', value=0.0) + image = image.movedim(1, -1) + + return image + + return image + + +# ═══════════════════════════════════════════════════════════════════════════ +# WanChunkedI2VSampler +# ═══════════════════════════════════════════════════════════════════════════ + +class WanChunkedI2VSampler: + """ + All-in-one chunked video generator for Wan I2V two-pass distilled models. + + Set **total_frames** and **chunk_frames** — the node handles everything: + encode → sample (high noise pass → low noise pass) → decode → + extract last frame as reference → repeat for the next chunk. + + Each chunk naturally conditions on the previous chunk's last frame, + so we just trim the overlapping first frame and concatenate. + + Supports two modes: + • **Single-pass**: Connect one model. All steps run in one shot. + • **Two-pass**: Connect a second model to model_b. + model runs steps 0→split_step, model_b runs split_step→end. + (Useful for distilled workflows with different LoRA weights per pass.) + """ + + @classmethod + def INPUT_TYPES(s): + schedulers = _get_scheduler_list() + return { + "required": { + "model": ("WANVIDEOMODEL", { + "tooltip": "Wan I2V model (2.1 or 2.2). Runs all steps in single-pass. In two-pass mode, runs the first steps (0 → split_step)." + }), + "vae": ("WANVAE", { + "tooltip": "Wan VAE model" + }), + "start_image": ("IMAGE", { + "tooltip": "Reference image for the first chunk" + }), + "total_frames": ("INT", { + "default": 161, "min": 5, "max": 9999, "step": 1, + "tooltip": "Total frames to generate. Output may be slightly more (rounded to fill the last chunk), then trimmed." + }), + "chunk_frames": ("INT", { + "default": 81, "min": 5, "max": 241, "step": 4, + "tooltip": "Frames per chunk. Auto-normalized to Wan's 4n+1 rule (5, 9, 13, … 77, 81, 85 …)." + }), + "width": ("INT", { + "default": 832, "min": 128, "max": 2048, "step": 16, + "tooltip": "Output video width in pixels." + }), + "height": ("INT", { + "default": 480, "min": 128, "max": 2048, "step": 16, + "tooltip": "Output video height in pixels." + }), + "keep_proportion": (["stretch", "crop", "pad"], { + "default": "crop", + "tooltip": "How to resize the input images to match the target width/height." + }), + "crop_position": (["center", "top", "bottom", "left", "right"], { + "default": "center", + "tooltip": "Where to crop or pad the image." + }), + "steps": ("INT", { + "default": 8, "min": 1, "max": 100, "step": 1, + "tooltip": "Total denoising steps." + }), + "cfg": ("FLOAT", { + "default": 2.0, "min": 0.0, "max": 30.0, "step": 0.01, + "tooltip": "CFG guidance scale. In two-pass mode this applies to the first pass only." + }), + "shift": ("FLOAT", { + "default": 11.0, "min": 0.0, "max": 1000.0, "step": 0.01, + "tooltip": "Scheduler shift parameter." + }), + "seed": ("INT", { + "default": 0, "min": 0, "max": 0xffffffffffffffff, + "tooltip": "Random seed. Same seed used for every chunk — variation comes from different reference images." + }), + "force_offload": ("BOOLEAN", { + "default": True, + "tooltip": "Offload model to CPU after each sampling pass. Highly recommended." + }), + "scheduler": (schedulers, { + "default": "euler", + "tooltip": "Noise scheduler algorithm." + }), + "enable_vae_tiling": ("BOOLEAN", { + "default": False, + "tooltip": "Enable VAE tiling for lower-VRAM decode. May introduce tile seam artifacts." + }), + }, + "optional": { + "model_b": ("WANVIDEOMODEL", { + "tooltip": "Optional second model for two-pass sampling. Connect this to split denoising across two models (e.g. different LoRA weights). Leave disconnected for single-pass." + }), + "split_step": ("INT", { + "default": 4, "min": 1, "max": 99, "step": 1, + "tooltip": "Two-pass only: model runs steps 0 → split_step, model_b runs split_step → end. Ignored in single-pass." + }), + "cfg_b": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 30.0, "step": 0.01, + "tooltip": "Two-pass only: CFG for the second pass (model_b). Ignored in single-pass." + }), + "text_embeds": ("WANVIDEOTEXTEMBEDS", { + "tooltip": "Text conditioning from WanVideoTextEncode." + }), + "clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", { + "tooltip": "Optional CLIP vision embeddings from WanVideoClipVisionEncode." + }), + "feta_args": ("FETAARGS", { + "tooltip": "FETA arguments. Recommended: weight=2, start=0, end=1." + }), + "context_options": ("WANVIDCONTEXT", { + "tooltip": "Context windowing options. In two-pass mode only applied to the first pass (model)." + }), + "noise_aug_strength": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Noise augmentation for I2V encoding. 0 = no augmentation." + }), + "start_latent_strength": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Strength of the start image in latent space." + }), + "end_latent_strength": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "FLF only: Strength of the end image conditioning in latent space. " + "Controls how strongly the last chunk is pulled toward the end image. " + "Only effective when end_image is connected." + }), + "riflex_freq_index": ("INT", { + "default": 0, "min": 0, "max": 1000, "step": 1, + "tooltip": "RIFLEX frequency index. 0 = disabled." + }), + "tile_x": ("INT", { + "default": 272, "min": 40, "max": 4096, "step": 8, + "tooltip": "VAE decode tile width (only when tiling enabled)." + }), + "tile_y": ("INT", { + "default": 272, "min": 40, "max": 4096, "step": 8, + "tooltip": "VAE decode tile height (only when tiling enabled)." + }), + "tile_stride_x": ("INT", { + "default": 144, "min": 32, "max": 2040, "step": 8, + "tooltip": "VAE decode tile stride X (only when tiling enabled)." + }), + "tile_stride_y": ("INT", { + "default": 128, "min": 32, "max": 2040, "step": 8, + "tooltip": "VAE decode tile stride Y (only when tiling enabled)." + }), + "end_image": ("IMAGE", { + "tooltip": "Optional target image for the last frame — enables FLF (First-Last-Frame) mode. " + "If a batch of images is provided, they will be distributed as target keyframes " + "across the chunks (e.g., 3 images for a 3-chunk video)." + }), + "end_blend_chunks": ("INT", { + "default": 0, "min": 0, "max": 40, "step": 1, + "tooltip": "FLF only: Number of FRAMES to crossfade at the FLF boundary. " + "The last chunk (with end_image) overlaps with the previous chunk " + "by this many frames, and a smooth alpha ramp blends between them. " + "0 = hard cut (no crossfade). Try 9–17 for smooth transitions " + "(~0.5–1 sec at 16fps). The last chunk is automatically enlarged " + "to the next valid 4n+1 size to accommodate the overlap." + }), + }, + } + + RETURN_TYPES = ("IMAGE", "INT", "INT",) + RETURN_NAMES = ("video", "total_frames", "chunks_generated",) + FUNCTION = "process" + CATEGORY = "VideoChunkTools" + DESCRIPTION = ( + "All-in-one chunked video generator for Wan I2V models.\n\n" + "Set total_frames and chunk_frames — the node handles everything:\n" + "encode → sample → decode → extract reference → repeat.\n\n" + "Supports single-pass (one model) or two-pass (connect model_b).\n\n" + "FLF (First-Last-Frame): Connect an end_image to make the video " + "transition toward a target end frame. For multi-chunk videos, " + "set end_blend_chunks (e.g. 9–17) for a smooth crossfade overlap.\n\n" + "Uses the proven concat approach: each chunk's first frame naturally " + "matches the previous chunk's last frame, so we just drop it and concatenate.\n\n" + "Requires: ComfyUI-WanVideoWrapper" + ) + + def process( + self, + model, vae, start_image, + total_frames, chunk_frames, width, height, + keep_proportion, crop_position, + steps, cfg, + shift, seed, force_offload, scheduler, + enable_vae_tiling, + # Optional + model_b=None, + split_step=4, + cfg_b=1.0, + text_embeds=None, + clip_embeds=None, + feta_args=None, + context_options=None, + noise_aug_strength=0.0, + start_latent_strength=1.0, + end_latent_strength=1.0, + riflex_freq_index=0, + tile_x=272, tile_y=272, + tile_stride_x=144, tile_stride_y=128, + end_image=None, + end_blend_chunks=0, + ): + # ── Determine mode ── + two_pass = model_b is not None + + if two_pass and split_step >= steps: + raise ValueError( + f"Two-pass mode: split_step ({split_step}) must be less than steps ({steps}). " + f"High noise runs steps 0→{split_step}, low noise runs {split_step}→{steps}." + ) + + # ── Resize Input Images ── + start_image = resize_image_to_target(start_image, width, height, keep_proportion, crop_position) + if end_image is not None: + end_image = resize_image_to_target(end_image, width, height, keep_proportion, crop_position) + + # ── Import WanVideoWrapper classes (lazy, from global registry) ── + Encoder = _wan_class("WanVideoImageToVideoEncode") + Sampler = _wan_class("WanVideoSampler") + Decoder = _wan_class("WanVideoDecode") + + encoder_inst = Encoder() + sampler_inst = Sampler() + decoder_inst = Decoder() + + # ── Normalize chunk_frames to Wan's 4n+1 rule ── + chunk_frames = ((chunk_frames - 1) // 4) * 4 + 1 + chunk_frames = max(chunk_frames, 5) + + # ── Estimate number of chunks ── + # First chunk contributes chunk_frames. + # Each subsequent chunk contributes (chunk_frames - 1) unique frames + # because its first frame overlaps with the previous chunk's last. + if total_frames <= chunk_frames: + num_chunks = 1 + else: + remaining = total_frames - chunk_frames + unique_per_extra_chunk = chunk_frames - 1 + num_chunks = 1 + math.ceil(remaining / max(1, unique_per_extra_chunk)) + + expected_output = chunk_frames + max(0, num_chunks - 1) * (chunk_frames - 1) + + # ── FLF crossfade: enlarge last chunk for smooth overlap ── + flf_crossfade = 0 + last_chunk_frames = chunk_frames + if end_image is not None and end_blend_chunks > 0 and num_chunks > 1: + flf_crossfade = min(end_blend_chunks, (chunk_frames - 1) // 2) + needed = chunk_frames + flf_crossfade + n = (needed - 1 + 3) // 4 # round up to next valid 4n+1 + last_chunk_frames = 4 * n + 1 + + mode_label = f"TWO-PASS (split at step {split_step})" if two_pass else "SINGLE-PASS" + print(f"\n{'=' * 65}") + print(f" WanChunkedI2V — Rolling-Reference Chunked Generation") + print(f"{'=' * 65}") + print(f" Mode : {mode_label}") + print(f" Requested frames : {total_frames}") + print(f" Chunk size : {chunk_frames} frames (4n+1)") + print(f" Chunks planned : {num_chunks}") + print(f" Expected output : {expected_output} frames (trimmed to {total_frames})") + print(f" Resolution : {width} × {height}") + if two_pass: + print(f" Steps : {steps} (model 0→{split_step}, model_b {split_step}→end)") + print(f" CFG : model={cfg}, model_b={cfg_b}") + else: + print(f" Steps : {steps}") + print(f" CFG : {cfg}") + print(f" Seed : {seed}") + print(f" Scheduler : {scheduler}") + print(f" Force offload : {force_offload}") + if end_image is not None: + B = end_image.shape[0] + if B > 1: + print(f" FLF Mode : ON (Multi-Keyframe: {B} target images)") + else: + print(f" FLF Mode : ON (end_image → last chunk only)") + if flf_crossfade > 0: + print(f" FLF Crossfade : {flf_crossfade} frames overlap") + print(f" Last chunk size : {last_chunk_frames} frames (enlarged for crossfade)") + else: + print(f" FLF Crossfade : 0 (hard cut at last chunk boundary)") + print(f" End Lat Strength : {end_latent_strength}") + else: + print(f" FLF Mode : OFF") + print(f"{'=' * 65}\n") + + # ── Progress bar (chunk-level) ── + pbar = ProgressBar(num_chunks) + + all_frames = [] + total_generated = 0 + reference_image = start_image + flf_saved_tail = None # saved frames for FLF crossfade + + for chunk_idx in range(num_chunks): + # Safety: stop if we already have enough frames + if total_generated >= total_frames: + break + + chunk_label = f"Chunk {chunk_idx + 1}/{num_chunks}" + print(f"\n ── {chunk_label} ──") + + # ============================================================ + # Step 1 — ENCODE reference image → image_embeds + # ============================================================ + is_last_chunk = (chunk_idx == num_chunks - 1) + + # Multi-Keyframe FLF: Map a batch of end_images to the chunks. + # If 1 image is provided, it applies to the last chunk. + # If N images are provided, they apply to the last N chunks. + chunk_end_image = None + if end_image is not None: + B = end_image.shape[0] + target_idx = B - num_chunks + chunk_idx + if 0 <= target_idx < B: + chunk_end_image = end_image[target_idx:target_idx+1] + + flf_this_chunk = (chunk_end_image is not None) + current_chunk_frames = last_chunk_frames if (is_last_chunk and flf_crossfade > 0) else chunk_frames + + encode_label = "Encoding ref + FLF end target" if flf_this_chunk else "Encoding reference image" + if flf_this_chunk and is_last_chunk and flf_crossfade > 0: + encode_label += f" ({current_chunk_frames}f enlarged)" + print(f" [1/4] {encode_label} …") + + encode_kwargs = dict( + width=width, + height=height, + num_frames=current_chunk_frames, + force_offload=force_offload, + noise_aug_strength=noise_aug_strength, + start_latent_strength=start_latent_strength, + end_latent_strength=end_latent_strength, + start_image=reference_image, + clip_embeds=clip_embeds, + vae=vae, + ) + if flf_this_chunk: + encode_kwargs["end_image"] = chunk_end_image + + image_embeds = encoder_inst.process(**encode_kwargs)[0] + + # ============================================================ + # Step 2 (+3) — SAMPLE + # ============================================================ + if two_pass: + # ── TWO-PASS: model → model_b ── + print(f" [2/4] Sampling pass 1 (steps 0 → {split_step}) …") + high_latent = sampler_inst.process( + model=model, + image_embeds=image_embeds, + steps=steps, + cfg=cfg, + shift=shift, + seed=seed, + force_offload=force_offload, + scheduler=scheduler, + riflex_freq_index=riflex_freq_index, + text_embeds=text_embeds, + feta_args=feta_args, + context_options=context_options, + start_step=0, + end_step=split_step, + add_noise_to_samples=False, + )[0] + + print(f" [3/4] Sampling pass 2 (steps {split_step} → end) …") + final_latent = sampler_inst.process( + model=model_b, + image_embeds=image_embeds, + steps=steps, + cfg=cfg_b, + shift=shift, + seed=seed, + force_offload=force_offload, + scheduler=scheduler, + riflex_freq_index=riflex_freq_index, + text_embeds=text_embeds, + samples=high_latent, + feta_args=feta_args, + context_options=None, + start_step=split_step, + end_step=-1, + add_noise_to_samples=False, + )[0] + + del high_latent + else: + # ── SINGLE-PASS: one model, all steps ── + print(f" [2/3] Sampling (all {steps} steps) …") + final_latent = sampler_inst.process( + model=model, + image_embeds=image_embeds, + steps=steps, + cfg=cfg, + shift=shift, + seed=seed, + force_offload=force_offload, + scheduler=scheduler, + riflex_freq_index=riflex_freq_index, + text_embeds=text_embeds, + feta_args=feta_args, + context_options=context_options, + add_noise_to_samples=False, + )[0] + + # ============================================================ + # Step 3/4 — DECODE latent → pixel frames + # ============================================================ + decode_label = "4/4" if two_pass else "3/3" + print(f" [{decode_label}] VAE decoding …") + chunk_images = decoder_inst.decode( + vae=vae, + samples=final_latent, + enable_vae_tiling=enable_vae_tiling, + tile_x=tile_x, + tile_y=tile_y, + tile_stride_x=tile_stride_x, + tile_stride_y=tile_stride_y, + )[0] # IMAGE tensor: (T, H, W, C) + + raw_frame_count = chunk_images.shape[0] + print(f" Decoded {raw_frame_count} frames") + + # ── Trim overlapping first frame (all chunks after the first) ── + if chunk_idx > 0: + chunk_images = chunk_images[1:] + print(f" Trimmed overlap → {chunk_images.shape[0]} unique frames") + + # ── FLF crossfade handling ── + is_penultimate = (chunk_idx == num_chunks - 2) + + if is_penultimate and flf_crossfade > 0 and end_image is not None: + # Save tail frames for crossfade with upcoming FLF chunk. + # Set reference to an earlier frame so the FLF chunk + # regenerates the overlap zone from that point forward. + cf = min(flf_crossfade, chunk_images.shape[0] - 2) + flf_saved_tail = chunk_images[-cf:].clone() + reference_image = chunk_images[-(cf + 1)].unsqueeze(0).clone() + + # Remove tail from this chunk's output (crossfade replaces it) + chunk_images = chunk_images[:-cf] + print(f" FLF: saved {cf} tail frames, ref set {cf+1} frames back") + + all_frames.append(chunk_images) + total_generated += chunk_images.shape[0] + + elif is_last_chunk and flf_crossfade > 0 and flf_saved_tail is not None: + # FLF chunk overlaps with saved tail — crossfade them. + cf = flf_saved_tail.shape[0] + flf_head = chunk_images[:cf].clone() + flf_rest = chunk_images[cf:] + + # Ease-in (cubic) alpha ramp: stays near 0 for most frames, + # then ramps up steeply at the end. This keeps the crossfade + # zone looking like the natural chunk with only the last few + # frames quickly transitioning to FLF content — far less + # visible "double exposure" than a linear ramp. + # Linear: 0.20, 0.40, 0.60, 0.80 (long ghosting) + # Cubic: 0.01, 0.06, 0.22, 0.51 (minimal ghosting) + alphas = torch.linspace(0.0, 1.0, cf + 2)[1:-1] + alphas = alphas ** 3 # cubic ease-in + for i in range(cf): + a = alphas[i].item() + flf_head[i] = (1.0 - a) * flf_saved_tail[i] + a * flf_head[i] + + print(f" FLF: crossfaded {cf} overlap frames (cubic ease-in)") + + all_frames.append(flf_head) + all_frames.append(flf_rest) + total_generated += cf + flf_rest.shape[0] + del flf_saved_tail, flf_head + + else: + # Normal chunk — append and extract reference + all_frames.append(chunk_images) + total_generated += chunk_images.shape[0] + if chunk_idx < num_chunks - 1: + reference_image = chunk_images[-1:].clone() + print(f" Extracted last frame as next reference") + + # ── Memory cleanup ── + del image_embeds, final_latent, chunk_images + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + pbar.update(1) + print(f" {chunk_label} DONE (total so far: {total_generated} frames)") + + # ══════════════════════════════════════════════════════════════ + # Concatenate all chunks + # ══════════════════════════════════════════════════════════════ + video = torch.cat(all_frames, dim=0) + + # Trim to requested length (last chunk may overshoot) + if video.shape[0] > total_frames: + video = video[:total_frames] + print(f"\n Trimmed {total_generated} → {total_frames} frames") + + final_count = video.shape[0] + chunks_done = len(all_frames) + + print(f"\n{'=' * 65}") + print(f" WanChunkedI2V COMPLETE") + print(f" Output : {final_count} frames") + print(f" Chunks used : {chunks_done}") + print(f"{'=' * 65}\n") + + return (video, final_count, chunks_done,) + + +# ═══════════════════════════════════════════════════════════════════════════ +# WanVideoChunkCalculator +# ═══════════════════════════════════════════════════════════════════════════ + +class WanVideoChunkCalculator: + """ + Simple utility node to calculate the exact total_frames needed for a + specific number of chunks, taking into account the 1-frame overlap. + """ + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "num_chunks": ("INT", { + "default": 1, "min": 1, "max": 100, "step": 1, + "tooltip": "Number of chunks you want to generate." + }), + "chunk_frames": ("INT", { + "default": 81, "min": 5, "max": 241, "step": 4, + "tooltip": "Frames per chunk. Auto-normalized to Wan's 4n+1 rule." + }), + } + } + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("total_frames",) + FUNCTION = "calculate" + CATEGORY = "VideoChunkTools" + DESCRIPTION = "Calculates the exact total_frames needed for a specific number of chunks." + + def calculate(self, num_chunks, chunk_frames): + # Normalize to 4n+1 + chunk_frames = ((chunk_frames - 1) // 4) * 4 + 1 + chunk_frames = max(chunk_frames, 5) + + if num_chunks <= 1: + total_frames = chunk_frames + else: + # First chunk is full size, subsequent chunks add (chunk_frames - 1) + total_frames = chunk_frames + (num_chunks - 1) * (chunk_frames - 1) + + return (total_frames,) + + +# ── Registration ────────────────────────────────────────────────────────── + +NODE_CLASS_MAPPINGS = { + "WanChunkedI2VSampler": WanChunkedI2VSampler, + "WanVideoChunkCalculator": WanVideoChunkCalculator, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WanChunkedI2VSampler": "Wan Chunked I2V Sampler ♾️", + "WanVideoChunkCalculator": "Wan Chunk Calculator 🧮", +}