From 8e3ab999f05071d85105880d3c430880cd7f85db Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Fri, 27 Mar 2026 19:51:51 +0100 Subject: [PATCH] fix: load VAE state dict with strict=False vae.ckpt is a full training checkpoint containing discriminator, STFT loss modules, and EMA wrappers that are absent from the inference AudioAutoencoder. strict=False ignores these training-only keys while still loading all encoder/decoder/bottleneck weights correctly. Co-Authored-By: Claude Sonnet 4.6 --- nodes/model_loader.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nodes/model_loader.py b/nodes/model_loader.py index 5d92977..e0ef4c1 100644 --- a/nodes/model_loader.py +++ b/nodes/model_loader.py @@ -105,7 +105,10 @@ class PrismAudioModelLoader: vae_state[k[len(prefix):]] = v else: vae_state[k] = v - model.pretransform.load_state_dict(vae_state) + # strict=False: vae.ckpt is a training checkpoint that also contains + # discriminator, loss modules, and EMA wrappers not present in the + # inference AudioAutoencoder — ignore those extra keys. + model.pretransform.load_state_dict(vae_state, strict=False) # Apply precision: DiT + conditioners in user-selected dtype, # but keep VAE (pretransform) in fp32 to avoid NaN from snake activations in fp16