From 4697e8720ab66057c84a2782fe03ba471013e0b9 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 5 Aug 2026 19:35:18 -0400 Subject: [PATCH 1/2] feat(minimax-h3): tiny-VAE (taeh3) step previews with x-hat-0 estimates The step preview previously projected the RAW noisy latents x_t through a uniform grayscale channel mean, so nothing recognizable appeared until the last couple of steps. Three changes: - The denoise loop now hands the step callback the scheduler's own predicted-clean estimate (x0 = x_t + sigma*v, data-ward velocity) of the generated video rows, which is decodable at every step. - The denoise invocation decodes that estimate with madebyollin's taeh3 tiny decoder (MIT, ~23 MB one-time download via the remote-model cache; vendored decoder-only, bit-exact against the reference implementation). A full-size 768x1344 preview frame costs ~16 ms per 50 s step. Verified against the real H3 video VAE that taeh3 consumes the NORMALIZED latent space (the denoise loop's space): real-VAE encode -> taeh3 decode reconstructs at ~21 dB, unnormalized latents produce garbage. - The no-download fallback path gets real latent->RGB factors + bias, least-squares fitted against the released H3 VAE encoder (rms ~0.09), replacing the grayscale placeholder. Co-Authored-By: Claude Fable 5 --- .../app/invocations/minimax_h3_denoise.py | 119 +++++++++++--- invokeai/app/util/step_callback.py | 41 ++++- invokeai/backend/minimax_h3/denoise.py | 18 ++- invokeai/backend/minimax_h3/taehv_decoder.py | 152 ++++++++++++++++++ invokeai/frontend/web/openapi.json | 2 +- tests/backend/minimax_h3/test_denoise.py | 20 +++ .../backend/minimax_h3/test_taehv_decoder.py | 86 ++++++++++ 7 files changed, 404 insertions(+), 34 deletions(-) create mode 100644 invokeai/backend/minimax_h3/taehv_decoder.py create mode 100644 tests/backend/minimax_h3/test_taehv_decoder.py diff --git a/invokeai/app/invocations/minimax_h3_denoise.py b/invokeai/app/invocations/minimax_h3_denoise.py index 2c63fd27d4f..f2389e05017 100644 --- a/invokeai/app/invocations/minimax_h3_denoise.py +++ b/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( diff --git a/invokeai/app/util/step_callback.py b/invokeai/app/util/step_callback.py index c3a5d578b00..f70dfda9cb0 100644 --- a/invokeai/app/util/step_callback.py +++ b/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}") diff --git a/invokeai/backend/minimax_h3/denoise.py b/invokeai/backend/minimax_h3/denoise.py index 717f39bdd53..13d83f86fdc 100644 --- a/invokeai/backend/minimax_h3/denoise.py +++ b/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 diff --git a/invokeai/backend/minimax_h3/taehv_decoder.py b/invokeai/backend/minimax_h3/taehv_decoder.py new file mode 100644 index 00000000000..5a08aa657e3 --- /dev/null +++ b/invokeai/backend/minimax_h3/taehv_decoder.py @@ -0,0 +1,152 @@ +"""Tiny AutoEncoder decoder for MiniMax H3 latents ("taeh3"), used for denoising previews. + +Adapted from madebyollin's TAEHV (MIT license), decoder half only, pinned to the same commit +as the published H3 weights: +https://github.com/madebyollin/taehv/blob/62f7591f59dfbb4c3c02b7a621d180a9eeaba26c/taehv.py + +The decoder consumes latents in the NORMALIZED space — ``(z - latents_mean) / latents_std``, +i.e. exactly the space the denoise loop operates in — and emits RGB frames in ``[0, 1]``. +Verified against the real H3 video VAE: a real-VAE encode round-tripped through this decoder +reconstructs at ~21 dB PSNR, while feeding *unnormalized* latents produces garbage, so no +mean/std conversion belongs in the preview path. + +Scale factors mirror the full VAE: 16x spatial (8x upsampling + 2x pixel shuffle), 4x +temporal (TGrow strides 1/2/2), 24 latent channels. The ``decoder.*`` ``nn.Sequential`` +indices are a state-dict contract with the released ``taeh3.safetensors`` — do not reorder. +""" + +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch import nn + +# Weights are pinned to the commit whose taehv.py this module mirrors. The GitHub "raw" +# redirect resolves to the LFS object; the download cache keys on this URL. +TAEH3_PREVIEW_MODEL_URL = ( + "https://github.com/madebyollin/taehv/raw/62f7591f59dfbb4c3c02b7a621d180a9eeaba26c/safetensors/taeh3.safetensors" +) + +TAEH3_LATENT_CHANNELS = 24 +TAEH3_TEMPORAL_UPSCALE = 4 +_PIXEL_SHUFFLE_PATCH = 2 + + +class _Clamp(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.tanh(x / 3) * 3 + + +class _MemBlock(nn.Module): + """Residual block whose conv also sees the previous frame's input (the temporal memory).""" + + def __init__(self, n_in: int, n_out: int): + super().__init__() + self.conv = nn.Sequential( + nn.Conv2d(n_in * 2, n_out, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(n_out, n_out, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(n_out, n_out, 3, padding=1), + ) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor, past: torch.Tensor) -> torch.Tensor: + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + + +class _TGrow(nn.Module): + """Temporal upsampling: a 1x1 conv to ``stride`` channel groups, reshaped into frames.""" + + def __init__(self, n_f: int, stride: int): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f * stride, 1, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _nt, c, h, w = x.shape + return self.conv(x).reshape(-1, c, h, w) + + +class TAEH3Decoder(nn.Module): + """The taeh3 decoder: ``[N, 24, T, h, w]`` normalized latents -> ``[N, 4T, 3, 16h, 16w]`` RGB.""" + + def __init__(self) -> None: + super().__init__() + n_f = (256, 128, 64, 64) + self.decoder = nn.Sequential( + _Clamp(), + nn.Conv2d(TAEH3_LATENT_CHANNELS, n_f[0], 3, padding=1), + nn.ReLU(inplace=True), + _MemBlock(n_f[0], n_f[0]), + _MemBlock(n_f[0], n_f[0]), + _MemBlock(n_f[0], n_f[0]), + nn.Upsample(scale_factor=2), + _TGrow(n_f[0], 1), + nn.Conv2d(n_f[0], n_f[1], 3, padding=1, bias=False), + _MemBlock(n_f[1], n_f[1]), + _MemBlock(n_f[1], n_f[1]), + _MemBlock(n_f[1], n_f[1]), + nn.Upsample(scale_factor=2), + _TGrow(n_f[1], 2), + nn.Conv2d(n_f[1], n_f[2], 3, padding=1, bias=False), + _MemBlock(n_f[2], n_f[2]), + _MemBlock(n_f[2], n_f[2]), + _MemBlock(n_f[2], n_f[2]), + nn.Upsample(scale_factor=2), + _TGrow(n_f[2], 2), + nn.Conv2d(n_f[2], n_f[3], 3, padding=1, bias=False), + nn.ReLU(inplace=True), + nn.Conv2d(n_f[3], 3 * _PIXEL_SHUFFLE_PATCH**2, 3, padding=1), + ) + + @classmethod + def load_model(cls, path: Path) -> "TAEH3Decoder": + """Loader for ``context.models.load_remote_model``: reads ``taeh3.safetensors``. + + The file also carries the (unused) encoder; only ``decoder.*`` keys are consumed, + strictly. Weights stay in their stored float16. + """ + from safetensors.torch import load_file + + state_dict = {k: v for k, v in load_file(path).items() if k.startswith("decoder.")} + model = cls() + model.to(torch.float16) + model.load_state_dict(state_dict, strict=True) + model.eval() + return model + + @torch.no_grad() + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """Decode ``[N, C, T, h, w]`` normalized latents to ``[N, 4T, 3, 16h, 16w]`` RGB in [0, 1]. + + Straight decoder pass with parallel temporal-memory handling — none of the full VAE's + 17-frame chunk alignment, so the leading ``4T - (4T - 3)``-ish warmup frames decode + against empty memory: for previews, read the LAST frame, not the first. + """ + if latents.ndim != 5 or latents.shape[1] != TAEH3_LATENT_CHANNELS: + raise ValueError(f"Expected [N, {TAEH3_LATENT_CHANNELS}, T, h, w] latents, got {list(latents.shape)}.") + weight_dtype = self.decoder[1].weight.dtype + x = latents.permute(0, 2, 1, 3, 4).to(weight_dtype) # NTCHW + n = x.shape[0] + x = x.flatten(0, 1) + for block in self.decoder: + if isinstance(block, _MemBlock): + nt, c, h, w = x.shape + t = nt // n + # Each frame's memory is the previous frame's block input; frame 0 sees zeros. + past = F.pad(x.view(n, t, c, h, w), (0, 0, 0, 0, 0, 0, 1, 0))[:, :t].reshape(x.shape) + x = block(x, past) + else: + x = block(x) + frames = F.pixel_shuffle(x, _PIXEL_SHUFFLE_PATCH).clamp_(0, 1) + return frames.view(n, -1, *frames.shape[1:]) + + def decode_preview_frame(self, latents: torch.Tensor) -> torch.Tensor: + """Decode a small latent window and return the best single frame, ``[3, H, W]`` in [0, 1]. + + The last frame of the window has the most temporal-memory context (the earlier frames + are causal warmup), so it is the one worth showing. + """ + return self.decode(latents)[0, -1] diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 4e9bb1f2b67..8bc159ff6e1 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -65851,7 +65851,7 @@ "tags": ["latents", "video", "audio", "minimax"], "title": "Denoise - MiniMax H3", "type": "object", - "version": "1.0.0", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/MiniMaxH3DenoiseOutput" } diff --git a/tests/backend/minimax_h3/test_denoise.py b/tests/backend/minimax_h3/test_denoise.py index 67469f9ba33..62a2778dd27 100644 --- a/tests/backend/minimax_h3/test_denoise.py +++ b/tests/backend/minimax_h3/test_denoise.py @@ -91,6 +91,26 @@ def test_step_callback_called_per_step(tiny_transformer): assert calls == [(i + 1, NUM_EVALS) for i in range(NUM_EVALS)] +def test_step_callback_receives_pred_x0_of_generated_rows(tiny_transformer): + """The callback gets the x-hat-0 estimate of the GENERATED video rows: condition rows are + excluded, dtype is float32, and at the final step (sigma_next = 0) the Euler update reduces + to x_next = x0, so the last callback payload must equal the returned generated rows.""" + state = _state(with_keyframe=True) + num_condition_rows = state.layout.num_condition_video_rows + num_generated_rows = state.video_rows.shape[0] - num_condition_rows + seen: list[torch.Tensor] = [] + prompt_embeds = torch.randn(1, 3, TINY_CONFIG["text_dim"], generator=torch.Generator().manual_seed(3)) + video_rows, _ = denoise( + tiny_transformer, + state, + prompt_embeds, + step_callback=lambda step, total, rows: seen.append(rows.clone()), + ) + assert all(rows.shape == (num_generated_rows, video_rows.shape[1]) for rows in seen) + assert all(rows.dtype == torch.float32 for rows in seen) + torch.testing.assert_close(seen[-1], video_rows[num_condition_rows:].to(torch.float32)) + + def test_cancellation_raises(tiny_transformer): state = _state() prompt_embeds = torch.randn(1, 3, TINY_CONFIG["text_dim"]) diff --git a/tests/backend/minimax_h3/test_taehv_decoder.py b/tests/backend/minimax_h3/test_taehv_decoder.py new file mode 100644 index 00000000000..0008b4b5864 --- /dev/null +++ b/tests/backend/minimax_h3/test_taehv_decoder.py @@ -0,0 +1,86 @@ +"""Tests for the vendored taeh3 preview decoder (structure, IO contract, loader).""" + +import pytest +import torch +from safetensors.torch import save_file + +from invokeai.backend.minimax_h3.taehv_decoder import TAEH3_LATENT_CHANNELS, TAEH3Decoder + + +@pytest.fixture(scope="module") +def decoder() -> TAEH3Decoder: + torch.manual_seed(0) + model = TAEH3Decoder() + model.eval() + return model + + +def test_decode_shapes_and_range(decoder): + """[N, 24, T, h, w] -> [N, 4T, 3, 16h, 16w], clamped to [0, 1].""" + latents = torch.randn(1, TAEH3_LATENT_CHANNELS, 2, 4, 6, generator=torch.Generator().manual_seed(1)) + frames = decoder.decode(latents) + assert frames.shape == (1, 8, 3, 64, 96) + assert frames.min() >= 0.0 and frames.max() <= 1.0 + + +def test_decode_preview_frame_is_last(decoder): + latents = torch.randn(1, TAEH3_LATENT_CHANNELS, 2, 4, 4, generator=torch.Generator().manual_seed(2)) + frame = decoder.decode_preview_frame(latents) + assert frame.shape == (3, 64, 64) + assert torch.equal(frame, decoder.decode(latents)[0, -1]) + + +def test_decode_rejects_bad_shapes(decoder): + with pytest.raises(ValueError, match="Expected"): + decoder.decode(torch.randn(1, 16, 2, 4, 4)) + with pytest.raises(ValueError, match="Expected"): + decoder.decode(torch.randn(TAEH3_LATENT_CHANNELS, 2, 4, 4)) + + +def test_state_dict_matches_released_checkpoint_layout(decoder): + """Pin the nn.Sequential indices/shapes against the released taeh3.safetensors layout. + + These keys and shapes were read from the file at the pinned commit; if this test fails, the + vendored structure has drifted and the real checkpoint will no longer strict-load. + """ + sd = decoder.state_dict() + expected = { + "decoder.1.weight": (256, 24, 3, 3), + "decoder.3.conv.0.weight": (256, 512, 3, 3), + "decoder.7.conv.weight": (256, 256, 1, 1), # TGrow stride 1 + "decoder.8.weight": (128, 256, 3, 3), + "decoder.13.conv.weight": (256, 128, 1, 1), # TGrow stride 2 + "decoder.14.weight": (64, 128, 3, 3), + "decoder.19.conv.weight": (128, 64, 1, 1), # TGrow stride 2 + "decoder.20.weight": (64, 64, 3, 3), + "decoder.22.weight": (12, 64, 3, 3), # 3 RGB x 2x2 pixel-shuffle patch + } + for key, shape in expected.items(): + assert key in sd, f"missing {key}" + assert tuple(sd[key].shape) == shape, f"{key}: {tuple(sd[key].shape)} != {shape}" + # 4x temporal upscale total: product of TGrow strides (1, 2, 2). + assert sd["decoder.13.conv.weight"].shape[0] // sd["decoder.13.conv.weight"].shape[1] == 2 + assert sd["decoder.19.conv.weight"].shape[0] // sd["decoder.19.conv.weight"].shape[1] == 2 + + +def test_load_model_ignores_encoder_keys(tmp_path, decoder): + """The released file bundles an encoder; load_model must strict-load only decoder.* keys.""" + sd = {k: v.to(torch.float16) for k, v in decoder.state_dict().items()} + sd["encoder.1.weight"] = torch.zeros(64, 12, 3, 3, dtype=torch.float16) + path = tmp_path / "taeh3.safetensors" + save_file(sd, str(path)) + + loaded = TAEH3Decoder.load_model(path) + assert not loaded.training + assert loaded.decoder[1].weight.dtype == torch.float16 + frames = loaded.decode(torch.randn(1, TAEH3_LATENT_CHANNELS, 1, 4, 4)) + assert frames.shape == (1, 4, 3, 64, 64) + + +def test_load_model_rejects_missing_keys(tmp_path, decoder): + sd = {k: v.to(torch.float16) for k, v in decoder.state_dict().items()} + sd.pop("decoder.22.weight") + path = tmp_path / "incomplete.safetensors" + save_file(sd, str(path)) + with pytest.raises(RuntimeError, match="Missing"): + TAEH3Decoder.load_model(path) From bb536047e2389b85388eee07ec1e910156d92020 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 16 Aug 2026 09:41:10 -0400 Subject: [PATCH 2/2] fix(minimax-h3): bump the bundled workflows' denoise node to 1.1.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_bundled_workflow_nodes_match_invocation_registry asserts every bundled workflow embeds the current version of each node it uses. This branch bumps minimax_h3_denoise to 1.1.0 for the taeh3 preview inputs, while the two H3 default workflows — added by the workflows PR, which merged first — still embedded 1.0.0. Neither change is wrong on its own; the mismatch only became reachable once both were on the same tree, which is why it surfaced in CI after the sync rather than on either branch alone. Only the minimax_h3_denoise node's version field changes in each file (the other embedded node versions are already current). Verified with the full backend suite this time rather than a subset: 4597 passed, 120 skipped, 5 deselected, 8 xfailed. Co-Authored-By: Claude Opus 5 (1M context) --- .../default_workflows/First Frame to Video - MiniMax H3.json | 2 +- .../default_workflows/Text to Video - MiniMax H3.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/invokeai/app/services/workflow_records/default_workflows/First Frame to Video - MiniMax H3.json b/invokeai/app/services/workflow_records/default_workflows/First Frame to Video - MiniMax H3.json index b7ae0ebd5e9..dccc8e4aef2 100644 --- a/invokeai/app/services/workflow_records/default_workflows/First Frame to Video - MiniMax H3.json +++ b/invokeai/app/services/workflow_records/default_workflows/First Frame to Video - MiniMax H3.json @@ -234,7 +234,7 @@ "type": "invocation", "data": { "id": "9d3b9faf-405b-7e6f-bf4d-5c6a7f8e9d04", - "version": "1.0.0", + "version": "1.1.0", "nodePack": "invokeai", "label": "", "notes": "", diff --git a/invokeai/app/services/workflow_records/default_workflows/Text to Video - MiniMax H3.json b/invokeai/app/services/workflow_records/default_workflows/Text to Video - MiniMax H3.json index d01f50a86aa..01df4fa8567 100644 --- a/invokeai/app/services/workflow_records/default_workflows/Text to Video - MiniMax H3.json +++ b/invokeai/app/services/workflow_records/default_workflows/Text to Video - MiniMax H3.json @@ -156,7 +156,7 @@ "type": "invocation", "data": { "id": "9d3b9faf-405b-7e6f-bf4d-5c6a7f8e9d04", - "version": "1.0.0", + "version": "1.1.0", "nodePack": "invokeai", "label": "", "notes": "",