debug: fix VAE load_state_dict diagnostic — load into .model directly

AutoencoderPretransform.load_state_dict() doesn't return IncompatibleKeys.
Load into pretransform.model (AudioAutoencoder) to get the return value
and see actual missing/unexpected key counts.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-27 21:56:06 +01:00
parent f2705b3063
commit 11457fc27a
+3 -1
View File
@@ -125,7 +125,9 @@ class PrismAudioModelLoader:
# strict=False: vae.ckpt is a training checkpoint that also contains # strict=False: vae.ckpt is a training checkpoint that also contains
# discriminator, loss modules, and EMA wrappers not present in the # discriminator, loss modules, and EMA wrappers not present in the
# inference AudioAutoencoder — ignore those extra keys. # inference AudioAutoencoder — ignore those extra keys.
vae_result = model.pretransform.load_state_dict(vae_state, strict=False) # Load directly into the inner AudioAutoencoder to get IncompatibleKeys back
# (AutoencoderPretransform.load_state_dict doesn't return the result)
vae_result = model.pretransform.model.load_state_dict(vae_state, strict=False)
print(f"[PrismAudio] VAE load: missing={len(vae_result.missing_keys)}, unexpected={len(vae_result.unexpected_keys)}", flush=True) print(f"[PrismAudio] VAE load: missing={len(vae_result.missing_keys)}, unexpected={len(vae_result.unexpected_keys)}", flush=True)
if vae_result.missing_keys: if vae_result.missing_keys:
print(f"[PrismAudio] VAE missing (first 10): {vae_result.missing_keys[:10]}", flush=True) print(f"[PrismAudio] VAE missing (first 10): {vae_result.missing_keys[:10]}", flush=True)