Try to fix some tiled VAE issues.

This commit is contained in:
comfyanonymous 2026-07-08 19:49:38 -04:00
parent fd92e58eb6
commit 6c639e2a93
3 changed files with 16 additions and 18 deletions

View File

@ -39,7 +39,7 @@ def _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap, temporal
temporal_size = int(temporal_size)
if temporal_size <= 0:
return 0
return None
temporal_overlap = max(0, int(temporal_overlap or 0))
temporal_overlap = min(temporal_overlap, temporal_size - 1)
@ -1535,22 +1535,21 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
return x
def decode_tiled(self, z, tile_x=32, tile_y=32, overlap=8, tile_t=None, overlap_t=None):
# SeedVR2's causal VAE owns temporal via the MemoryState cache; temporal
# slicing breaks that continuity (empirically corrupts decode), so the VAE
# tiling knobs (tile_t / overlap_t) are discarded and temporal stays whole.
# SeedVR2's causal VAE owns temporal via the MemoryState cache; external
# temporal tiling breaks that continuity, so only spatial tiling is applied.
sf = self.spatial_downsample_factor
seedvr2_tiling = {
"enable_tiling": True,
"tile_size": (tile_y * sf, tile_x * sf),
"tile_overlap": (overlap * sf, overlap * sf),
"temporal_size": 0,
"temporal_overlap": 0,
"temporal_size": None,
"temporal_overlap": None,
}
return self.decode(z, seedvr2_tiling=seedvr2_tiling)
def encode_tiled(self, x, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
# Temporal tiling knobs are discarded; the causal VAE owns temporal (slicing
# breaks MemoryState continuity), so temporal stays whole.
# External temporal tiling knobs are discarded; the causal VAE keeps its
# own internal MemoryState slicing.
if tile_y is None:
tile_y = 512
if tile_x is None:
@ -1569,8 +1568,8 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
self,
tile_size=(tile_y, tile_x),
tile_overlap=(overlap_y, overlap_x),
temporal_size=0,
temporal_overlap=0,
temporal_size=None,
temporal_overlap=None,
encode=True,
)

View File

@ -304,9 +304,8 @@ def test_vaedecode_tiled_spatial_applies_temporal_discarded(monkeypatch):
temporal_overlap=4,
)
# Spatial inputs flow through; temporal inputs are discarded — SeedVR2 owns
# temporal via the MemoryState causal cache, so VAEDecodeTiled's temporal
# knobs are no-ops at the wrapper.
# Spatial inputs flow through; temporal inputs are discarded as public tiling
# knobs, but SeedVR2's internal MemoryState causal slicing is left intact.
assert vae.first_stage_model.calls == [
{
"shape": (1, _LATENT_CHANNELS, 2, 4, 5),
@ -314,8 +313,8 @@ def test_vaedecode_tiled_spatial_applies_temporal_discarded(monkeypatch):
"enable_tiling": True,
"tile_size": (512, 512),
"tile_overlap": (64, 64),
"temporal_size": 0,
"temporal_overlap": 0,
"temporal_size": None,
"temporal_overlap": None,
},
}
]

View File

@ -19,7 +19,7 @@ from comfy.ldm.seedvr.vae import MemoryState, tiled_vae # noqa: E402
_LATENT_CHANNELS = seedvr_vae_mod.SEEDVR2_LATENT_CHANNELS
def test_runtime_decode_zero_temporal_size_disables_slicing_for_call():
def test_runtime_decode_zero_temporal_size_preserves_model_slicing():
class StubVAEModel(torch.nn.Module):
def __init__(self):
super().__init__()
@ -54,8 +54,8 @@ def test_runtime_decode_zero_temporal_size_disables_slicing_for_call():
encode=False,
)
assert vae.decode_min_sizes == [5]
assert vae.memory_states == [MemoryState.DISABLED]
assert vae.decode_min_sizes == [2]
assert vae.memory_states == [MemoryState.INITIALIZING, MemoryState.ACTIVE]
assert vae.slicing_latent_min_size == 2