Skip to content

Commit 25b6d1d

Browse files
authored
wan: vae: Fix light/color change (#13101)
There was an issue where the resample split was too early and dropped one of the rolling convolutions a frame early. This is most noticable as a lighting/color change between pixel frames 5->6 (latent 2->3), or as a lighting change between the first and last frame in an FLF wan flow.
1 parent 11c15d8 commit 25b6d1d

File tree

1 file changed

+9
-9
lines changed

1 file changed

+9
-9
lines changed

comfy/ldm/wan/vae.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -376,23 +376,23 @@ def run_up(self, layer_idx, x_ref, feat_cache, feat_idx, out_chunks):
376376
return
377377

378378
layer = self.upsamples[layer_idx]
379-
if isinstance(layer, Resample) and layer.mode == 'upsample3d' and x.shape[2] > 1:
380-
for frame_idx in range(x.shape[2]):
379+
if feat_cache is not None:
380+
x = layer(x, feat_cache, feat_idx)
381+
else:
382+
x = layer(x)
383+
384+
if isinstance(layer, Resample) and layer.mode == 'upsample3d' and x.shape[2] > 2:
385+
for frame_idx in range(0, x.shape[2], 2):
381386
self.run_up(
382-
layer_idx,
383-
[x[:, :, frame_idx:frame_idx + 1, :, :]],
387+
layer_idx + 1,
388+
[x[:, :, frame_idx:frame_idx + 2, :, :]],
384389
feat_cache,
385390
feat_idx.copy(),
386391
out_chunks,
387392
)
388393
del x
389394
return
390395

391-
if feat_cache is not None:
392-
x = layer(x, feat_cache, feat_idx)
393-
else:
394-
x = layer(x)
395-
396396
next_x_ref = [x]
397397
del x
398398
self.run_up(layer_idx + 1, next_x_ref, feat_cache, feat_idx, out_chunks)

0 commit comments

Comments
 (0)