From 15ac9b4c9cc0df0846ec6aedaf4d58b74dce188b Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 15:59:25 +0700 Subject: [PATCH] Decode final sampled latent in VAE gate --- tools/compare_vae_pixels.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tools/compare_vae_pixels.py b/tools/compare_vae_pixels.py index d27c42d..c613b93 100644 --- a/tools/compare_vae_pixels.py +++ b/tools/compare_vae_pixels.py @@ -21,14 +21,19 @@ def load_reference_frames(path: Path, frames: int) -> torch.Tensor: def main() -> None: parser = argparse.ArgumentParser() - parser.add_argument("--latent", default="/artifacts/capture/output_11.pt") + parser.add_argument("--latent", default="/artifacts/fl2va-sampler-reference/step_11.pt") parser.add_argument("--frames", type=int, default=22) parser.add_argument("--reference-frames", type=Path, default=Path("/artifacts/fl2va-sampler-reference-e679178/reference_frames")) parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/direct-vae-same-latent.mp4")) args = parser.parse_args() state = torch.load(args.latent, map_location="cuda", weights_only=False) - latent = state["video"].to("cuda") + if "denoised" in state: + video_shape = (1, 24, 7, 12, 20) + video_count = torch.tensor(video_shape).prod().item() + latent = state["denoised"].to("cuda").reshape(-1)[:video_count].reshape(video_shape) + else: + latent = state["video"].to("cuda") vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval() with torch.inference_mode(): decoded = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, : args.frames]