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

extensions_built_in/captioner/AceStepCaptioner.py

@@ -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(
+from extensions_built_in.llm_models.src.thinker import (  # noqa: E402  shared fast paths
-    module, query, key, value, attention_mask, dropout=0.0, scaling=None, is_causal=None, **kwargs
+    OstrisQwen25OmniAudioEncoder,
-):
+    OstrisQwen25OmniThinker,
-    """sdpa for prefill; for single-query decode against the static cache sdpa's
+    ostris_sdpa_attention_forward,
-    masked (memory-efficient) kernel takes ~1ms per layer, 27ms/token on the 7B.
+    prepare_thinker,
-    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)
+# output layouts:
-AttentionMaskInterface.register("ostris_sdpa", sdpa_mask)
+#   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):
+def clean_lyrics(text: str) -> str:
-    """Upstream flattens every 2s window of the whole clip into one sequence
+    """Strip the transcriber's '# Languages ... # Lyrics' header, keep the sectioned lyrics."""
-    and attends with a dense [S, S] block-diagonal mask (S = 15000 for 300s):
+    if "# Lyrics" in text:
-    O(S^2) memory and time for attention that is only ever within a window.
+        text = text.split("# Lyrics", 1)[1]
-    Batch the windows instead: [num_windows, 100, d] with a key-padding mask
+    elif text.lstrip().startswith("# Languages"):
-    for the short tail window. Same math, ~150x less attention work."""
+        text = ""
-
+    text = text.strip()
-    def forward(self, input_features, feature_lens=None, aftercnn_lens=None, **kwargs):
+    return text if text else "[Instrumental]"
-        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)
         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:
+    @staticmethod
-        """CPU side of one file (runs in a worker thread): librosa analysis,
+    def _load_audio(file_path: str):
-        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()
+        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

@@ -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

@@ -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

@@ -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

@@ -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;