Comfy-Org / Comfy-Org/ComfyUI

LTX-2 Audio VAE: state dict key mapping missing for official checkpoint

Open
#11,789 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.