ComfyUI/comfy/ldm/seedvr/constants.py
John Pollock 093d17c587 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.

<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.
2026-07-04 13:05:19 -05:00

49 lines
2.6 KiB
Python

"""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
SEEDVR2_7B_MLP_CHUNK = 8192
SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS = 4096 # partial-RoPE application token-chunk.
SEEDVR2_LATENT_CHANNELS = 16
SEEDVR2_COLOR_MEM_HEADROOM = 0.75
SEEDVR2_LAB_SCALE_MULTIPLIER = 13
SEEDVR2_WAVELET_SCALE_MULTIPLIER = 10 # per-frame byte multiplier, wavelet path.
SEEDVR2_ADAIN_SCALE_MULTIPLIER = 6
BYTEDANCE_VAE_SCALING_FACTOR = 0.9152 # configs_3b/main.yaml:57.
BYTEDANCE_VAE_SHIFTING_FACTOR = 0.0
BYTEDANCE_VAE_CONV_MEM_GIB = 0.5
BYTEDANCE_VAE_NORM_MEM_GIB = 0.5
BYTEDANCE_LOGVAR_CLAMP_MIN = -30.0 # video_vae_v3/modules/types.py:28.
BYTEDANCE_LOGVAR_CLAMP_MAX = 20.0 # video_vae_v3/modules/types.py:28.
BYTEDANCE_GN_CHUNKS_FP16 = 4 # causal_inflation_lib.py:351 (GroupNorm chunk count, fp16).
BYTEDANCE_GN_CHUNKS_FP32 = 2 # causal_inflation_lib.py:351 (GroupNorm chunk count, fp32).
BYTEDANCE_BLOCK_OUT_CHANNELS = (128, 256, 512, 512) # s8_c16_t4_inflation_sd3.yaml:7-11.
BYTEDANCE_SLICING_SAMPLE_MIN = 4 # s8_c16_t4_inflation_sd3.yaml:22 (slicing_sample_min_size).
BYTEDANCE_VAE_TEMPORAL_DOWNSAMPLE = 4 # infer.py:230 (temporal_downsample_factor); the 4n+1 factor.
BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE = 8 # infer.py:231 (spatial_downsample_factor).
BYTEDANCE_720P_REF_AREA = 45 * 80 # dit_v2/window.py:32 (720p reference area for window scaling).
BYTEDANCE_MAX_TEMPORAL_WINDOW = 30 # dit_v2/window.py:35 (max temporal window frames).
BYTEDANCE_ROPE_MAX_FREQ = 256 # dit_v2/rope.py:31 (pixel-RoPE max frequency).
BYTEDANCE_SINUSOIDAL_DIM = 256 # dit_3b/nadit.py:120 (timestep sinusoidal embed dim).
ROPE_THETA = 10000 # RoPE base; Su et al., "RoFormer", arXiv:2104.09864.
CIELAB_DELTA = 6.0 / 29.0 # CIE 15 (delta).
CIELAB_KAPPA = (29.0 / 3.0) ** 3 # CIE 15 (kappa).
D65_WHITE_X = 0.95047 # CIE D65 standard illuminant Xn (Yn = 1).
D65_WHITE_Z = 1.08883 # CIE D65 standard illuminant Zn.
WAVELET_DECOMP_LEVELS = 5 # wavelet color-fix decomposition depth (GIMP/Krita; StableSR).