Problem
After adding fallback from missing vae/encoder and vae/decoder subdirectories to the unified Hugging Face vae/ directory, image-to-video generation gets past VAE file discovery but fails during VAE encoder forward:
✓ Transformer loaded
Auto-detected STG blocks: [29] (model=2)
Traceback (most recent call last):
File "mlx_video/models/ltx_2/generate.py", line 2107, in generate_video
stage1_image_latent = vae_encoder(...)
File "mlx_video/models/ltx_2/video_vae/video_vae.py", line 347, in __call__
sample = down_block(sample, causal=True)
File "mlx_video/models/ltx_2/video_vae/sampling.py", line 103, in __call__
return x_conv + x_in
ValueError: [broadcast_shapes] Shapes (1,4,1,32,32) and (1,6,1,32,32) cannot be broadcast.
Root Cause
The unified Lightricks/LTX-2 VAE safetensors file uses unprefixed Diffusers component keys:
encoder.conv_in.conv.weight
encoder.down_blocks.0.resnets.0.conv1.conv.weight
decoder.conv_in.conv.weight
decoder.up_blocks.0.resnets.0.conv1.conv.weight
Current sanitizers only accept prefixed keys:
vae.encoder.*
vae.decoder.*
A local key comparison showed:
VideoEncoder.sanitize(unified_vae_weights) -> 0 sanitized keys
model_keys -> 94 missing keys
So the unified VAE file is found, but the encoder/decoder weights are not loaded. The runtime then executes an effectively uninitialized/default VAE path and fails with a shape mismatch.
Expected Behavior
VideoEncoder.sanitize() and VideoDecoder.sanitize() should accept both layouts:
vae.encoder.* / vae.decoder.*
encoder.* / decoder.*
and normalize them into the local MLX module key names before loading.
Local Evidence
Failed log:
outputs/logs/horse-avatar-ltx2-manual.log
Unified VAE file:
.../Lightricks/LTX-2/.../vae/diffusion_pytorch_model.safetensors
Problem
After adding fallback from missing
vae/encoderandvae/decodersubdirectories to the unified Hugging Facevae/directory, image-to-video generation gets past VAE file discovery but fails during VAE encoder forward:Root Cause
The unified
Lightricks/LTX-2VAE safetensors file uses unprefixed Diffusers component keys:Current sanitizers only accept prefixed keys:
A local key comparison showed:
So the unified VAE file is found, but the encoder/decoder weights are not loaded. The runtime then executes an effectively uninitialized/default VAE path and fails with a shape mismatch.
Expected Behavior
VideoEncoder.sanitize()andVideoDecoder.sanitize()should accept both layouts:and normalize them into the local MLX module key names before loading.
Local Evidence
Failed log:
Unified VAE file: