Honor low_vram in ACE-Step load_model (fix OOM during latent caching)

[SanDiegoDude]

Summary

When training ACE-Step (1.5 / 1.5 XL) with low_vram: true, latent caching OOMs on longer audio clips even though the models are small enough to fit otherwise.

Root cause is in AceStep15Model.load_model (extensions_built_in/audio_models/ace_step/ace_step_15_model.py). Under low_vram, the transformer is moved to the CPU after quantization, but load_model then unconditionally moves both the transformer and text encoder back to the GPU:

if self.model_config.low_vram:
    self.print_and_status_update("Moving transformer to CPU")
    self.model.to("cpu")
...
# move back to device
self.model.to(device)        # <- defeats low_vram
self.text_encoder.to(device)
self.vae.to(device)

So low_vram provides no VRAM savings during latent caching. Caching encodes each clip through the VAE while the full quantized DiT + text encoder are still resident, and a longer clip OOMs in the VAE encoder:

cache_latents_all_latents -> encode_audio -> vae.encode -> self.encoder(x)
  -> SnakeBeta.forward: x + (1.0 / (b + 1e-9)) * torch.sin(x * a).pow(2)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.82 GiB ...

(observed on a 24 GB card with a ~5.5 minute training clip; the ~19.5 GB resident was the model weights, not the VAE activation.)

Fix

Keep the transformer and text encoder on the CPU when low_vram is set, and rely on the existing on-demand .to(device) moves already present in get_prompt_embeds, get_noise_prediction, and generate_single_audio. Only the VAE (needed for latent caching) stays resident. Non-low_vram behavior is unchanged.

Notes / scope

• This matches the apparent design intent (the other code paths already move submodules on demand, and layer_offloading raises NotImplementedError).
• This addresses the caching-time OOM specifically. A fuller low_vram story could also move the transformer back to CPU after sampling/generation; that's intentionally left out of this minimal fix and can be a follow-up.

Test plan

• ACE-Step 1.5 XL, low_vram: true, 24 GB GPU, dataset with clips up to ~5.5 min: latent caching now completes (peaks at a few GB instead of ~19.5 GB) and training proceeds.
• Confirm low_vram: false path is unchanged (weights load directly to device).

Made with Cursor