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

Match upstream H3 VAE temporal decode

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:29:31 +07:00
parent d8fd04bce4
commit d45773053e
1 changed files with 60 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

76
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -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)
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))
frames = tokens * self.vae_ratio_t
for split in range(2):
part = max(0, min((split + 1) * 20, frames) - split * 20 - self.frame_pre_padding)
token_start = index * self.tokens_chunk_size
token_end = token_start + self.tokens_chunk_size + self.token_overlap
clip_token_len = max(0, min(token_end, z_len) - min(token_start, z_len))
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
return total + final_overlap - self._decode_temporal_pad_frames(z_len, pad_tokens)
final_overlap_frames = part
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])
first = clip[:, :, :20, :, :][:, :, self.frame_pre_padding:]
tail = clip[:, :, 20:40, :, :][:, :, self.frame_pre_padding:]
clip = self._adaptive_decode(z[:, :, index * self.tokens_chunk_size:index * self.tokens_chunk_size + self.tokens_chunk_size + self.token_overlap])
for split in range(split_count):
frame_start = split * chunk_dec
frame_end = min(frame_start + chunk_dec, clip.shape[2])
part = clip[:, :, frame_start:frame_end][:, :, self.frame_pre_padding:]
if split == 0:
if overlap is not None:
first = self.blend(overlap, first, self.frame_overlap, -3)
output.append(first)
overlap = tail
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)
return torch.cat(output, dim=2)[:, :, :output_frames]
write(overlap)
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)