Match upstream H3 VAE temporal decode

src/h3_blackwell_runtime/vae_decoder.py

@@ -229,7 +229,9 @@ class MiniMaxH3VideoVAE(nn.Module):
         shape[dim] = extent
         position = torch.arange(extent, device=b.device, dtype=b.dtype).view(shape)
         blended = a.narrow(dim, a.shape[dim] - extent, extent) * (1 - position / extent) + b.narrow(dim, 0, extent) * (position / extent)
-        return torch.cat((blended, b.narrow(dim, extent, b.shape[dim] - extent)), dim=dim)
+        if extent < b.shape[dim]:
+            return torch.cat((blended, b.narrow(dim, extent, b.shape[dim] - extent)), dim=dim)
+        return blended
 
     def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
         height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio
@@ -267,22 +269,34 @@ class MiniMaxH3VideoVAE(nn.Module):
     def _decode_temporal_pad_frames(self, z_len: int, pad_tokens: int) -> int:
         if pad_tokens <= 0:
             return 0
-        return sum(1 if (z_len - pad_tokens + index) % self.tokens_chunk_size == 0 else self.vae_ratio_t for index in range(pad_tokens))
+        intra_tail = self.clip_length % self.vae_ratio_t
+        if intra_tail == 0:
+            return pad_tokens * self.vae_ratio_t
+        z_len_before_pad = z_len - pad_tokens
+        return sum(intra_tail if (z_len_before_pad + index) % self.tokens_chunk_size == 0 else self.vae_ratio_t for index in range(pad_tokens))
 
     def _decode_temporal_frame_plan(self, z_len: int, chunks: int, pad_tokens: int) -> int:
-        total, final_overlap = 0, 0
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+        total_frames, final_overlap_frames = 0, 0
         for index in range(chunks):
-            tokens = max(0, min(index * 5 + 8, z_len) - min(index * 5, z_len))
+            token_start = index * self.tokens_chunk_size
-            frames = tokens * self.vae_ratio_t
+            token_end = token_start + self.tokens_chunk_size + self.token_overlap
-            for split in range(2):
+            clip_token_len = max(0, min(token_end, z_len) - min(token_start, z_len))
-                part = max(0, min((split + 1) * 20, frames) - split * 20 - self.frame_pre_padding)
+            clip_frame_len = clip_token_len * self.vae_ratio_t
+            for split in range(split_count):
+                frame_start = split * chunk_dec
+                frame_end = min(frame_start + chunk_dec, clip_frame_len)
+                part = max(0, frame_end - frame_start - self.frame_pre_padding)
                 if split == 0:
-                    total += part
+                    total_frames += part
                 else:
-                    final_overlap = part
+                    final_overlap_frames = part
-        return total + final_overlap - self._decode_temporal_pad_frames(z_len, pad_tokens)
+        return total_frames + final_overlap_frames - self._decode_temporal_pad_frames(z_len, pad_tokens)
 
     def decode_temporal(self, z: torch.Tensor) -> torch.Tensor:
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
         pseudo_tokens = z.shape[2] + self.token_drop
         pad_tokens = (-pseudo_tokens) % self.tokens_chunk_size
         pseudo_tokens += pad_tokens
@@ -293,18 +307,44 @@ class MiniMaxH3VideoVAE(nn.Module):
         if pad_tokens:
             z = torch.cat((z, z[:, :, -1:].expand(-1, -1, pad_tokens, -1, -1)), dim=2)
         output_frames = self._decode_temporal_frame_plan(z.shape[2], chunks, pad_tokens)
-        output, overlap = [], None
+        output = None
+        overlap = None
+        write_pos = 0
+
+        def write(part: torch.Tensor) -> None:
+            nonlocal output, write_pos
+            if part.shape[2] <= 0:
+                return
+            if output is None:
+                shape = list(part.shape)
+                shape[2] = output_frames
+                output = torch.empty(shape, dtype=part.dtype, device=part.device)
+            copy_frames = min(part.shape[2], max(0, output.shape[2] - write_pos))
+            if copy_frames > 0:
+                output[:, :, write_pos:write_pos + copy_frames].copy_(part[:, :, :copy_frames])
+                write_pos += copy_frames
+
         for index in range(chunks):
-            clip = self._adaptive_decode(z[:, :, index * 5:index * 5 + 8])
+            clip = self._adaptive_decode(z[:, :, index * self.tokens_chunk_size:index * self.tokens_chunk_size + self.tokens_chunk_size + self.token_overlap])
-            first = clip[:, :, :20, :, :][:, :, self.frame_pre_padding:]
+            for split in range(split_count):
-            tail = clip[:, :, 20:40, :, :][:, :, self.frame_pre_padding:]
+                frame_start = split * chunk_dec
-            if overlap is not None:
+                frame_end = min(frame_start + chunk_dec, clip.shape[2])
-                first = self.blend(overlap, first, self.frame_overlap, -3)
+                part = clip[:, :, frame_start:frame_end][:, :, self.frame_pre_padding:]
-            output.append(first)
+                if split == 0:
-            overlap = tail
+                    if overlap is not None:
+                        part = self.blend(overlap, part, self.frame_overlap, -3)
+                        overlap = None
+                    write(part)
+                else:
+                    overlap = part.contiguous()
+            if index == chunks - 1 and overlap is not None:
+                write(overlap)
+                overlap = None
         if overlap is not None:
-            output.append(overlap)
+            write(overlap)
-        return torch.cat(output, dim=2)[:, :, :output_frames]
+        if output is None:
+            raise RuntimeError("VAE temporal decode produced no frames")
+        return output
 
     def _adaptive_decode(self, z: torch.Tensor) -> torch.Tensor:
         return self.tiled_decode(z) if self.tiling else self._decode_pixels(z)