From af6458d1b54bd04ac00d0509033193f1ebef5e63 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 28 Apr 2026 13:39:20 -0600 Subject: [PATCH] Enable caching of ACE step latents. --- .../ace_step/ace_step_15_model.py | 4 +- .../audio_models/ace_step/src/model.py | 50 +++++++++++++++++++ .../audio_models/ace_step/src/pipeline.py | 6 ++- toolkit/dataloader_mixins.py | 9 ++-- 4 files changed, 63 insertions(+), 6 deletions(-) diff --git a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py index 4dc089c..430f82b 100644 --- a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py +++ b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py @@ -158,6 +158,8 @@ class AceStep15Model(BaseAudioModel): tokenizer=self.tokenizer, scheduler=self.get_train_scheduler(), ) + if self.model_config.low_vram: + self.pipeline.do_tiled_decoding = True def get_prompt_embeds(self, prompt: str) -> PromptEmbeds: if isinstance(prompt, str): @@ -325,7 +327,7 @@ class AceStep15Model(BaseAudioModel): self.vae.to(device) output = self.vae.encode(audio_tensor.to(device=device, dtype=dtype)) # transpose from [B, 64, T] to [B, T, 64] for DiT - output = output.transpose(1, 2) + output = output.transpose(1, 2).contiguous() return output diff --git a/extensions_built_in/audio_models/ace_step/src/model.py b/extensions_built_in/audio_models/ace_step/src/model.py index 60a364f..8bf9ffa 100644 --- a/extensions_built_in/audio_models/ace_step/src/model.py +++ b/extensions_built_in/audio_models/ace_step/src/model.py @@ -1081,6 +1081,7 @@ class OobleckVAE(nn.Module): ] self.decoder = _SeqWrap(*dec) self.bottleneck = VAEBottleneck() + self.upscale_factor = math.prod(strides) def encode(self, x): return self.bottleneck.encode(self.encoder(x)) @@ -1088,6 +1089,55 @@ class OobleckVAE(nn.Module): def decode(self, x): return self.decoder(self.bottleneck.decode(x)) + def tiled_decode(self, x, tile_seconds=10.0, overlap_seconds=1.0): + """VRAM-light decode: split the latent into ~tile_seconds tiles with + overlap_seconds of overlap, decode each tile independently, and + linearly crossfade the overlapping audio regions.""" + z = self.bottleneck.decode(x) + tile_frames = max(1, round(tile_seconds * LATENT_RATE)) + overlap_frames = max(1, round(overlap_seconds * LATENT_RATE)) + if overlap_frames >= tile_frames: + raise ValueError("overlap_seconds must be smaller than tile_seconds") + + T = z.shape[-1] + if T <= tile_frames: + return self.decoder(z) + + step = tile_frames - overlap_frames + fade_len = overlap_frames * self.upscale_factor + out_T = T * self.upscale_factor + + out = None + ramp = None + write_pos = 0 + + for i, start in enumerate(range(0, T, step)): + end = min(start + tile_frames, T) + decoded = self.decoder(z[..., start:end]) + + if out is None: + out = decoded.new_zeros(decoded.shape[0], decoded.shape[1], out_T) + ramp = torch.linspace(0, 1, fade_len, device=decoded.device, dtype=decoded.dtype) + + if i == 0: + n = decoded.shape[-1] + out[..., :n] = decoded + write_pos = n + else: + blend_start = write_pos - fade_len + out[..., blend_start:blend_start + fade_len] = ( + out[..., blend_start:blend_start + fade_len] * (1 - ramp) + + decoded[..., :fade_len] * ramp + ) + tail = decoded.shape[-1] - fade_len + out[..., write_pos:write_pos + tail] = decoded[..., fade_len:] + write_pos += tail + + if end == T: + break + + return out[..., :write_pos] + @property def device(self): return next(self.parameters()).device diff --git a/extensions_built_in/audio_models/ace_step/src/pipeline.py b/extensions_built_in/audio_models/ace_step/src/pipeline.py index def452a..57edff6 100644 --- a/extensions_built_in/audio_models/ace_step/src/pipeline.py +++ b/extensions_built_in/audio_models/ace_step/src/pipeline.py @@ -36,6 +36,7 @@ class AceStep15Pipeline: self.text_encoder: TextEncoder = text_encoder self.tokenizer: AutoTokenizer = tokenizer self.scheduler = scheduler + self.do_tiled_decoding = False def to(self, *args, **kwargs): self.transformer.to(*args, **kwargs) @@ -158,6 +159,9 @@ class AceStep15Pipeline: xt = xt - vt * (tv - t_sched_t[i + 1].item()) # VAE decode - wav = self.vae.decode(xt.transpose(1, 2)) # [1, 2, samples] + if self.do_tiled_decoding: + wav = self.vae.tiled_decode(xt.transpose(1, 2)) # [1, 2, samples] + else: + wav = self.vae.decode(xt.transpose(1, 2)) # [1, 2, samples] wav = wav[0, :, : int(duration * SAMPLE_RATE)] return wav diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index cb03d8d..ffd2f86 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -1820,6 +1820,9 @@ class LatentCachingFileItemDTOMixin: item["audio_normalize"] = True if self.dataset_config.audio_preserve_pitch: item["audio_preserve_pitch"] = True + if self.is_audio_model: + item["is_audio_model"] = True + item["sample_rate"] = self.sample_rate return item def get_latent_path(self: 'FileItemDTO', recalculate=False): @@ -1882,8 +1885,6 @@ class LatentCachingMixin: self.latent_cache = {} def cache_latents_all_latents(self: 'AiToolkitDataset'): - if self.is_audio_model: - raise Exception("Audio models are not supported for latent caching yet") with accelerator.main_process_first(): print_acc(f"Caching latents for {self.dataset_path}") # cache all latents to disk @@ -1949,8 +1950,8 @@ class LatentCachingMixin: if to_disk: state_dict['first_frame_latent'] = first_frame_latent.clone().detach().cpu() - # audio - if file_item.audio_data is not None: + # audio (video+audio models only — audio-only models already encoded above via encode_images) + if not self.is_audio_model and file_item.audio_data is not None: audio_latent = self.sd.encode_audio([file_item.audio_data]).squeeze(0) if to_disk: state_dict['audio_latent'] = audio_latent.clone().detach().cpu()