evan/ai-toolkit
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix and optimize Ace Step captioner

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-09-15 06:49:06 -06:00
1 parent 59c7de0539
commit 07de70940a
1 file changed
+408 -103
Download Patch File Download Diff File Expand all files Collapse all files
extensions_built_in/captioner/AceStepCaptioner.py
+408 -103
View File
@@ -6,9 +6,19 @@ except ImportError:
librosa = None
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
from transformers import AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
Qwen2_5OmniAudioEncoder,
Qwen2_5OmniThinkerForConditionalGeneration,
)
from transformers.modeling_outputs import BaseModelOutputWithPooling
from collections import OrderedDict
import traceback
from optimum.quanto import freeze
from toolkit.basic import flush
@@ -27,6 +37,16 @@ TARGET_SAMPLE_RATE = 16000
CAPTIONER_ID = "ACE-Step/acestep-captioner"
TRANSCRIBER_ID = "ACE-Step/acestep-transcriber"
LYRICS_PROMPT = "*Task* Transcribe this audio in detail"
CAPTION_PROMPT = "*Task* Describe this music in detail. Include genre, mood, instrumentation, tempo feel, and vocal style if present."
# same budget the Qwen2.5-Omni top-level generate used (thinker_max_new_tokens)
MAX_NEW_TOKENS = 1024
# fixed cache length under compiled decode: the processor caps audio at 300s
# (7500 audio tokens) so prompt + MAX_NEW_TOKENS always fits; a constant keeps
# the static kv cache (and the compiled decode graph) at one shape for every file
STATIC_MAX_LENGTH = 8704
# Key profiles for Krumhansl-Schmuckler key detection
MAJOR_PROFILE = np.array(
[6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]
@@ -105,6 +125,204 @@ def analyze_audio(audio_path):
}
def ostris_sdpa_attention_forward(
module, query, key, value, attention_mask, dropout=0.0, scaling=None, is_causal=None, **kwargs
):
"""sdpa for prefill; for single-query decode against the static cache sdpa's
masked (memory-efficient) kernel takes ~1ms per layer, 27ms/token on the 7B.
A grouped matmul reads the cache once instead."""
if query.shape[2] != 1:
return sdpa_attention_forward(
module, query, key, value, attention_mask, dropout=dropout, scaling=scaling, is_causal=is_causal, **kwargs
)
b, hq, _, d = query.shape
hk = key.shape[1]
if scaling is None:
scaling = d**-0.5
q = query.view(b, hk, hq // hk, d)
scores = torch.matmul(q, key.transpose(-1, -2)) * scaling # [B, hk, g, L]
if attention_mask is not None:
mask = attention_mask[:, :, -1:, : key.shape[2]]
if mask.dtype == torch.bool:
scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
else:
scores = scores + mask
probs = torch.softmax(scores.float(), dim=-1).to(query.dtype)
out = torch.matmul(probs, value) # [B, hk, g, d]
return out.view(b, hq, 1, d).transpose(1, 2).contiguous(), None
AttentionInterface.register("ostris_sdpa", ostris_sdpa_attention_forward)
AttentionMaskInterface.register("ostris_sdpa", sdpa_mask)
class OstrisQwen25OmniAudioEncoder(Qwen2_5OmniAudioEncoder):
"""Upstream flattens every 2s window of the whole clip into one sequence
and attends with a dense [S, S] block-diagonal mask (S = 15000 for 300s):
O(S^2) memory and time for attention that is only ever within a window.
Batch the windows instead: [num_windows, 100, d] with a key-padding mask
for the short tail window. Same math, ~150x less attention work."""
def forward(self, input_features, feature_lens=None, aftercnn_lens=None, **kwargs):
chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
chunk_lengths = torch.full(
(chunk_num.sum(),), self.n_window * 2, dtype=torch.long, device=feature_lens.device
)
tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
chunk_lengths = torch.where(chunk_lengths == 0, self.n_window * 2, chunk_lengths)
chunk_list = input_features.split(chunk_lengths.tolist(), dim=1)
padded_feature, padded_mask, padded_mask_after_cnn = self.padded_and_mask_function(
chunk_list, chunk_lengths, padding_value=0, padding_side="right"
)
x = F.gelu(self.conv1(padded_feature)) * padded_mask
x = F.gelu(self.conv2(x)).transpose(1, 2)
x = x + self.positional_embedding.positional_embedding[: x.shape[1], :].unsqueeze(0).to(x.dtype)
key_mask = None
if not bool(padded_mask_after_cnn.all()):
key_mask = padded_mask_after_cnn[:, None, None, :]
n, seq_len, dim = x.shape
for layer in self.layers:
attn = layer.self_attn
residual = x
h = layer.self_attn_layer_norm(x)
q = attn.q_proj(h).view(n, seq_len, attn.num_heads, -1).transpose(1, 2)
k = attn.k_proj(h).view(n, seq_len, attn.num_heads, -1).transpose(1, 2)
v = attn.v_proj(h).view(n, seq_len, attn.num_heads, -1).transpose(1, 2)
o = F.scaled_dot_product_attention(q, k, v, attn_mask=key_mask, scale=attn.scaling)
x = residual + attn.out_proj(o.transpose(1, 2).reshape(n, seq_len, dim))
residual = x
h = layer.final_layer_norm(x)
x = residual + layer.fc2(layer.activation_fn(layer.fc1(h)))
hidden_states = x[padded_mask_after_cnn]
outs = []
for each in hidden_states.split(aftercnn_lens.tolist(), dim=0):
each = self.avg_pooler(each.transpose(0, 1)).transpose_(0, 1)
outs.append(self.proj(self.ln_post(each)))
return BaseModelOutputWithPooling(last_hidden_state=torch.cat(outs, dim=0))
class OstrisQwen25OmniThinker(Qwen2_5OmniThinkerForConditionalGeneration):
"""Thinker with static-cache-safe MRoPE handling (same fix as the Qwen3-Omni
captioner): under a compileable cache generate hands forward a 4D mask, but
upstream's rope block does `1 - attention_mask` and get_rope_index needs the
2D padding mask. Prefill uses the true 2D mask stashed by the caller; decode
continues from the cache length with no data-dependent ops."""
_pad_mask_2d = None
# per-file media shapes must not reach the compiled decode graph or it
# recompiles for every file
_PREFILL_ONLY_KEYS = (
"input_features",
"feature_attention_mask",
"audio_feature_lengths",
"pixel_values",
"pixel_values_videos",
"image_grid_thw",
"video_grid_thw",
"video_second_per_grid",
)
def prepare_inputs_for_generation(self, *args, **kwargs):
model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
ids = model_inputs.get("input_ids", None)
if ids is not None and ids.shape[1] == 1:
for key in self._PREFILL_ONLY_KEYS:
model_inputs.pop(key, None)
return model_inputs
def forward(
self,
input_ids=None,
input_features=None,
pixel_values=None,
pixel_values_videos=None,
image_grid_thw=None,
video_grid_thw=None,
attention_mask=None,
feature_attention_mask=None,
audio_feature_lengths=None,
position_ids=None,
past_key_values=None,
use_audio_in_video=None,
video_second_per_grid=None,
**kwargs,
):
if position_ids is None and input_ids is not None:
if input_ids.shape[1] > 1 or self.rope_deltas is None:
mask2d = (
attention_mask
if isinstance(attention_mask, torch.Tensor) and attention_mask.dim() == 2
else self._pad_mask_2d
)
if mask2d is None:
mask2d = torch.ones_like(input_ids)
mask2d = mask2d.long()
if mask2d.shape[1] != input_ids.shape[1]:
mask2d = mask2d[:, : input_ids.shape[1]]
audio_lens = (
feature_attention_mask.sum(1)
if feature_attention_mask is not None
else audio_feature_lengths
)
delta0 = (1 - mask2d).sum(dim=-1).unsqueeze(1)
position_ids, rope_deltas = self.get_rope_index(
input_ids,
image_grid_thw,
video_grid_thw,
mask2d,
use_audio_in_video or False,
audio_lens,
video_second_per_grid,
)
self.rope_deltas = rope_deltas - delta0
else:
batch_size, seq_length = input_ids.shape
# static cache length is a device tensor: stays sync-free
past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
pos = (
torch.arange(seq_length, device=input_ids.device).view(1, -1)
+ past_len
+ self.rope_deltas.to(input_ids.device)
)
position_ids = pos.unsqueeze(0).expand(3, batch_size, seq_length)
return super().forward(
input_ids=input_ids,
input_features=input_features,
pixel_values=pixel_values,
pixel_values_videos=pixel_values_videos,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
attention_mask=attention_mask,
feature_attention_mask=feature_attention_mask,
audio_feature_lengths=audio_feature_lengths,
position_ids=position_ids,
past_key_values=past_key_values,
use_audio_in_video=use_audio_in_video,
video_second_per_grid=video_second_per_grid,
**kwargs,
)
def prepare_thinker(full: Qwen2_5OmniForConditionalGeneration):
"""Strip a loaded Qwen2.5-Omni to its thinker with the fast paths attached."""
# drop the talker before anything touches the GPU
full.disable_talker()
model = full.thinker
model.__class__ = OstrisQwen25OmniThinker
model.audio_tower.__class__ = OstrisQwen25OmniAudioEncoder
model.model.config._attn_implementation = "ostris_sdpa"
model.generation_config.eos_token_id = full.generation_config.eos_token_id
model.generation_config.pad_token_id = (
full.generation_config.pad_token_id
if full.generation_config.pad_token_id is not None
else 151643
)
return model
class AceStepCaptionConfig(CaptionConfig):
def __init__(self, **kwargs):
super().__init__(**kwargs)
@@ -118,54 +336,61 @@ class AceStepCaptioner(BaseCaptioner):
def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
super(AceStepCaptioner, self).__init__(process_id, job, config, **kwargs)
def load_model(self):
self.print_and_status_update("Loading transcriber model")
self.model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
self.caption_config.model_name_or_path,
def _load_thinker(self, name_or_path: str, label: str):
"""Load a Qwen2.5-Omni checkpoint and keep only its thinker."""
self.print_and_status_update(f"Loading {label} model")
full = Qwen2_5OmniForConditionalGeneration.from_pretrained(
name_or_path,
dtype=self.torch_dtype,
device_map="cpu",
)
self.model.to(self.device_torch)
self.model.disable_talker()
model = prepare_thinker(full)
model.to(self.device_torch)
if self.caption_config.quantize:
self.print_and_status_update("Quantizing transcriber model")
quantize(self.model, weights=get_qtype(self.caption_config.qtype))
freeze(self.model)
self.print_and_status_update(f"Quantizing {label} model")
quantize(model, weights=get_qtype(self.caption_config.qtype))
freeze(model)
flush()
self.processor = Qwen2_5OmniProcessor.from_pretrained(
self.caption_config.model_name_or_path
)
processor = Qwen2_5OmniProcessor.from_pretrained(name_or_path)
if self.caption_config.low_vram:
self.model.to("cpu")
model.to("cpu")
return model, processor
def load_model(self):
self.model, self.processor = self._load_thinker(
self.caption_config.model_name_or_path, "transcriber"
)
self.model2 = None
self.processor2 = None
if self.caption_config.fixed_caption is not None:
# load captioner model
self.print_and_status_update("Loading captioner model")
self.model2 = Qwen2_5OmniForConditionalGeneration.from_pretrained(
self.caption_config.model_name_or_path2,
dtype=self.torch_dtype,
device_map="cpu",
if self.caption_config.fixed_caption is None:
self.model2, self.processor2 = self._load_thinker(
self.caption_config.model_name_or_path2, "captioner"
)
self.model2.to(self.device_torch)
self.model2.disable_talker()
if self.caption_config.quantize:
self.print_and_status_update("Quantizing captioner model")
quantize(self.model2, weights=get_qtype(self.caption_config.qtype))
freeze(self.model2)
flush()
self.processor2 = Qwen2_5OmniProcessor.from_pretrained(
self.caption_config.model_name_or_path2,
)
if self.caption_config.low_vram:
self.model2.to("cpu")
flush()
def run_qwen_audio(self, model, processor, audio_data, sr, prompt_text):
"""Run a Qwen2.5-Omni model on audio with a text prompt."""
def maybe_compile_models(self):
"""CUDA-graph decode via a static kv cache (see Qwen3OmniCaptioner):
eager HF decode on the 7B thinker is launch-bound at ~17 tok/s."""
if not self.caption_config.compile:
return
if self.caption_config.low_vram:
print("[AITK] low_vram is on; skipping compiled decode.")
return
import importlib.util
if importlib.util.find_spec("triton") is None:
print("[AITK] compile requested but triton is not installed, skipping.")
return
for model in (self.model, self.model2):
if model is not None:
model.generation_config.cache_implementation = "static"
print(
"[AITK] Compiled decode enabled (static cache + cuda graphs). "
"The first file per model compiles (~2 min cold, faster once cached)."
)
@staticmethod
def _make_inputs(processor, audio_data, prompt_text):
conversation = [
{
"role": "user",
@@ -178,92 +403,172 @@ class AceStepCaptioner(BaseCaptioner):
text = processor.apply_chat_template(
conversation, add_generation_prompt=True, tokenize=False
)
inputs = processor(
return processor(
text=text,
audio=[audio_data],
images=None,
videos=None,
return_tensors="pt",
padding=True,
sampling_rate=sr,
sampling_rate=TARGET_SAMPLE_RATE,
)
inputs = inputs.to(model.device).to(model.dtype)
text_ids = model.generate(**inputs, return_audio=False)
output = processor.batch_decode(
text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
result = output[0]
marker = "assistant\n"
if marker in result:
result = result[result.rfind(marker) + len(marker) :]
return result.strip()
def get_audio_lyrics(self, audio_data: torch.Tensor) -> str:
if self.caption_config.low_vram and self.model2.device != torch.device("cpu"):
# move captioner to cpu
def _prep_file(self, file_path: str) -> dict:
"""CPU side of one file (runs in a worker thread): librosa analysis,
decode + resample, and the mel/tokenize pass for each model."""
analysis = analyze_audio(file_path)
waveform, sr = torchaudio.load(file_path)
if waveform.shape[0] > 1:
waveform = waveform.mean(dim=0, keepdim=True)
if sr != TARGET_SAMPLE_RATE:
waveform = torchaudio.functional.resample(waveform, sr, TARGET_SAMPLE_RATE)
audio_data = waveform.squeeze(0).numpy()
item = {
"file": file_path,
"analysis": analysis,
"lyrics_inputs": self._make_inputs(self.processor, audio_data, LYRICS_PROMPT),
}
if self.model2 is not None:
item["caption_inputs"] = self._make_inputs(
self.processor2, audio_data, CAPTION_PROMPT
)
return item
def _generate(self, model, processor, inputs) -> str:
inputs = inputs.to(model.device).to(model.dtype)
# a generate that dies between static-cache creation and its first
# forward leaves a half-built cache that breaks every later call
stale_cache = getattr(model, "_cache", None)
if stale_cache is not None and not stale_cache.is_initialized:
del model._cache
model._pad_mask_2d = inputs.get("attention_mask", None)
input_len = inputs["input_ids"].shape[1]
gen_kwargs = {"max_new_tokens": MAX_NEW_TOKENS}
if model.generation_config.cache_implementation == "static":
if input_len + 16 < STATIC_MAX_LENGTH:
from transformers.generation import MaxLengthCriteria, StoppingCriteriaList
gen_kwargs = {
"max_length": STATIC_MAX_LENGTH,
"stopping_criteria": StoppingCriteriaList(
[MaxLengthCriteria(max_length=min(input_len + MAX_NEW_TOKENS, STATIC_MAX_LENGTH))]
),
}
else:
# prompt would not fit the fixed cache; eager for this file
gen_kwargs["cache_implementation"] = "dynamic"
generated = model.generate(**inputs, **gen_kwargs)
output = processor.batch_decode(
generated[:, input_len:],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0]
return output.strip()
def get_audio_lyrics(self, inputs) -> str:
if (
self.caption_config.low_vram
and self.model2 is not None
and self.model2.device != torch.device("cpu")
):
self.model2.to("cpu")
# move lyric model if needed
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
return self._generate(self.model, self.processor, inputs)
prompt_text = "*Task* Transcribe this audio in detail"
return self.run_qwen_audio(
self.model, self.processor, audio_data, TARGET_SAMPLE_RATE, prompt_text
)
def get_audio_caption(self, audio_data: torch.Tensor) -> str:
def get_audio_caption(self, inputs) -> str:
if self.caption_config.low_vram and self.model.device != torch.device("cpu"):
# move lyricmodel to cpu
self.model.to("cpu")
# move captioner model if needed
if self.model2.device == torch.device("cpu"):
self.model2.to(self.device_torch)
prompt_text = "*Task* Describe this music in detail. Include genre, mood, instrumentation, tempo feel, and vocal style if present."
return self.run_qwen_audio(
self.model2, self.processor2, audio_data, TARGET_SAMPLE_RATE, prompt_text
return self._generate(self.model2, self.processor2, inputs)
def _caption_item(self, item: dict) -> str:
analysis = item["analysis"]
lyrics = self.get_audio_lyrics(item["lyrics_inputs"])
language = "en"
if "# Languages" in lyrics and "# Lyrics" in lyrics:
language = lyrics.split("# Languages")[1].split("# Lyrics")[0]
# remove newlines and extra spaces from language
language = language.replace("\n", "").strip()
lyrics = lyrics.split("# Lyrics")[1].strip()
if self.caption_config.fixed_caption is not None:
caption = self.caption_config.fixed_caption
else:
caption = self.get_audio_caption(item["caption_inputs"])
output = f"<CAPTION>\n{caption}\n</CAPTION>\n"
output += f"<LYRICS>\n{lyrics}\n</LYRICS>\n"
output += f"<BPM>{analysis['bpm']}</BPM>\n"
output += f"<KEYSCALE>{analysis['keyscale']}</KEYSCALE>\n"
output += f"<TIMESIGNATURE>{analysis['timesignature']}</TIMESIGNATURE>\n"
output += f"<DURATION>{analysis['duration']}</DURATION>\n"
output += f"<LANGUAGE>{language}</LANGUAGE>"
return output
def run_caption_loop(self):
"""Worker threads run librosa + decode + feature extraction ahead of
the GPU so the main thread only moves tensors and generates."""
import concurrent.futures
from collections import deque
import tqdm as tqdm_mod
pbar = tqdm_mod.tqdm(
total=len(self.file_paths), desc="Captioning files", unit="file", smoothing=0.9
)
def finish(file_path, caption):
if caption is not None:
self.save_caption_for_file(file_path, caption)
self.step_num += 1
self.update_step()
pbar.update(1)
executor = concurrent.futures.ThreadPoolExecutor(
max_workers=max(1, int(self.caption_config.num_workers))
)
try:
futures = deque()
file_iter = iter(self.file_paths)
lookahead = max(2, int(self.caption_config.num_workers)) + 1
for _ in range(lookahead):
path = next(file_iter, None)
if path is None:
break
futures.append((path, executor.submit(self._prep_file, path)))
while futures:
if self.is_ui_captioner:
self.maybe_stop()
if self.is_stopping:
break
path, fut = futures.popleft()
nxt = next(file_iter, None)
if nxt is not None:
futures.append((nxt, executor.submit(self._prep_file, nxt)))
try:
item = fut.result()
except Exception as e:
print(f"Error preprocessing {path}: {e}")
finish(path, None)
continue
try:
caption = self._caption_item(item)
except Exception as e:
print(f"Error processing {path}: {e}")
traceback.print_exc()
caption = None
finish(path, caption)
finally:
executor.shutdown(wait=False, cancel_futures=True)
pbar.close()
def get_caption_for_file(self, file_path: str) -> str:
try:
# analyze audio with librosa
analysis = analyze_audio(file_path)
# load audio with torchaudio for transcription
waveform, sr = torchaudio.load(file_path)
waveform = waveform.to(self.device_torch)
if waveform.shape[0] > 1:
waveform = waveform.mean(dim=0, keepdim=True)
if sr != TARGET_SAMPLE_RATE:
waveform = torchaudio.functional.resample(
waveform, sr, TARGET_SAMPLE_RATE
)
audio_data = waveform.squeeze(0).cpu().numpy()
# get the lyrics from the audio
lyrics = self.get_audio_lyrics(audio_data)
language = "en"
if "# Languages" in lyrics and "# Lyrics" in lyrics:
language = lyrics.split("# Languages")[1].split("# Lyrics")[0]
# remove newlines and extra spaces from language
language = language.replace("\n", "").strip()
lyrics = lyrics.split("# Lyrics")[1].strip()
# get the caption from the audio
if self.caption_config.fixed_caption is not None:
caption = self.caption_config.fixed_caption
else:
caption = self.get_audio_caption(audio_data)
output = f"<CAPTION>\n{caption}\n</CAPTION>\n"
output += f"<LYRICS>\n{lyrics}\n</LYRICS>\n"
output += f"<BPM>{analysis['bpm']}</BPM>\n"
output += f"<KEYSCALE>{analysis['keyscale']}</KEYSCALE>\n"
output += f"<TIMESIGNATURE>{analysis['timesignature']}</TIMESIGNATURE>\n"
output += f"<DURATION>{analysis['duration']}</DURATION>\n"
output += f"<LANGUAGE>{language}</LANGUAGE>"
return output
return self._caption_item(self._prep_file(file_path))
except Exception as e:
print(f"Error processing {file_path}: {e}")
traceback.print_exc()
return None