mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-06 22:51:18 +08:00
093d17c587
## 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.
<details>
<summary>API workflow JSON</summary>
```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"}
}
```
</details>
## 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.
63 lines
3.5 KiB
Python
63 lines
3.5 KiB
Python
"""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)
|