From 210855c625f4662b6bd6617b6e29fb5bee7e3d16 Mon Sep 17 00:00:00 2001
From: Daniel Maddern <daniel.maddern@illuvium.io>
Date: Thu, 13 Aug 2026 15:53:15 +0700
Subject: [PATCH] Stream H3 VAE tiled decode canvas

---
 src/h3_blackwell_runtime/vae_decoder.py | 50 ++++++++++++++-----------
 1 file changed, 29 insertions(+), 21 deletions(-)

diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py
index 88d1790..781dc57 100644
--- a/src/h3_blackwell_runtime/vae_decoder.py
+++ b/src/h3_blackwell_runtime/vae_decoder.py
@@ -202,29 +202,37 @@ class MiniMaxH3VideoVAE(nn.Module):
         return torch.cat((blended, b.narrow(dim, extent, b.shape[dim] - extent)), dim=dim)
 
     def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
-        y_starts, y_lengths, y_overlaps = self.split_tiles(z.shape[-2] * self.vae_ratio)
-        x_starts, x_lengths, x_overlaps = self.split_tiles(z.shape[-1] * self.vae_ratio)
-        rows: list[list[torch.Tensor]] = []
-        for y, height in zip(y_starts, y_lengths):
-            row = []
-            for x, width in zip(x_starts, x_lengths):
-                row.append(self._decode_pixels(z[..., y // 16:(y + height) // 16, x // 16:(x + width) // 16]))
-            rows.append(row)
-        result_rows = []
-        for row_index, row in enumerate(rows):
-            result_row = []
-            for index, tile in enumerate(row):
+        height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio
+        y_starts, y_lengths, y_overlaps = self.split_tiles(height)
+        x_starts, x_lengths, x_overlaps = self.split_tiles(width)
+        canvas = None
+        row_tails = []
+        output_y = 0
+        for row_index, (y, tile_height) in enumerate(zip(y_starts, y_lengths)):
+            next_row_tails = []
+            left_tail = None
+            output_x = 0
+            for column_index, (x, tile_width) in enumerate(zip(x_starts, x_lengths)):
+                tile = self._decode_pixels(z[..., y // self.vae_ratio:(y + tile_height) // self.vae_ratio, x // self.vae_ratio:(x + tile_width) // self.vae_ratio])
+                if row_index < len(y_starts) - 1:
+                    next_row_tails.append(tile[..., -y_overlaps[row_index]:, :].clone())
+                next_left_tail = tile[..., :, -x_overlaps[column_index]:].clone() if column_index < len(x_starts) - 1 else None
                 if row_index:
-                    tile = self.blend(rows[row_index - 1][index], tile, y_overlaps[row_index - 1], -2)
-                if result_row:
-                    tile = self.blend(row[index - 1], tile, x_overlaps[index - 1], -1)
-                if row_index < len(rows) - 1:
+                    tile = self.blend(row_tails[column_index], tile, y_overlaps[row_index - 1], -2)
+                if column_index:
+                    tile = self.blend(left_tail, tile, x_overlaps[column_index - 1], -1)
+                left_tail = next_left_tail
+                if row_index < len(y_starts) - 1:
                     tile = tile[..., :-y_overlaps[row_index], :]
-                if index < len(row) - 1:
-                    tile = tile[..., :, :-x_overlaps[index]]
-                result_row.append(tile)
-            result_rows.append(torch.cat(result_row, dim=-1))
-        return torch.cat(result_rows, dim=-2)
+                if column_index < len(x_starts) - 1:
+                    tile = tile[..., :, :-x_overlaps[column_index]]
+                if canvas is None:
+                    canvas = torch.empty(*tile.shape[:-2], height, width, dtype=tile.dtype, device=tile.device)
+                canvas[..., output_y:output_y + tile.shape[-2], output_x:output_x + tile.shape[-1]].copy_(tile)
+                output_x += tile.shape[-1]
+            row_tails = next_row_tails
+            output_y += tile.shape[-2]
+        return canvas
 
     def _decode_temporal_pad_frames(self, z_len: int, pad_tokens: int) -> int:
         if pad_tokens <= 0: