diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index 06b8823..ef6d940 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -162,8 +162,10 @@ class MiniMaxH3VideoVAE(nn.Module): super().__init__() self.vae_ratio, self.vae_ratio_t = 16, 4 self.clip_length, self.token_drop = 17, 3 - self.tokens_chunk_size, self.token_overlap = 5, 3 - self.frame_pre_padding, self.frame_overlap = 3, 9 + self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t + self.tokens_chunk_size = math.ceil(self.clip_length / self.vae_ratio_t) + self.token_overlap = (-self.token_drop) % self.tokens_chunk_size + self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0) self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64 self.post_quant_conv = nn.Conv3d(24, 24, 1, device=device) self.decoder = ViT3DDecoder(device=device)