From 6b2526a4d7216fef0c09c5f9011a166f0645acd0 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 15 Sep 2026 12:48:15 -0600 Subject: [PATCH] Allow Ace Step captioner to work with different output formats like yue2 --- .../captioner/AceStepCaptioner.py | 251 ++++-------------- scripts/convert_vllm_to_comfy.py | 8 + ui/src/components/CaptionSimpleJob.tsx | 11 + ui/src/helpers/captionOptions.ts | 11 +- ui/src/types.ts | 1 + 5 files changed, 75 insertions(+), 207 deletions(-) diff --git a/extensions_built_in/captioner/AceStepCaptioner.py b/extensions_built_in/captioner/AceStepCaptioner.py index b84cb13..cac935d 100644 --- a/extensions_built_in/captioner/AceStepCaptioner.py +++ b/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( - 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 - ////// 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"]) diff --git a/scripts/convert_vllm_to_comfy.py b/scripts/convert_vllm_to_comfy.py index 85c7c29..6aec27c 100644 --- a/scripts/convert_vllm_to_comfy.py +++ b/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, } diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index f7866df..f268044 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/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 = ({ jobConfig, setJobConfig, gpuIDs, se /> )} + {additionalSections.includes('caption.caption_format') && ( +
+ setJobConfig(value, 'config.process[0].caption.caption_format')} + options={captionFormatOptions} + /> +
+ )} {additionalSections.includes('caption.fixed_caption') && (