wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-27 02:42:31 +08:00
Code Issues Actions 10 Packages Projects Releases Wiki Activity

Cap encode chunks fix.

Browse Source
This commit is contained in:
Talmaj Marinc 2026-04-14 17:00:28 +02:00
parent e962c3f846
commit 9ca7cdb17e
1 changed files with 12 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
comfy/ldm/cogvideo/vae.py
View File

@ -453,15 +453,21 @@ class AutoencoderKLCogVideoX(nn.Module):
def encode(self, x): def encode(self, x):
t = x.shape[2] t = x.shape[2]
frame_batch = self.num_sample_frames_batch_size frame_batch = self.num_sample_frames_batch_size
# ceil so remainder frames get their own chunk instead of inflating the first one remainder = t % frame_batch
num_batches = max(-(-t // frame_batch), 1)
conv_cache = None conv_cache = None
enc = [] enc = []
for i in range(num_batches):
start = i * frame_batch # Process remainder frames first so only the first chunk can have an
end = min((i + 1) * frame_batch, t) # odd temporal dimension — where Downsample3D's first-frame-special
chunk, conv_cache = self.encoder(x[:, :, start:end], conv_cache=conv_cache) # handling in temporal compression is actually correct.
if remainder > 0:
chunk, conv_cache = self.encoder(x[:, :, :remainder], conv_cache=conv_cache)
enc.append(chunk.to(x.device)) enc.append(chunk.to(x.device))
for start in range(remainder, t, frame_batch):
chunk, conv_cache = self.encoder(x[:, :, start:start + frame_batch], conv_cache=conv_cache)
enc.append(chunk.to(x.device))
enc = torch.cat(enc, dim=2) enc = torch.cat(enc, dim=2)
mean, _ = enc.chunk(2, dim=1) mean, _ = enc.chunk(2, dim=1)
return mean return mean