daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Stream H3 VAE tiled decode canvas

Browse source
This commit is contained in:
Daniel Maddern 2026-08-13 15:53:15 +07:00
parent 830f67ead2
commit 210855c625
1 changed files with 29 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -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: