Enable caching of ACE step latents.

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
 
 

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

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

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
+                    # audio (video+audio models only — audio-only models already encoded above via encode_images)
-                    if file_item.audio_data is not None:
+                    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()