diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index 67763e3..36d01b7 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/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)) - 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:] - if overlap is not None: - first = self.blend(overlap, first, self.frame_overlap, -3) - output.append(first) - overlap = tail + 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: + 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)