From 093d17c587fb628a8ddfe470203cf0baf3314524 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Jul 2026 13:05:19 -0500 Subject: [PATCH] Add SeedVR2 temporal chunk and merge nodes ## Summary Sampling a long SeedVR2 video in one pass runs out of VRAM: the DiT working set grows linearly with latent frames times pixel area, so a 100 frame clip at 4x upscale needs more memory than any consumer card has. This PR adds two workflow-level nodes that split the latent into overlapping temporal chunks and recombine the sampled chunks with a Hann crossfade. The executor's list handling runs the stock KSampler once per chunk, so the sampler itself is untouched. - **Chunk SeedVR2 Latent** splits the latent on the temporal axis. `frames_per_chunk` is in pixel frames on the 4n+1 grid, `temporal_overlap` sets how many latent frames adjacent chunks share, and `chunking_mode=auto` solves the chunk size from measured free VRAM and the latent's own dimensions. The node outputs the effective overlap so the merge is wired, not typed. - **Merge SeedVR2 Latent Chunks** recombines the sampled chunks in order, crossfading each shared region with a Hann window (flat shoulders on the outer thirds, fade across the middle third). Zero overlap is a plain concatenation, bit-identical to `torch.cat`. ## Changes - Added 2 nodes and 1 crossfade helper to `comfy_extras/nodes_seedvr.py` (+201 lines). - Added 4 chunk-law constants to `comfy/ldm/seedvr/constants.py` (+10 lines). - Added pytest unit tests in `tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py` (+66 lines): chunk geometry, 4n+1 and mode validation, overlap clamping, temporal noise mask slicing (5-D and 4-D), the auto law including batch scaling, crossfade weights and blend direction, round trip, and metadata handling. Everything outside the two new nodes is byte-identical to the base branch. The new constants are read only by the chunk node, and workflows that do not use these nodes take no new code path. ## Auto mode calibration The auto law is `max_latent_frames = (free_GiB - reserved - margin) / (0.30 * megapixels)`, calibrated on an RTX 5090 (32 GB) with the 3b fp16 model: a 17-cell resolution sweep plus temporal bisection located the activation wall and confirmed it is linear in latent frames times pixel area (the same total-voxel budget holds from 1.5:1 through 24:1 aspect ratios and under transposition). The margin is four standard deviations of the measured run-to-run spread, which costs about one latent frame of chunk size and makes an out-of-memory failure a lottery ticket rather than a coin flip. Manual mode bypasses the law entirely. On a 32 GB card, a 640x480x100 input at 4x upscale sampled as a single chunk allocates past 31 GiB and dies; auto mode picks 49 frame chunks (three chunks with overlap) and the full pipeline completes in about 240 seconds. The same law stands down on small inputs: a 320x240x100 clip runs as a single chunk because it fits. ## Example workflow Load a video, 4x upscale, auto chunking, temporal overlap 3. Expected output for a 640x480x100 input: 2560x1920, 100 frames, seams invisible at the default overlap.
API workflow JSON ```json { "14": {"inputs": {"vae_name": "ema_vae_fp16.safetensors"}, "class_type": "VAELoader"}, "17": {"inputs": {"video": ["24", 0]}, "class_type": "GetVideoComponents"}, "20": {"inputs": {"fps": ["17", 2], "images": ["30", 0], "audio": ["17", 1]}, "class_type": "CreateVideo"}, "21": {"inputs": {"tile_size": 192, "overlap": 64, "temporal_size": 64, "temporal_overlap": 8, "pixels": ["27", 0], "vae": ["14", 0]}, "class_type": "VAEEncodeTiled"}, "22": {"inputs": {"tile_size": 256, "overlap": 32, "temporal_size": 64, "temporal_overlap": 8, "samples": ["33", 0], "vae": ["14", 0]}, "class_type": "VAEDecodeTiled"}, "23": {"inputs": {"unet_name": "seedvr2_3b_fp16.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"}, "24": {"inputs": {"file": "input.mp4", "video-preview": ""}, "class_type": "LoadVideo"}, "25": {"inputs": {"filename_prefix": "video/seedvr2_upscale", "format": "auto", "codec": "auto", "video": ["20", 0]}, "class_type": "SaveVideo"}, "26": {"inputs": {"resize_type": "scale by multiplier", "resize_type.multiplier": 4, "scale_method": "bicubic", "input": ["17", 0]}, "class_type": "ResizeImageMaskNode"}, "27": {"inputs": {"resized_images": ["26", 0]}, "class_type": "SeedVR2Preprocess"}, "28": {"inputs": {"model": ["23", 0], "vae_conditioning": ["32", 0]}, "class_type": "SeedVR2Conditioning"}, "29": {"inputs": {"seed": 5770521, "steps": 1, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": ["23", 0], "positive": ["28", 0], "negative": ["28", 1], "latent_image": ["32", 0]}, "class_type": "KSampler"}, "30": {"inputs": {"color_correction_method": "lab", "images": ["22", 0], "original_resized_images": ["26", 0]}, "class_type": "SeedVR2PostProcessing"}, "32": {"inputs": {"frames_per_chunk": 21, "temporal_overlap": 3, "chunking_mode": "auto", "latent": ["21", 0]}, "class_type": "SeedVR2TemporalChunk"}, "33": {"inputs": {"temporal_overlap": ["32", 1], "latent_chunks": ["29", 0]}, "class_type": "SeedVR2TemporalMerge"} } ```
## Prior art - Reference implementation: https://github.com/ByteDance-Seed/SeedVR - Community precedent for temporal chunking with blended reassembly: https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler Chunk boundaries are a mathematical compromise: the model attends within a chunk, so different chunkings produce different outputs. The overlap crossfade hides the seam; power users can widen or zero the overlap from the workflow. --- comfy/ldm/seedvr/constants.py | 10 + comfy_extras/nodes_seedvr.py | 187 ++++++++++++++++++ .../test_seedvr2_temporal_chunk.py | 62 ++++++ 3 files changed, 259 insertions(+) create mode 100644 tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py diff --git a/comfy/ldm/seedvr/constants.py b/comfy/ldm/seedvr/constants.py index b8b300388..dca41aa54 100644 --- a/comfy/ldm/seedvr/constants.py +++ b/comfy/ldm/seedvr/constants.py @@ -1,5 +1,15 @@ """SeedVR2 constants.""" +# Temporal chunk-size law: the sampler's activation wall is linear in +# T_latent * pixel area (17-cell resolution sweep + T bisection, RTX 5090, 3b fp16): +# max_latent_frames = (free_GiB - RESERVED - K*SIGMA) / (GIB_PER_MPX_FRAME * megapixels) +# RESERVED covers model staging plus fixed CUDA/torch overhead; SIGMA is the measured +# run-to-run spread of the wall; K=4 trades ~10% smaller chunks for ~1e-5 OOM odds. +SEEDVR2_CHUNK_GIB_PER_MPX_FRAME = 0.30 +SEEDVR2_CHUNK_RESERVED_GIB = 8.5 +SEEDVR2_CHUNK_SIGMA_GIB = 0.55 +SEEDVR2_CHUNK_SIGMA_K = 4 + SEEDVR2_7B_VID_DIM = 3072 SEEDVR2_OOM_BACKOFF_DIVISOR = 2 SEEDVR2_DTYPE_BYTES_FLOOR = 4 diff --git a/comfy_extras/nodes_seedvr.py b/comfy_extras/nodes_seedvr.py index 51e412a87..bccc33b87 100644 --- a/comfy_extras/nodes_seedvr.py +++ b/comfy_extras/nodes_seedvr.py @@ -1,3 +1,5 @@ +import logging + from typing_extensions import override from comfy_api.latest import ComfyExtension, io import torch @@ -9,7 +11,12 @@ from comfy.ldm.seedvr.color_fix import ( wavelet_color_transfer, ) from comfy.ldm.seedvr.constants import ( + BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE, SEEDVR2_ADAIN_SCALE_MULTIPLIER, + SEEDVR2_CHUNK_GIB_PER_MPX_FRAME, + SEEDVR2_CHUNK_RESERVED_GIB, + SEEDVR2_CHUNK_SIGMA_GIB, + SEEDVR2_CHUNK_SIGMA_K, SEEDVR2_COLOR_MEM_HEADROOM, SEEDVR2_DTYPE_BYTES_FLOOR, SEEDVR2_LAB_SCALE_MULTIPLIER, @@ -406,6 +413,184 @@ class SeedVR2Conditioning(io.ComfyNode): return io.NodeOutput(positive, negative) +def _seedvr2_chunk_crossfade_weights(overlap, device, dtype): + """Descending previous-chunk weights across the overlap (next chunk gets ``1 - w``): a Hann fade over the middle third, flat shoulders on the outer thirds.""" + ramp = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype) + ramp = ((ramp - 1.0 / 3.0) / (1.0 / 3.0)).clamp(0.0, 1.0) + return 0.5 + 0.5 * torch.cos(torch.pi * ramp) + + +class SeedVR2TemporalChunk(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2TemporalChunk", + display_name="Chunk SeedVR2 Latent", + category="model/latent/batch", + description="Split a SeedVR2 video latent into overlapping temporal chunks small enough to sample one at a time within VRAM, wiring latent_chunks to both Apply SeedVR2 Conditioning and the sampler latent input before recombining with Merge SeedVR2 Latent Chunks.", + search_aliases=["seedvr2", "chunk", "temporal", "video upscale", "rebatch"], + inputs=[ + io.Latent.Input("latent", tooltip="The VAE-encoded SeedVR2 latent to split."), + io.Int.Input("frames_per_chunk", default=21, min=1, max=16384, step=4, + tooltip="Pixel frames per temporal chunk (4n+1: 1, 5, 9, 13, ...)."), + io.Int.Input("temporal_overlap", default=0, min=0, max=16384, + tooltip="Latent frames shared between adjacent chunks and crossfaded at merge; 0 = no overlap."), + io.Combo.Input("chunking_mode", options=["auto", "manual"], default="manual", + tooltip="manual = use frames_per_chunk exactly; auto = predict the largest chunk that fits free VRAM."), + ], + outputs=[ + io.Latent.Output(display_name="latent_chunks", is_output_list=True, + tooltip="The temporal chunks in sequence order."), + io.Int.Output(display_name="temporal_overlap", + tooltip="The effective latent-frame overlap between adjacent chunks, for Merge SeedVR2 Latent Chunks."), + ], + ) + + @classmethod + def execute(cls, latent, frames_per_chunk, temporal_overlap, chunking_mode) -> io.NodeOutput: + samples = latent["samples"] + if samples.ndim != 5: + raise ValueError( + f"SeedVR2TemporalChunk: expected a 5-D video latent (B, C, T, H, W); " + f"got shape {tuple(samples.shape)}." + ) + if samples.shape[1] != SEEDVR2_LATENT_CHANNELS: + raise ValueError( + f"SeedVR2TemporalChunk: expected {SEEDVR2_LATENT_CHANNELS} latent channels; " + f"got shape {tuple(samples.shape)}." + ) + if temporal_overlap < 0: + raise ValueError( + f"SeedVR2TemporalChunk: temporal_overlap must be >= 0; got {temporal_overlap}." + ) + if chunking_mode not in ("auto", "manual"): + raise ValueError( + f"SeedVR2TemporalChunk: chunking_mode must be 'auto' or 'manual'; " + f"got {chunking_mode!r}." + ) + t_latent = samples.shape[2] + t_pixel = 4 * (t_latent - 1) + 1 + + if chunking_mode == "auto": + free_gb = comfy.model_management.get_free_memory( + comfy.model_management.get_torch_device()) / (1024 ** 3) + mpx_per_frame = (samples.shape[0] * samples.shape[3] * samples.shape[4]) * (BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE ** 2) / 1e6 + budget_gb = free_gb - SEEDVR2_CHUNK_RESERVED_GIB - SEEDVR2_CHUNK_SIGMA_K * SEEDVR2_CHUNK_SIGMA_GIB + chunk_latent_max = max(1, int(budget_gb / (SEEDVR2_CHUNK_GIB_PER_MPX_FRAME * mpx_per_frame))) + frames_per_chunk = min(4 * (chunk_latent_max - 1) + 1, t_pixel) + logging.info( + "SeedVR2TemporalChunk auto: free=%.2fGiB, %.2fMpx -> frames_per_chunk=%d (t_pixel=%d).", + free_gb, mpx_per_frame, frames_per_chunk, t_pixel, + ) + elif frames_per_chunk < 1 or (frames_per_chunk - 1) % 4 != 0: + raise ValueError( + f"SeedVR2TemporalChunk: frames_per_chunk must be a 4n+1 pixel-frame count " + f"(1, 5, 9, 13, 17, 21, ...); got {frames_per_chunk}." + ) + + if t_pixel <= frames_per_chunk: + return io.NodeOutput([latent], 0) + + chunk_latent = (frames_per_chunk - 1) // 4 + 1 + temporal_overlap = min(temporal_overlap, chunk_latent - 1) + step = chunk_latent - temporal_overlap + + chunks = [] + for start in range(0, t_latent, step): + end = min(start + chunk_latent, t_latent) + chunk = latent.copy() + chunk["samples"] = samples[:, :, start:end].contiguous() + chunks.append(chunk) + if end >= t_latent: + break + return io.NodeOutput(chunks, temporal_overlap) + + +class SeedVR2TemporalMerge(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2TemporalMerge", + display_name="Merge SeedVR2 Latent Chunks", + category="model/latent/batch", + is_input_list=True, + description="Recombine sampled SeedVR2 temporal chunks into one latent, crossfading each overlap with a Hann window sized by the temporal_overlap wired from Chunk SeedVR2 Latent.", + search_aliases=["seedvr2", "merge", "temporal", "hann", "crossfade"], + inputs=[ + io.Latent.Input("latent_chunks", tooltip="The sampled temporal chunks in sequence order."), + io.Int.Input("temporal_overlap", default=0, min=0, max=16384, force_input=True, + tooltip="The temporal_overlap output of Chunk SeedVR2 Latent. 0 = plain concatenation."), + ], + outputs=[ + io.Latent.Output(display_name="latent", tooltip="The recombined full-length latent."), + ], + ) + + @classmethod + def execute(cls, latent_chunks, temporal_overlap) -> io.NodeOutput: + temporal_overlap = temporal_overlap[0] + if temporal_overlap < 0: + raise ValueError( + f"SeedVR2TemporalMerge: temporal_overlap must be >= 0; got {temporal_overlap}." + ) + chunks = [entry["samples"] for entry in latent_chunks] + first = chunks[0] + if first.ndim != 5: + raise ValueError( + f"SeedVR2TemporalMerge: expected 5-D video latents (B, C, T, H, W); " + f"chunk 0 has shape {tuple(first.shape)}." + ) + for i, chunk in enumerate(chunks[1:], start=1): + if chunk.shape[:2] != first.shape[:2] or chunk.shape[3:] != first.shape[3:]: + raise ValueError( + f"SeedVR2TemporalMerge: chunk {i} shape {tuple(chunk.shape)} does not " + f"match chunk 0 shape {tuple(first.shape)} outside the temporal axis." + ) + if i < len(chunks) - 1 and chunk.shape[2] != first.shape[2]: + raise ValueError( + f"SeedVR2TemporalMerge: chunk {i} has {chunk.shape[2]} latent frames but " + f"chunk 0 has {first.shape[2]}; only the final chunk may be shorter." + ) + + out = latent_chunks[0].copy() + out.pop("noise_mask", None) + + if len(chunks) == 1: + out["samples"] = first + return io.NodeOutput(out) + if temporal_overlap == 0: + out["samples"] = torch.cat(chunks, dim=2) + return io.NodeOutput(out) + + chunk_latent = first.shape[2] + step = chunk_latent - min(temporal_overlap, chunk_latent - 1) + t_total = step * (len(chunks) - 1) + chunks[-1].shape[2] + b, c, _, h, w = first.shape + merged = torch.empty((b, c, t_total, h, w), device=first.device, dtype=first.dtype) + + merged[:, :, :chunk_latent] = first + filled = chunk_latent + for i, chunk in enumerate(chunks[1:], start=1): + start = i * step + end = start + chunk.shape[2] + # Crossfade width is bounded by the previous fill frontier and by a runt + # final chunk shorter than the configured overlap. + fade = min(filled - start, chunk.shape[2]) + if fade > 0: + w_prev = _seedvr2_chunk_crossfade_weights( + fade, chunk.device, chunk.dtype).view(1, 1, fade, 1, 1) + merged[:, :, start:start + fade] = ( + merged[:, :, start:start + fade] * w_prev + chunk[:, :, :fade] * (1.0 - w_prev) + ) + merged[:, :, start + fade:end] = chunk[:, :, fade:] + else: + merged[:, :, start:end] = chunk + filled = end + + out["samples"] = merged + return io.NodeOutput(out) + + class SeedVRExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: @@ -413,6 +598,8 @@ class SeedVRExtension(ComfyExtension): SeedVR2Conditioning, SeedVR2Preprocess, SeedVR2PostProcessing, + SeedVR2TemporalChunk, + SeedVR2TemporalMerge, ] async def comfy_entrypoint() -> SeedVRExtension: diff --git a/tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py b/tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py new file mode 100644 index 000000000..8476cd1e0 --- /dev/null +++ b/tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py @@ -0,0 +1,62 @@ +"""SeedVR2 temporal chunk/merge node regression tests.""" + +import pytest +import torch + +from comfy.cli_args import args as cli_args +from comfy.ldm.seedvr.constants import SEEDVR2_LATENT_CHANNELS + +if not torch.cuda.is_available(): + cli_args.cpu = True + +import comfy.model_management # noqa: E402 +from comfy_extras.nodes_seedvr import SeedVR2TemporalChunk, SeedVR2TemporalMerge, _seedvr2_chunk_crossfade_weights # noqa: E402 + +def _latent(t_latent, h=8, w=8, b=1): + g = torch.Generator().manual_seed(7) + return {"samples": torch.randn(b, SEEDVR2_LATENT_CHANNELS, t_latent, h, w, generator=g)} + +def _split(latent, frames_per_chunk, temporal_overlap, chunking_mode="manual"): + return SeedVR2TemporalChunk.execute(latent, frames_per_chunk, temporal_overlap, chunking_mode).args + +def _merge(chunks, temporal_overlap): + return SeedVR2TemporalMerge.execute(chunks, [temporal_overlap]).args[0] + +def test_chunk_temporal_windows_and_validation(): + with pytest.raises(ValueError, match="4n\\+1"): + _split(_latent(9), 20, 0) + with pytest.raises(ValueError, match="5-D"): + _split({"samples": torch.zeros(1, SEEDVR2_LATENT_CHANNELS * 9, 8, 8)}, 21, 0) + with pytest.raises(ValueError, match="chunking_mode"): + _split(_latent(13), 21, 0, "adaptive") + latent = _latent(13) + chunks, overlap = _split(latent, 21, 2) # chunk_latent=6, step=4 -> [0:6], [4:10], [8:13] + assert overlap == 2 and [c["samples"].shape[2] for c in chunks] == [6, 6, 5] + assert all(torch.equal(c["samples"], latent["samples"][:, :, s:e]) for c, (s, e) in zip(chunks, [(0, 6), (4, 10), (8, 13)])) + assert len(_split(_latent(13), 21, 999)[0]) == 8 # overlap clamps to chunk_latent-1 -> step=1 + assert (r := _split(_latent(5), 21, 3)) and len(r[0]) == 1 and r[1] == 0 # t_pixel <= 21: passthrough + +def test_chunk_auto_mode_applies_vram_law(monkeypatch): + monkeypatch.setattr(comfy.model_management, "get_free_memory", lambda dev=None: 10.8 * (1024 ** 3)) + # budget = 10.8 - 8.5 - 4*0.55 = 0.1 GiB; 32x32 latent = 0.0655 Mpx -> chunk_latent = 5 + assert [c["samples"].shape[2] for c in _split(_latent(13, h=32, w=32), 1, 0, "auto")[0]] == [5, 5, 3] + assert _split(_latent(13, h=32, w=32, b=2), 1, 0, "auto")[0][0]["samples"].shape[2] == 2 # batch halves the chunk + +def test_merge_crossfade_and_reassembly(): + latent = _latent(13) + latent["noise_mask"] = torch.rand(1, 1, 13, 8, 8) + latent["batch_index"] = [0] + merged = _merge(_split(latent, 21, 0)[0], 0) + assert torch.equal(merged["samples"], latent["samples"]) + assert "noise_mask" not in merged and merged["batch_index"] == [0] + assert torch.allclose(_merge(_split(latent, 21, 3)[0], 3)["samples"], latent["samples"], atol=1e-6) + w = _seedvr2_chunk_crossfade_weights(3, merged["samples"].device, merged["samples"].dtype) + assert w[0] == 1.0 and w[-1] == 0.0 and torch.all(w[:-1] >= w[1:]) + ones, zeros = {"samples": torch.ones(1, SEEDVR2_LATENT_CHANNELS, 6, 8, 8)}, {"samples": torch.zeros(1, SEEDVR2_LATENT_CHANNELS, 6, 8, 8)} + fused = _merge([ones, zeros], 3)["samples"] # overlap equals w: prev fades out, next fades in + assert torch.equal(fused[:, :, 3:6], w.view(1, 1, 3, 1, 1).expand(1, SEEDVR2_LATENT_CHANNELS, 3, 8, 8)) + assert torch.equal(fused[:, :, :3], ones["samples"][:, :, :3]) and torch.equal(fused[:, :, 6:], zeros["samples"][:, :, :3]) + short = _split(latent, 21, 2)[0] + short[0]["samples"] = short[0]["samples"][:, :, :4] + with pytest.raises(ValueError, match="only the final chunk may be shorter"): + _merge(short, 2)