from typing import Optional try: import librosa except ImportError: librosa = None import threading import numpy as np import torch import torchaudio from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor from transformers.generation import LogitsProcessor, LogitsProcessorList from collections import OrderedDict import traceback from optimum.quanto import freeze from toolkit.basic import flush from toolkit.util.quantize import quantize, get_qtype from .BaseCaptioner import BaseCaptioner, CaptionConfig import transformers import logging import warnings # transformers.logging.set_verbosity_error() warnings.filterwarnings("ignore") logging.disable(logging.WARNING) 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 # lyric transcription loops on repeated musical phrases more than captioning TRANSCRIBER_REPETITION_PENALTY = 1.05 # whisper-style temperature fallback when greedy decode degenerates into a # token loop; the penalty alone cannot break a confident single-token loop TRANSCRIBER_RETRIES = ( {"do_sample": True, "temperature": 0.4, "repetition_penalty": 1.20}, {"do_sample": True, "temperature": 0.7, "repetition_penalty": 1.30}, {"do_sample": True, "temperature": 1.0, "repetition_penalty": 1.50}, ) # a loop is a unit of up to LOOP_MAX_PERIOD tokens repeated at least # LOOP_REPEATS times over a span of at least LOOP_MIN_SPAN tokens; sampled # decode jitters a token now and then, so the match is a fraction, not exact. # real outros repeat a line up to ~12 times (~128 tokens); loops never stop LOOP_MAX_PERIOD = 32 LOOP_REPEATS = 4 LOOP_MIN_SPAN = 160 LOOP_MATCH_FRACTION = 0.9 # on a detected loop the unit's tokens are banned for max(2*period, this) steps LOOP_BAN_STEPS = 8 # loops broken before decode is forced to stop LOOP_BREAKS_BEFORE_STOP = 2 # repeats of the looping unit kept when truncating LOOP_KEEP_REPEATS = 2 CAPTIONER_REPETITION_PENALTY = 1.05 # 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 # silence trim on the extracted vocal stem: frames below this fraction of the # peak frame RMS (-40 dB) at the head and tail are cut, keeping a short margin SILENCE_FRAME_SECONDS = 0.02 SILENCE_RMS_FRACTION = 0.01 SILENCE_MARGIN_SECONDS = 0.25 # 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] ) MINOR_PROFILE = np.array( [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17] ) KEY_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"] # ═══════════════════════════════════════════════════════════════════════════════ # Audio analysis (BPM, key, time signature) via librosa # ═══════════════════════════════════════════════════════════════════════════════ def analyze_audio(audio_path): """Extract BPM, key, and time signature from audio using librosa.""" if librosa is None: raise ImportError( "librosa is required for the AceStep captioner but is not " "installed (no numba/llvmlite wheels for this platform yet)." ) y, sr = librosa.load(audio_path, sr=22050, mono=True) duration = librosa.get_duration(y=y, sr=sr) # BPM tempo, _ = librosa.beat.beat_track(y=y, sr=sr) if hasattr(tempo, "__len__"): tempo = tempo[0] bpm = int(round(float(tempo))) # Key detection via chroma correlation with key profiles chroma = librosa.feature.chroma_cqt(y=y, sr=sr) chroma_avg = chroma.mean(axis=1) major_corrs = np.array( [np.corrcoef(np.roll(MAJOR_PROFILE, i), chroma_avg)[0, 1] for i in range(12)] ) minor_corrs = np.array( [np.corrcoef(np.roll(MINOR_PROFILE, i), chroma_avg)[0, 1] for i in range(12)] ) best_major_idx = major_corrs.argmax() best_minor_idx = minor_corrs.argmax() if major_corrs[best_major_idx] >= minor_corrs[best_minor_idx]: keyscale = f"{KEY_NAMES[best_major_idx]} major" else: keyscale = f"{KEY_NAMES[best_minor_idx]} minor" # Time signature estimation from beat strength pattern onset_env = librosa.onset.onset_strength(y=y, sr=sr) tempo_est, beats = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr) if len(beats) >= 8: beat_strengths = onset_env[beats] # Check 3/4 vs 4/4 by looking at periodicity of strong beats acf = np.correlate( beat_strengths - beat_strengths.mean(), beat_strengths - beat_strengths.mean(), mode="full", ) acf = acf[len(acf) // 2 :] if len(acf) > 6: # Look at autocorrelation peaks at lag 3 vs lag 4 score_3 = acf[3] if len(acf) > 3 else 0 score_4 = acf[4] if len(acf) > 4 else 0 timesig = "3" if score_3 > score_4 * 1.2 else "4" else: timesig = "4" else: timesig = "4" return { "bpm": bpm, "keyscale": keyscale, "timesignature": timesig, "duration": int(round(duration)), } from extensions_built_in.llm_models.src.thinker import ( # noqa: E402 shared fast paths OstrisQwen25OmniAudioEncoder, OstrisQwen25OmniThinker, ostris_sdpa_attention_forward, prepare_thinker, ) # 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") def find_loop_period(tokens) -> int: """Period of a unit repeated across the tail of the token ids, 0 if none.""" arr = np.asarray(tokens) for period in range(1, LOOP_MAX_PERIOD + 1): span = max(LOOP_REPEATS * period, LOOP_MIN_SPAN) if span > arr.shape[0]: continue tail = arr[-span:] if np.mean(tail[period:] == tail[:-period]) >= LOOP_MATCH_FRACTION: return period return 0 def find_loop_start(tokens, period: int) -> int: """Index where the periodic run at the tail of the token ids begins.""" # a mismatch run longer than this is real content, shorter is jitter tolerance = max(period, 4) start = len(tokens) - 1 mismatch_run = 0 while start - period >= 0: if tokens[start] == tokens[start - period]: mismatch_run = 0 else: mismatch_run += 1 if mismatch_run > tolerance: start += mismatch_run # back out of the real content break start -= 1 return start + 1 class LoopBreaker(LogitsProcessor): """Breaks token loops inside a single decode: bans the repeating unit's tokens for a stretch so the model has to move on, and forces EOS once it has looped LOOP_BREAKS_BEFORE_STOP times. Batch size 1 only.""" def __init__(self, prompt_len: int, eos_ids): self.prompt_len = prompt_len self.eos_ids = list(eos_ids) self.generated: list = [] self.banned: list = [] self.ban_until = 0 self.window_start = 0 # detection only looks at tokens after the last break self.breaks = 0 self.stopped = False # first loop: where it began, its period, and where decode resumed after the ban self.loop_start = 0 self.loop_period = 0 self.loop_resume = 0 def __call__(self, input_ids, scores): n = input_ids.shape[1] - self.prompt_len if n > len(self.generated): self.generated.extend(input_ids[0, self.prompt_len + len(self.generated):].tolist()) if self.stopped: return self._force_eos(scores) if n < self.ban_until: scores[:, self.banned] = float("-inf") return scores period = find_loop_period(self.generated[self.window_start:]) if not period: return scores self.breaks += 1 if self.breaks >= LOOP_BREAKS_BEFORE_STOP: self.stopped = True return self._force_eos(scores) self.banned = sorted(set(self.generated[-period:])) self.ban_until = n + max(2 * period, LOOP_BAN_STEPS) self.window_start = self.ban_until if self.breaks == 1: self.loop_start = find_loop_start(self.generated, period) self.loop_period = period self.loop_resume = self.ban_until scores[:, self.banned] = float("-inf") return scores def _force_eos(self, scores): scores[:] = float("-inf") scores[:, self.eos_ids] = 0.0 return scores 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) # transcribe a MelBandRoformer vocal stem instead of the full mix; the # captioner still hears the full mix self.extract_vocals_before_transcribe: bool = bool( kwargs.get("extract_vocals_before_transcribe", False) ) 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): caption_config_class = AceStepCaptionConfig caption_config: AceStepCaptionConfig def __init__(self, process_id: int, job, config: OrderedDict, **kwargs): super(AceStepCaptioner, self).__init__(process_id, job, config, **kwargs) self.separator = None # prep runs in worker threads; one separation on the GPU at a time self.separator_lock = threading.Lock() 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", ) model = prepare_thinker(full) model.to(self.device_torch) if self.caption_config.quantize: self.print_and_status_update(f"Quantizing {label} model") quantize(model, weights=get_qtype(self.caption_config.qtype)) freeze(model) flush() processor = Qwen2_5OmniProcessor.from_pretrained(name_or_path) if self.caption_config.low_vram: 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 None: self.model2, self.processor2 = self._load_thinker( self.caption_config.model_name_or_path2, "captioner" ) if self.caption_config.extract_vocals_before_transcribe: from toolkit.audio.melbandroformer import load_melbandroformer self.print_and_status_update("Loading vocal separator") self.separator = load_melbandroformer(device=self.device_torch) flush() 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", "content": [ {"type": "audio", "audio": "<|audio_bos|><|AUDIO|><|audio_eos|>"}, {"type": "text", "text": prompt_text}, ], } ] text = processor.apply_chat_template( conversation, add_generation_prompt=True, tokenize=False ) return processor( text=text, audio=[audio_data], images=None, videos=None, return_tensors="pt", padding=True, sampling_rate=TARGET_SAMPLE_RATE, ) @staticmethod def _to_model_audio(waveform: torch.Tensor, sr: int) -> np.ndarray: """[C, T] at any rate -> mono float array at TARGET_SAMPLE_RATE.""" 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) return waveform.squeeze(0).numpy() @staticmethod def _load_audio(file_path: str): waveform, sr = torchaudio.load(file_path) return AceStepCaptioner._to_model_audio(waveform, sr) @staticmethod def trim_silence(audio: np.ndarray, sample_rate: int = TARGET_SAMPLE_RATE) -> np.ndarray: """Cut leading and trailing silence, judged by frame RMS against the loudest frame.""" frame = max(1, int(sample_rate * SILENCE_FRAME_SECONDS)) n_frames = len(audio) // frame if n_frames < 2: return audio rms = np.sqrt(np.mean(audio[: n_frames * frame].reshape(n_frames, frame) ** 2, axis=1)) loud = np.flatnonzero(rms >= rms.max() * SILENCE_RMS_FRACTION) if len(loud) == 0: return audio margin = int(sample_rate * SILENCE_MARGIN_SECONDS) start = max(0, loud[0] * frame - margin) end = min(len(audio), (loud[-1] + 1) * frame + margin) return audio[start:end] def _extract_vocals(self, waveform: torch.Tensor, sr: int) -> np.ndarray: """Vocal stem of [C, T] audio as mono TARGET_SAMPLE_RATE, silence-trimmed.""" from toolkit.audio.melbandroformer import separate with self.separator_lock: vocals, _ = separate(self.separator, waveform, sr) return self.trim_silence(self._to_model_audio(vocals.cpu(), sr)) 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. Vocal separation is the one GPU step here, serialized by separator_lock.""" waveform, sr = torchaudio.load(file_path) audio_data = self._to_model_audio(waveform, sr) lyrics_audio = ( self._extract_vocals(waveform, sr) if self.separator is not None else audio_data ) item = { "file": file_path, # 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, lyrics_audio, LYRICS_PROMPT), } if self.model2 is not None: item["caption_inputs"] = self._make_inputs( self.processor2, audio_data, CAPTION_PROMPT ) return item @staticmethod def _place_inputs(model, inputs): """Move a processor batch onto the model; the thinker also reads its 2D pad mask from the model. Called again before every decode pass.""" inputs = inputs.to(model.device).to(model.dtype) model._pad_mask_2d = inputs.get("attention_mask", None) return inputs def _generate( self, model, processor, inputs, repetition_penalty: float, retries: tuple = (), break_loops: bool = False, file_path: str = "", max_new_tokens: int = MAX_NEW_TOKENS, ) -> str: inputs = self._place_inputs(model, inputs) # 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 input_len = inputs["input_ids"].shape[1] # greedy decode loops on repeated phrases without a penalty gen_kwargs = {"max_new_tokens": max_new_tokens, "repetition_penalty": repetition_penalty} if model.generation_config.cache_implementation == "static": if input_len + 16 < STATIC_MAX_LENGTH: from transformers.generation import MaxLengthCriteria, StoppingCriteriaList gen_kwargs = { "repetition_penalty": repetition_penalty, "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" eos_ids = model.generation_config.eos_token_id eos_ids = set(eos_ids if isinstance(eos_ids, (list, tuple)) else [eos_ids]) def strip_eos(tokens: torch.Tensor) -> torch.Tensor: end = tokens.shape[0] while end > 0 and int(tokens[end - 1]) in eos_ids: end -= 1 return tokens[:end] def run_pass(): breaker = LoopBreaker(input_len, eos_ids) if break_loops else None kwargs = dict(gen_kwargs) if breaker is not None: kwargs["logits_processor"] = LogitsProcessorList([breaker]) self._place_inputs(model, inputs) generated = model.generate(**inputs, **kwargs) return strip_eos(generated[0, input_len:]), breaker new_tokens, breaker = run_pass() first_tokens, first_breaker = new_tokens, breaker stopped = breaker is not None and breaker.stopped for attempt, retry in enumerate(retries, 1): if not stopped: break print( f"\n[AITK] {file_path}: transcription looped {breaker.breaks} times and was stopped, " f"retry {attempt}/{len(retries)} with temperature={retry['temperature']} " f"repetition_penalty={retry['repetition_penalty']}" ) gen_kwargs.update(retry) new_tokens, breaker = run_pass() stopped = breaker.stopped if stopped: # every pass re-looped after a break; keep the greedy pass up to its first loop keep = first_breaker.loop_start + LOOP_KEEP_REPEATS * first_breaker.loop_period new_tokens = first_tokens[:keep] print( f"\n[AITK] {file_path}: transcription looped on every pass; keeping the greedy " f"pass truncated at its first loop (period {first_breaker.loop_period} tokens)" ) elif breaker is not None and breaker.breaks: # drop the excess repeats and the tokens forced out under the ban keep = breaker.loop_start + LOOP_KEEP_REPEATS * breaker.loop_period new_tokens = torch.cat([new_tokens[:keep], new_tokens[breaker.loop_resume :]]) print( f"\n[AITK] {file_path}: broke a transcription loop " f"(period {breaker.loop_period} tokens) and continued" ) output = processor.batch_decode( new_tokens[None], skip_special_tokens=True, clean_up_tokenization_spaces=False, )[0] return output.strip() def get_audio_lyrics(self, inputs, file_path: str = "") -> str: if ( self.caption_config.low_vram and self.model2 is not None and self.model2.device != torch.device("cpu") ): self.model2.to("cpu") if self.model.device == torch.device("cpu"): self.model.to(self.device_torch) return self._generate( self.model, self.processor, inputs, TRANSCRIBER_REPETITION_PENALTY, retries=TRANSCRIBER_RETRIES, break_loops=True, file_path=file_path, ) def get_audio_caption(self, inputs, file_path: str = "") -> str: if self.caption_config.low_vram and self.model.device != torch.device("cpu"): self.model.to("cpu") if self.model2.device == torch.device("cpu"): self.model2.to(self.device_torch) return self._generate(self.model2, self.processor2, inputs, CAPTIONER_REPETITION_PENALTY) def _transcribe_item(self, item: dict) -> str: """Raw transcriber text for a prepped item: '# Languages ... # Lyrics ...'.""" return self.get_audio_lyrics(item["lyrics_inputs"], item["file"]) def _caption_item(self, item: dict) -> str: if self.caption_config.caption_format == "yue2": lyrics = clean_lyrics(self._transcribe_item(item)) if self.caption_config.fixed_caption is not None: caption = self.caption_config.fixed_caption else: caption = self.get_audio_caption(item["caption_inputs"], item["file"]) caption = " ".join(caption.split()) return f"{caption}\n[Lyrics]\n{lyrics}" analysis = item["analysis"] lyrics = self._transcribe_item(item) 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"], item["file"]) 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: 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