Greg
277f12c213
- New ChainTextEmbeds node: chains up to 8 WanVideoTextEncode outputs into an ordered TEXT_EMBED_SEQUENCE for per-chunk prompting - WanChunkedI2VSampler: new optional text_embed_sequence input overrides single text_embeds when connected - Chunk N uses embed N from the sequence; last embed repeats for any extra chunks beyond the sequence length - Console output shows which embed index is used per chunk - README updated with per-chunk prompting workflow diagram
776 lines
35 KiB
Python
776 lines
35 KiB
Python
# 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. Applied to ALL chunks uniformly."
|
||
}),
|
||
"text_embed_sequence": ("TEXT_EMBED_SEQUENCE", {
|
||
"tooltip": "Per-chunk text conditioning from ChainTextEmbeds. "
|
||
"Chunk 1 gets embed_1, chunk 2 gets embed_2, etc. "
|
||
"Last embed repeats for extra chunks. "
|
||
"Overrides text_embeds when connected."
|
||
}),
|
||
"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,
|
||
text_embed_sequence=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")
|
||
if text_embed_sequence is not None:
|
||
print(f" Text Embeds : PER-CHUNK sequence ({len(text_embed_sequence)} embeds for {num_chunks} chunks)")
|
||
elif text_embeds is not None:
|
||
print(f" Text Embeds : Single (same for all chunks)")
|
||
else:
|
||
print(f" Text Embeds : None")
|
||
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} ──")
|
||
|
||
# ── Resolve per-chunk text conditioning ──
|
||
if text_embed_sequence is not None:
|
||
embed_idx = min(chunk_idx, len(text_embed_sequence) - 1)
|
||
chunk_text_embeds = text_embed_sequence[embed_idx]
|
||
if len(text_embed_sequence) > 1:
|
||
print(f" Text embed : #{embed_idx + 1} of {len(text_embed_sequence)}")
|
||
else:
|
||
chunk_text_embeds = text_embeds
|
||
|
||
# ============================================================
|
||
# 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=chunk_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=chunk_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=chunk_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 🧮",
|
||
}
|