Normalize decoded Foley audio waveform shape

ltx_foley_v2a/nodes.py

@@ -151,9 +151,21 @@ class LTXFoleyAudioVAEDecode:
         if audio_latent.is_nested:
             audio_latent = audio_latent.unbind()[-1]
 
-        audio = audio_vae.decode(audio_latent).movedim(-1, 1).to(audio_latent.device)
+        audio = audio_vae.decode(audio_latent).to(audio_latent.device)
+        if audio.ndim == 2:
+            audio = audio.unsqueeze(1)
+        elif audio.ndim != 3:
+            raise ValueError(f"Expected decoded audio to have 2 or 3 dimensions, got shape {tuple(audio.shape)}")
+
+        if audio.shape[1] in (1, 2, 6):
+            waveform = audio
+        elif audio.shape[2] in (1, 2, 6):
+            waveform = audio.movedim(-1, 1)
+        else:
+            raise ValueError(f"Could not infer audio channel dimension from decoded shape {tuple(audio.shape)}")
+
         output_sample_rate = _audio_vae_model(audio_vae).output_sample_rate
-        return ({"waveform": audio, "sample_rate": int(output_sample_rate)},)
+        return ({"waveform": waveform, "sample_rate": int(output_sample_rate)},)
 
 
 NODE_CLASS_MAPPINGS = {