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

Cap encode chunks at the configured frame batch size.

Browse Source
This commit is contained in:
Talmaj Marinc 2026-04-14 15:05:27 +02:00
parent 3e961f9960
commit e962c3f846
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -453,13 +453,13 @@ 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
num_batches = max(t // frame_batch, 1) # ceil so remainder frames get their own chunk instead of inflating the first one
num_batches = max(-(-t // frame_batch), 1)
conv_cache = None conv_cache = None
enc = [] enc = []
for i in range(num_batches): for i in range(num_batches):
remaining = t % frame_batch start = i * frame_batch
start = frame_batch * i + (0 if i == 0 else remaining) end = min((i + 1) * frame_batch, t)
end = frame_batch * (i + 1) + remaining
chunk, conv_cache = self.encoder(x[:, :, start:end], conv_cache=conv_cache) chunk, conv_cache = self.encoder(x[:, :, start:end], conv_cache=conv_cache)
enc.append(chunk.to(x.device)) enc.append(chunk.to(x.device))
enc = torch.cat(enc, dim=2) enc = torch.cat(enc, dim=2)