LTX-2 Audio VAE: state dict key mapping missing for official checkpoint
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Summary
The official `ltx2_audio_vae.safetensors` checkpoint from Lightricks uses different state dict key names than what `AudioVAE.load_state_dict()` expects, causing the model to fail loading.
## Problem
When loading `ltx2_audio_vae.safetensors`, the following keys don't match:
**VAE keys:**
- Checkpoint has: `latents_mean`, `latents_std`
- Expected: `per_channel_statistics.mean-of-means`, `per_channel_statistics.std-of-means`
**Vocoder keys:**
- Checkpoint has: `conv_in`, `conv_out`, `resnets`, `upsamplers`
- Expected: `conv_pre`, `conv_post`, `resblocks`, `ups`
## Proposed Fix
Add key mapping in `AudioVAE.load_state_dict()` to support both formats:
```python
# Map latents_mean/latents_std to per_channel_statistics buffer names
if "latents_mean" in vae_sd and "latents_std" in vae_sd:
vae_sd["per_channel_statistics.mean-of-means"] = vae_sd.pop("latents_mean")
vae_sd["per_channel_statistics.std-of-means"] = vae_sd.pop("latents_std")
# Map vocoder key names from diffusers format to HiFi-GAN format
vocoder_sd = utils.state_dict_prefix_replace(vocoder_sd, {
"conv_in.": "conv_pre.",
"conv_out.": "conv_post.",
"resnets.": "resblocks.",
"upsamplers.": "ups.",
})
```
## Reference Implementation
I've implemented this fix in my fork: https://github.com/m0nk1111/ComfyUI/commit/15387800
## Environment
- ComfyUI version: latest master (dc202a2e)
- Checkpoint: `ltx2_audio_vae.safetensors` from Lightricks (via Hugging Face)
Contributor guide
Assessment
This issue has not been assessed yet.