Fix ACE-Step audio sample saving (bf16 dtype + waveform shape)

[SanDiegoDude]

Summary

Saving generated audio samples for ACE-Step (1.5 and 1.5 XL) currently fails before any sample is written, due to two issues in GenerateImageConfig.save_image in toolkit/config_modules.py:

1. bf16 dtype — with the default train.dtype: bf16, the generated waveform is bfloat16, which torchaudio/ffmpeg cannot encode:

   ValueError: No format found for dtype torch.bfloat16; dtype must be one of
   [torch.uint8, torch.int16, torch.int32, torch.int64, torch.float32, torch.float64].
   

2. Waveform shape — the ACE-Step pipeline already returns a [channels, time] tensor (it squeezes the batch dim internally), but save_image indexes image[0] again, collapsing it to 1D. ffmpeg then fails with:

   RuntimeError: Failed to create input filter:
   "time_base=1/48000:sample_rate=48000:sample_fmt=flt:channel_layout=0x0" (Invalid argument)
   

The fix casts the waveform to float32 and normalizes it to a 2D [channels, time] tensor (handling 1D/2D/3D inputs) before calling torchaudio.save.

Test plan

• Train ACE-Step 1.5 XL with low_vram: true, train.dtype: bf16, and sampling enabled; baseline samples now write as valid 180s MP3s and training proceeds.
• Sanity-check a non-audio (image) model still saves correctly (this branch only touches the audio output_ext path).

Made with Cursor