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

Allow Ace Step captioner to work with different output formats like yue2

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-09-15 12:48:15 -06:00
1 parent 8bf12e476e
commit 6b2526a4d7
5 files changed
+75 -207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

extensions_built_in/captioner/AceStepCaptioner.py
+45 -206
View File
@@ -6,17 +6,8 @@ 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
@@ -126,207 +117,40 @@ 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
from extensions_built_in.llm_models.src.thinker import ( # noqa: E402 shared fast paths
OstrisQwen25OmniAudioEncoder,
OstrisQwen25OmniThinker,
ostris_sdpa_attention_forward,
prepare_thinker,
)
AttentionInterface.register("ostris_sdpa", ostris_sdpa_attention_forward)
AttentionMaskInterface.register("ostris_sdpa", sdpa_mask)
# output layouts:
# ace_step - <CAPTION>/<LYRICS>/<BPM>/<KEYSCALE>/<TIMESIGNATURE>/<DURATION>/<LANGUAGE> tags
# yue2 - description paragraph, then a [Lyrics] line and the sectioned lyrics
# (YuE2's native prompt layout; no librosa analysis needed)
CAPTION_FORMATS = ("ace_step", "yue2")
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
def clean_lyrics(text: str) -> str:
"""Strip the transcriber's '# Languages ... # Lyrics' header, keep the sectioned lyrics."""
if "# Lyrics" in text:
text = text.split("# Lyrics", 1)[1]
elif text.lstrip().startswith("# Languages"):
text = ""
text = text.strip()
return text if text else "[Instrumental]"
class AceStepCaptionConfig(CaptionConfig):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.fixed_caption: Optional[str] = kwargs.get("fixed_caption", None)
self.caption_format: str = kwargs.get("caption_format", "ace_step") or "ace_step"
if self.caption_format not in CAPTION_FORMATS:
raise ValueError(
f"caption_format must be one of {CAPTION_FORMATS}, got {self.caption_format!r}"
)
class AceStepCaptioner(BaseCaptioner):
@@ -413,19 +237,25 @@ class AceStepCaptioner(BaseCaptioner):
sampling_rate=TARGET_SAMPLE_RATE,
)
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)
@staticmethod
def _load_audio(file_path: str):
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()
return waveform.squeeze(0).numpy()
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."""
audio_data = self._load_audio(file_path)
item = {
"file": file_path,
"analysis": analysis,
# BPM/key/time signature only appear in the ace_step layout
"analysis": analyze_audio(file_path)
if self.caption_config.caption_format == "ace_step"
else None,
"lyrics_inputs": self._make_inputs(self.processor, audio_data, LYRICS_PROMPT),
}
if self.model2 is not None:
@@ -484,6 +314,15 @@ class AceStepCaptioner(BaseCaptioner):
return self._generate(self.model2, self.processor2, inputs)
def _caption_item(self, item: dict) -> str:
if self.caption_config.caption_format == "yue2":
lyrics = clean_lyrics(self.get_audio_lyrics(item["lyrics_inputs"]))
if self.caption_config.fixed_caption is not None:
caption = self.caption_config.fixed_caption
else:
caption = self.get_audio_caption(item["caption_inputs"])
caption = " ".join(caption.split())
return f"{caption}\n[Lyrics]\n{lyrics}"
analysis = item["analysis"]
lyrics = self.get_audio_lyrics(item["lyrics_inputs"])
scripts/convert_vllm_to_comfy.py
+8
View File
@@ -120,8 +120,16 @@ class Qwen3OmniHandler(ArchHandler):
return key[len("thinker.") :]
class Qwen25OmniHandler(Qwen3OmniHandler):
"""Qwen2.5-Omni: same layout as Qwen3-Omni (thinker = text model + visual +
audio_tower + lm_head; talker / token2wav dropped). The thinker also
carries the speaker dict; only thinker.* tensors are kept."""
ARCH_HANDLERS = {
"Qwen3OmniMoeForConditionalGeneration": Qwen3OmniHandler,
"Qwen2_5OmniForConditionalGeneration": Qwen25OmniHandler,
"Qwen2_5OmniModel": Qwen25OmniHandler,
}
ui/src/components/CaptionSimpleJob.tsx
+11
View File
@@ -12,6 +12,7 @@ import { CaptionJobConfig } from '@/types';
import { handleCaptionerTypeChange } from '@/helpers/captionJobConfig';
import {
batchSizeOptions,
captionFormatOptions,
captionerTypes,
defaultQtype,
groupedCaptionerTypes,
@@ -93,6 +94,16 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
/>
</div>
)}
{additionalSections.includes('caption.caption_format') && (
<div className="mt-4">
<SelectInput
label="Caption Format"
value={jobConfig.config.process[0].caption.caption_format || 'ace_step'}
onChange={value => setJobConfig(value, 'config.process[0].caption.caption_format')}
options={captionFormatOptions}
/>
</div>
)}
{additionalSections.includes('caption.fixed_caption') && (
<div className="mt-4">
<TextInput
ui/src/helpers/captionOptions.ts
+10 -1
View File
@@ -1,7 +1,7 @@
import { GroupedSelectOption, SelectOption } from "@/types";
type CaptionGroup = 'image' | 'music' | 'video' | 'image/video/sound';
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
export interface CaptionOption {
name: string;
@@ -92,6 +92,8 @@ export const captionerTypes: CaptionOption[] = [
'config.process[0].caption.model_name_or_path': ['ACE-Step/acestep-transcriber', defaultNameOrPath],
'config.process[0].caption.model_name_or_path2': ['ACE-Step/acestep-captioner', undefined],
'config.process[0].caption.extensions': [extensionsAudio, defaultExtensions],
'config.process[0].caption.caption_format': ['ace_step', undefined],
'config.process[0].caption.compile': [true, false],
},
name_or_path_options: [
{ value: 'ACE-Step/acestep-transcriber', label: 'ACE-Step/acestep-transcriber' },
@@ -102,6 +104,7 @@ export const captionerTypes: CaptionOption[] = [
additionalSections: [
'caption.model_name_or_path2',
'caption.fixed_caption',
'caption.caption_format',
],
},
{
@@ -230,6 +233,12 @@ export const quantizationOptions: SelectOption[] = [
{ value: 'uint2', label: '2 bit' },
];
// AceStepCaptioner output layouts (caption.caption_format)
export const captionFormatOptions: SelectOption[] = [
{ value: 'ace_step', label: 'ACE-Step (caption, lyrics, bpm, key, time signature, duration)' },
{ value: 'yue2', label: 'YuE2 (description, then [Lyrics] block)' },
];
export const batchSizeOptions: SelectOption[] = [
{ value: '1', label: '1 (default)' },
{ value: '2', label: '2' },
ui/src/types.ts
+1
View File
@@ -322,6 +322,7 @@ export interface CaptionProcessConfig {
max_res?: number;
max_new_tokens?: number;
fixed_caption?: string;
caption_format?: string;
caption_extension?: string;
thinking?: boolean;
batch_size?: number;