mirror of
https://github.com/THUDM/CogVideo.git
synced 2026-06-01 00:38:16 +08:00
Fix patch_size_t padding calculation for frame count alignment
This commit is contained in:
parent
7a1af71545
commit
6c74057905
@ -109,10 +109,12 @@ class CogVideoXT2VLoraTrainer(Trainer):
|
|||||||
|
|
||||||
patch_size_t = self.state.transformer_config.patch_size_t
|
patch_size_t = self.state.transformer_config.patch_size_t
|
||||||
if patch_size_t is not None:
|
if patch_size_t is not None:
|
||||||
ncopy = latent.shape[2] % patch_size_t
|
remainder = latent.shape[2] % patch_size_t
|
||||||
|
ncopy = (patch_size_t - remainder) % patch_size_t
|
||||||
# Copy the first frame ncopy times to match patch_size_t
|
# Copy the first frame ncopy times to match patch_size_t
|
||||||
first_frame = latent[:, :, :1, :, :] # Get first frame [B, C, 1, H, W]
|
if ncopy > 0:
|
||||||
latent = torch.cat([first_frame.repeat(1, 1, ncopy, 1, 1), latent], dim=2)
|
first_frame = latent[:, :, :1, :, :] # Get first frame [B, C, 1, H, W]
|
||||||
|
latent = torch.cat([first_frame.repeat(1, 1, ncopy, 1, 1), latent], dim=2)
|
||||||
assert latent.shape[2] % patch_size_t == 0
|
assert latent.shape[2] % patch_size_t == 0
|
||||||
|
|
||||||
batch_size, num_channels, num_frames, height, width = latent.shape
|
batch_size, num_channels, num_frames, height, width = latent.shape
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user