diff --git a/extensions_built_in/captioner/AceStepCaptioner.py b/extensions_built_in/captioner/AceStepCaptioner.py
index e1d77b3..b84cb13 100644
--- a/extensions_built_in/captioner/AceStepCaptioner.py
+++ b/extensions_built_in/captioner/AceStepCaptioner.py
@@ -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"\n{caption}\n\n"
+ output += f"\n{lyrics}\n\n"
+ output += f"{analysis['bpm']}\n"
+ output += f"{analysis['keyscale']}\n"
+ output += f"{analysis['timesignature']}\n"
+ output += f"{analysis['duration']}\n"
+ output += f"{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"\n{caption}\n\n"
- output += f"\n{lyrics}\n\n"
- output += f"{analysis['bpm']}\n"
- output += f"{analysis['keyscale']}\n"
- output += f"{analysis['timesignature']}\n"
- output += f"{analysis['duration']}\n"
- output += f"{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