LTX-2 I2V fails because VAE encoder/decoder weights are stored in unified vae directory

[andyelka-creator]

Problem

After config loading and transformer key sanitization were fixed, an LTX-2 image-to-video run reaches transformer loading successfully but fails when loading the VAE encoder.

Command context:

mlx_video.ltx_2.generate \
  --pipeline distilled \
  --height 512 \
  --width 512 \
  --num-frames 97 \
  --model-repo Lightricks/LTX-2 \
  --image /Users/andrejeliseev/Documents/horse1.png

Failure:

✓ Transformer loaded
Auto-detected STG blocks: [29] (model=2)
Traceback (most recent call last):
  File "mlx_video/models/ltx_2/generate.py", line 2098, in generate_video
    vae_encoder = VideoEncoder.from_pretrained(model_path / "vae" / "encoder")
  File "mlx_video/models/ltx_2/video_vae/video_vae.py", line 445, in from_pretrained
    raise FileNotFoundError(f"No safetensors files found in {model_path}")
FileNotFoundError: No safetensors files found in .../Lightricks/LTX-2/.../vae/encoder

Root Cause

The current generation code assumes converted split VAE directories:

vae/encoder/*.safetensors
vae/decoder/*.safetensors

The Hugging Face Lightricks/LTX-2 snapshot stores the video VAE as a unified directory:

vae/config.json
vae/diffusion_pytorch_model.safetensors

VideoEncoder.sanitize() and VideoDecoder.sanitize() already know how to extract vae.encoder.* and vae.decoder.* keys from unified weights, but from_pretrained() never falls back from the missing vae/encoder or vae/decoder subdirectory to the parent vae directory.

Expected Behavior

When a requested VAE component subdirectory has no safetensors, the loader should support the unified Hugging Face vae/ layout and still load the encoder/decoder component from vae/diffusion_pytorch_model.safetensors.

Local Evidence

vae/config.json
vae/diffusion_pytorch_model.safetensors

Latest failed local log:

outputs/logs/horse-avatar-ltx2-manual.log

[andyelka-creator]

Opened draft PR #31 with a focused fix for this issue.

Validation performed locally:

• Added regression coverage for missing vae/encoder and vae/decoder subdirectories falling back to unified parent vae/ safetensors.
• Verified real downloaded Lightricks/LTX-2 VAE component loading:
  • VideoEncoder.from_pretrained(.../vae/encoder) -> video encoder unified fallback ok
  • VideoDecoder.from_pretrained(.../vae/decoder, strict=False) -> video decoder unified fallback ok
• Relevant pytest subset: 20 passed, 2 warnings.

PR: #31