Update h3_momentum.py

h3_momentum.py

@@ -31,6 +31,7 @@ from diffusers.modular_pipelines.minimax_h3.before_denoise import (
     MiniMaxH3SetTimestepsStep,
 )
 from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.encoders import encode_vae_condition
 from diffusers.modular_pipelines.minimax_h3.denoise import (
     MiniMaxH3DenoiseLoopWrapper,
     MiniMaxH3LoopDenoiser,
@@ -133,11 +134,17 @@ class MiniMaxH3MomentumConditionStep(ModularPipelineBlocks):
             channels = components.vae_latent_channels
             patch_size = components.patch_size
 
-            # (num_frames, 3, H, W) in [0, 1] -> (1, 3, num_frames, H, W), the video VAE's own input convention.
-            frames = block_state.given_video.to(device=device, dtype=components.vae.dtype)
-            frames = frames.permute(1, 0, 2, 3)[None]
-            posterior = components.vae.encode(frames, return_dict=False)[0]
-            momentum_latents = posterior.mode()  # (1, channels, k, latent_height, latent_width)
+            # (num_frames, 3, H, W) uint8 -> (1, 3, num_frames, H, W), `encode_vae_condition`'s own documented
+            # input convention — the same helper `MiniMaxH3KeyframeVaeEncoderStep` uses for a keyframe, so the
+            # ImageNet pixel normalization, the seeded posterior *sample* (not `.mode()`), the float16 rounding,
+            # and the final latents_mean/latents_std normalization are all exactly what the released model was
+            # actually conditioned on. Reproducing this by hand was the bug in the first version of this block —
+            # not the row-packing math, which stays as derived.
+            pixels = block_state.given_video.to(device=device).permute(1, 0, 2, 3)[None]
+            momentum_latents = encode_vae_condition(
+                components.vae, pixels, components.pixel_mean, components.pixel_std,
+                components.keyframe_encode_seed,
+            ).to(device=device)  # (1, channels, k, latent_height, latent_width)
 
             given_rows = pack_video_rows(momentum_latents, patch_size, channels).to(block_state.latents.dtype)
             num_momentum_rows = given_rows.shape[0]