[feat] Add an optional MLX TAEH3 preview decoder - #1794
Merged
Conversation
Read each BF16 tensor in one bounded file operation and expand it to FP32 in MLX before constructing the decoder-layer graph. Keep embedding lookups row-scoped and retain one resident language-model layer. On the 36 GiB M4 Max, complete layer-50 conditioning falls from about 79-82 seconds to 15-16 seconds. Reversed-order trials and same-runtime comparisons preserve the exact hidden features on MLX 0.31.2 and 0.32.2. This is an uncached conditioner gain, not a full-generation claim. Cover all BF16 storage patterns, source ownership, FP16/FP32 fallback, row bounds, short reads and invalid offsets. The 23 focused cases and configured pre-commit checks pass on the tested runtimes.
Apply the gathered AdaLN scale and shift through one reusable shapeless MLX graph. Keep the normalization, table values, row indices and output dtypes unchanged. Eight-pair real-block tests reduced latency by 2.56% on MLX 0.31.2 and 1.52% on MLX 0.32.2. Order-reversed four-step trials preserved every video and audio velocity and both final latents bit-for-bit. The compiled runs were 6.7% to 11.3% faster, although active paging makes the isolated block measurement the more stable performance result. An independent seed-2027 eager replay also preserved the final video and audio latents. Cover FP32, BF16 and FP16 hidden/table combinations and strided inputs. Focused tests pass on both runtime versions, the tiny H3 parity suite passes, and configured pre-commit checks pass.
Port the released tiny H3 decoder to channels-last MLX convolutions with causal state carried across temporal chunks. Preserve normalized latent inputs and the 37-latent to 124-frame mapping. Keep the full VAE default. Add pinned safetensors loading, CLI selection, a decode benchmark, parity and dispatch tests, and documentation of reconstruction-quality limits.
mergify
Bot
added
type: feat
Contributor
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add an opt-in, fully MLX-native TAEH3 video decoder for H3 previews. The full H3 VAE remains the default. TAEH3 changes reconstruction quality, so this is separate from the exact runtime work in #1792.
The port consumes the same normalized diffusion latents, uses channels-last convolutions, and carries causal memory across bounded temporal chunks. It preserves H3's 37-latent to 124-frame mapping. No denoiser, scheduler, audio decoder, or checkpoint weights change when selecting TAEH3 alone.
--video-decode-backend taeh3selects the decoder. An immutable, SHA-256-verified safetensors checkpoint is downloaded on demand. An explicit local file supports offline use. No executable source is downloaded. A decode-only benchmark and usage/quality documentation accompany the implementation.Architecture and weights are Ollin Boer Bohan's MIT-licensed upstream TAEH3. The explicit PyTorch API is proposed separately in madebyollin/taehv#30; this MLX port does not require that PR to merge and does not claim MiniMax endorsement or newly trained weights.
Measurements
Apple M4 Max, 36 GiB, Python 3.12.13, MLX 0.32.2, FP32, five-latent chunks. Same saved production latents, 832x480, 124 frames:
One matched pair, not a repeated full-VAE baseline. Eight additional TAEH3-only trials had a 0.98 s median across seven warm repeats, range 0.96–0.99 s. The approximately 75x ratio is decoder-only; denoising is unchanged. The full-VAE trial overlapped brief CPU/pre-commit checks but no other MLX workload.
TAEH3 PSNR against the full VAE was 29.86 dB on that clip. Inspected frames show softer fine fur, fabric, and background detail. The remuxed clips have identical decoded audio hashes. This is not evidence of lossless reconstruction or a completed perceptual-quality evaluation.
A separate uncached native-resolution run with temporal
--fastplus TAEH3 completed in 205.47 s wall time. It generated 73 source frames at 832x480, then RIFE restored 124 frames with full-duration audio. Denoising took 181.14 s; TAEH3 0.64 s; RIFE 5.23 s. This is one combined-mode measurement, not native dense parity. Peak denoise allocation was 17.87 GiB, and system swap rose from 1231.19 to 2753.75 MiB. Human motion/speech acceptance remains outstanding.Validation
TAEH3_REFERENCE_DIR=/path/to/taehv python -m pytest fastvideo/tests/mlx/test_mlx_taeh3.py fastvideo/tests/mlx/test_mlx_minimax_h3*.py -q.atol=1e-5, rtol=1e-5, including the 37-frame latent fixture, multiple batches, chunk boundaries, and noncontiguous inputs.git diff --check: passed.Review notes
This branch is stacked on #1792, based on
d7d8755277ab3e9155b4816b6feecb2f3b10a984. Its first two inherited commits belong to that PR. Review the TAEH3 commit separately until #1792 merges. The existing performance PR remains untouched.Spatial fast remains an optional quality tradeoff and is not the preferred path. Its 95.75 s combined run is documented as an experiment, not a native-resolution speedup or recommended quality preset.