mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-15 02:49:18 +08:00
Try to fix some tiled VAE issues.
This commit is contained in:
parent
fd92e58eb6
commit
6c639e2a93
@ -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,
|
||||
)
|
||||
|
||||
|
||||
@ -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,
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user