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 98f18e7de7c..1bfae941c15 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -65167,7 +65167,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)