feat(minimax-h3): tiny-VAE (taeh3) step previews with x-hat-0 estimates

invokeai/app/invocations/minimax_h3_denoise.py

@@ -6,7 +6,10 @@
 guidance-distilled: no negative prompt, no CFG, one forward per step.
 """
 
+from contextlib import ExitStack
+
 import torch
+from PIL import Image
 from tqdm import tqdm
 
 from invokeai.app.invocations.baseinvocation import (
@@ -26,6 +29,7 @@
     OutputField,
 )
 from invokeai.app.invocations.model import MiniMaxH3TransformerField
+from invokeai.app.services.session_processor.session_processor_common import CanceledException
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.minimax_h3.denoise import denoise
 from invokeai.backend.minimax_h3.packing import (
@@ -43,7 +47,9 @@
     build_denoise_state,
     validate_num_frames,
 )
+from invokeai.backend.minimax_h3.taehv_decoder import TAEH3_PREVIEW_MODEL_URL, TAEH3Decoder
 from invokeai.backend.minimax_h3.transformer_minimax_h3 import MiniMaxH3Transformer3DModel
+from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig
 from invokeai.backend.model_manager.taxonomy import BaseModelType
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import MiniMaxH3ConditioningInfo
@@ -68,7 +74,7 @@ class MiniMaxH3DenoiseOutput(BaseInvocationOutput):
     title="Denoise - MiniMax H3",
     tags=["latents", "video", "audio", "minimax"],
     category="latents",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class MiniMaxH3DenoiseInvocation(BaseInvocation):
@@ -141,6 +147,21 @@ def _estimate_working_memory(layout: MiniMaxH3PackedSequence) -> int:
         estimated += 2 * GB
         return estimated
 
+    @staticmethod
+    def _load_preview_decoder(context: InvocationContext) -> LoadedModelWithoutConfig | None:
+        """Fetch (a one-time ~23 MB download) and load the taeh3 preview decoder.
+
+        Previews degrade gracefully to the linear latent->RGB projection when the download is
+        unavailable (offline installs), so any failure here is a warning, never an error.
+        """
+        try:
+            return context.models.load_remote_model(TAEH3_PREVIEW_MODEL_URL, TAEH3Decoder.load_model)
+        except Exception as e:
+            context.logger.warning(
+                f"MiniMax H3 preview decoder unavailable ({e}); previews fall back to latent projection."
+            )
+            return None
+
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> MiniMaxH3DenoiseOutput:
         validate_num_frames(self.num_frames)
@@ -200,38 +221,88 @@ def invoke(self, context: InvocationContext) -> MiniMaxH3DenoiseOutput:
 
         num_condition_video_rows = state.layout.num_condition_video_rows
 
-        def step_callback(step: int, total_steps: int, video_rows: torch.Tensor) -> None:
-            # Unpack the generated rows to a 5D grid and preview the middle temporal slice.
-            latents_5d = unpatchify_video_tokens(
-                video_rows[num_condition_video_rows:],
-                num_latent_frames,
-                latent_height,
-                latent_width,
-                MINIMAX_H3_VAE_LATENT_CHANNELS,
-                MINIMAX_H3_PATCH_SIZE,
-            )
-            context.util.sd_step_callback(
-                PipelineIntermediateState(
-                    step=step,
-                    order=1,
-                    total_steps=total_steps,
-                    timestep=0,
-                    latents=latents_5d[:, :, num_latent_frames // 2],
-                ),
-                BaseModelType.MiniMaxH3,
-            )
+        preview_failed = False
+
+        def make_step_callback(preview_decoder: TAEH3Decoder | None):
+            def step_callback(step: int, total_steps: int, pred_x0_video_rows: torch.Tensor) -> None:
+                if context.util.is_canceled():
+                    raise CanceledException
+                # Unpack the generated rows' x-hat-0 estimate to a 5D grid.
+                latents_5d = unpatchify_video_tokens(
+                    pred_x0_video_rows,
+                    num_latent_frames,
+                    latent_height,
+                    latent_width,
+                    MINIMAX_H3_VAE_LATENT_CHANNELS,
+                    MINIMAX_H3_PATCH_SIZE,
+                )
+                nonlocal preview_failed
+                if preview_decoder is not None and not preview_failed:
+                    try:
+                        # A two-latent-frame window ending at the middle frame: the extra frame is
+                        # causal warmup so the decoder's temporal memory is warm for the shown frame.
+                        mid = num_latent_frames // 2
+                        window = latents_5d[:, :, max(0, mid - 1) : mid + 1]
+                        frame = preview_decoder.decode_preview_frame(window)
+                        image = Image.fromarray(frame.mul(255).round().byte().permute(1, 2, 0).cpu().numpy())
+                        context.util.signal_progress(
+                            "Denoising MiniMax H3 audio-video", step / total_steps, image, (self.width, self.height)
+                        )
+                        return
+                    except CanceledException:
+                        raise
+                    except Exception:
+                        preview_failed = True
+                        context.logger.warning(
+                            "MiniMax H3 preview decode failed; falling back to latent-projection previews.",
+                            exc_info=True,
+                        )
+                # Fallback: linear latent->RGB projection of the middle temporal slice.
+                context.util.sd_step_callback(
+                    PipelineIntermediateState(
+                        step=step,
+                        order=1,
+                        total_steps=total_steps,
+                        timestep=0,
+                        latents=latents_5d[:, :, num_latent_frames // 2],
+                    ),
+                    BaseModelType.MiniMaxH3,
+                )
+
+            return step_callback
 
         estimated_working_memory = self._estimate_working_memory(state.layout)
         transformer_info = context.models.load(self.transformer.transformer)
-        with transformer_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, transformer):
+        # The preview decoder's cache record is created after the transformer's RAM load (whose
+        # make_room could otherwise drop the unlocked 23 MB record) and locked before the
+        # transformer's VRAM lock, so the partial load accounts for it. Previews are strictly
+        # best-effort: no failure here may abort the generation.
+        preview_decoder_info = self._load_preview_decoder(context)
+        with ExitStack() as stack:
+            preview_decoder: TAEH3Decoder | None = None
+            if preview_decoder_info is not None:
+                try:
+                    preview_model = stack.enter_context(preview_decoder_info)
+                    assert isinstance(preview_model, TAEH3Decoder)
+                    preview_decoder = preview_model
+                except Exception:
+                    context.logger.warning(
+                        "Could not lock the MiniMax H3 preview decoder; previews fall back to latent projection.",
+                        exc_info=True,
+                    )
+            step_callback = make_step_callback(preview_decoder)
+
+            _, transformer = stack.enter_context(
+                transformer_info.model_on_device(working_mem_bytes=estimated_working_memory)
+            )
             assert isinstance(transformer, MiniMaxH3Transformer3DModel)
             context.util.signal_progress("Denoising MiniMax H3 audio-video")
             # steps counts sigma grid points (terminal included) -> steps-1 model evaluations.
             progress = tqdm(total=len(state.timesteps), desc=f"Denoising MiniMax H3 ({self.num_frames} frames)")
 
-            def callback_with_progress(step: int, total_steps: int, video_rows: torch.Tensor) -> None:
+            def callback_with_progress(step: int, total_steps: int, pred_x0_video_rows: torch.Tensor) -> None:
                 progress.update(1)
-                step_callback(step, total_steps, video_rows)
+                step_callback(step, total_steps, pred_x0_video_rows)
 
             try:
                 video_rows, audio_rows = denoise(

invokeai/app/util/step_callback.py

@@ -257,11 +257,39 @@
 
 WAN22_LATENT_RGB_BIAS = [0.0317, -0.0878, -0.1388]
 
-# MiniMax H3's video VAE has 24 latent channels and 16x spatial downscale. No community RGB
-# projection exists yet, so previews use a uniform channel-mean (grayscale) fallback.
-# TODO(minimax-h3): generate real factors with scripts/generate_vae_linear_approximation.py
-# against the H3 video VAE once weights are available locally.
-MINIMAX_H3_LATENT_RGB_FACTORS = [[1.0 / 24.0, 1.0 / 24.0, 1.0 / 24.0] for _ in range(24)]
+# MiniMax H3's video VAE: 24 latent channels, 16x spatial downscale. Least-squares fit of
+# NORMALIZED posterior-mean latents against 16x-downscaled RGB in [-1, 1], over real photos
+# plus synthetic gradients/patches, using the released H3 video VAE encoder (fit rms ~0.09).
+# This is the fallback path only — when the taeh3 preview decoder is available, the denoise
+# node decodes previews with it instead.
+MINIMAX_H3_LATENT_RGB_FACTORS = [
+    [-0.0127, -0.0944, -0.1146],
+    [-0.0083, 0.0638, -0.0942],
+    [0.3635, 0.4082, 0.1479],
+    [0.2079, 0.1357, -0.5101],
+    [0.0178, 0.3250, -0.3183],
+    [0.0567, 0.2060, -0.2453],
+    [0.0343, -0.0136, -0.0482],
+    [0.0079, 0.0299, -0.0814],
+    [0.0220, 0.0043, 0.0158],
+    [0.2984, 0.0988, 0.1576],
+    [-0.0066, 0.0184, 0.1134],
+    [-0.0794, -0.0416, 0.0628],
+    [0.0419, -0.0184, 0.0618],
+    [-0.0274, 0.0420, -0.0235],
+    [-0.0231, -0.0312, 0.0310],
+    [0.0089, 0.0368, -0.0387],
+    [0.0126, 0.0085, -0.0299],
+    [-0.0187, 0.0028, 0.0194],
+    [0.0264, -0.0304, 0.0089],
+    [0.0512, 0.0168, 0.0110],
+    [0.0168, -0.0357, -0.0001],
+    [0.0063, -0.0116, -0.0509],
+    [-0.0237, -0.0347, 0.0324],
+    [-0.0099, 0.0042, -0.0358],
+]
+
+MINIMAX_H3_LATENT_RGB_BIAS = [0.1189, 0.1415, -0.0034]
 
 
 def sample_to_lowres_estimated_image(
@@ -373,8 +401,9 @@ def diffusion_step_callback(
             latent_rgb_factors = WAN_LATENT_RGB_FACTORS
             latent_rgb_bias = WAN_LATENT_RGB_BIAS
     elif base_model == BaseModelType.MiniMaxH3:
-        # 24-ch H3 video VAE; grayscale channel-mean fallback until real factors exist.
+        # 24-ch H3 video VAE; factors fitted against the released encoder (see constants above).
         latent_rgb_factors = MINIMAX_H3_LATENT_RGB_FACTORS
+        latent_rgb_bias = MINIMAX_H3_LATENT_RGB_BIAS
     else:
         raise ValueError(f"Unsupported base model: {base_model}")
 

invokeai/backend/minimax_h3/denoise.py

@@ -28,8 +28,10 @@ def denoise(
         transformer: The FL2VA transformer.
         state: The prepared denoise state (rows, layout, schedules).
         prompt_embeds: The layer-50 Qwen3-VL hidden states, shape ``(1, num_text_tokens, text_dim)``.
-        step_callback: Called after every step with ``(step_index, total_steps, video_rows)`` —
-            the current video rows including conditioning rows, for previews.
+        step_callback: Called after every step with ``(step_index, total_steps, pred_x0_video_rows)``
+            — the step's *predicted-clean* (x-hat-0) estimate of the GENERATED video rows
+            (conditioning rows excluded), float32, for previews. Unlike the noisy running
+            latents, the prediction is decodable at every step.
         is_canceled: Polled once per step; a True return raises ``KeyboardInterrupt``-free
             cancellation by letting the caller's exception type propagate from the callback.
 
@@ -66,6 +68,15 @@ def denoise(
             return_dict=False,
         )
 
+        pred_x0_video_rows: torch.Tensor | None = None
+        if step_callback is not None:
+            # The scheduler's own denoised estimate (`x0 = x_t + sigma * v`, data-ward velocity),
+            # taken BEFORE the in-place Euler update below overwrites x_t.
+            sigma_video = 1.0 - t.to(torch.float32)
+            pred_x0_video_rows = latents[num_condition_video_rows:].to(torch.float32) + sigma_video * noise_pred[
+                0, num_condition_video_rows:
+            ].to(torch.float32)
+
         latents[num_condition_video_rows:] = state.scheduler.step(
             noise_pred[0, num_condition_video_rows:].float(),
             t,
@@ -80,6 +91,7 @@ def denoise(
         )[0]
 
         if step_callback is not None:
-            step_callback(i + 1, total_steps, latents)
+            assert pred_x0_video_rows is not None
+            step_callback(i + 1, total_steps, pred_x0_video_rows)
 
     return latents, audio_latents