diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py index 470e448..07280e7 100644 --- a/extensions_built_in/audio_models/yue2/yue2_model.py +++ b/extensions_built_in/audio_models/yue2/yue2_model.py @@ -4,9 +4,8 @@ Weights: the Comfy-Org all-in-one checkpoint (``checkpoints/yue2_3b_bf16.safeten resolved under MODELS_PATH in ComfyUI's folder layout and downloaded there when missing. The official audio -> semantic-token tokenizer is unreleased; training conditioning comes from the community head (MERT-v2-FullSong + classifier) in -``Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4``. Its companion NAR adapter is -merged into the base NAR on load (``model_kwargs.merge_nar_lora``, default true) since the -head's tokens are a dialect the stock renderer was not trained on. +``Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4``. Its companion NAR adapter can +be merged into the base NAR on load (``model_kwargs.merge_nar_lora``, default false). Training: per song the latent cache carries VAE latents plus the head's codec tokens; the "text embedding" is the AR prompt prefix (instruction + tags + lyrics) as token embeddings. @@ -144,7 +143,7 @@ class YuE2AudioModel(BaseAudioModel): self.ar_lr_multiplier = float(kw.get("ar_lr_multiplier", 1.0)) self._loss_log_every = int(kw.get("loss_log_every", 25)) self._loss_log_step = 0 - self.merge_nar_lora = bool(kw.get("merge_nar_lora", True)) + self.merge_nar_lora = bool(kw.get("merge_nar_lora", False)) self.nar_lora_path = kw.get("nar_lora_path", f"{HEAD_REPO}/{NAR_LORA_FILE}") self.semantic_head_path = kw.get("semantic_head_path", f"{HEAD_REPO}/{HEAD_FILE}") self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0)) diff --git a/extensions_built_in/captioner/Qwen3OmniCaptioner.py b/extensions_built_in/captioner/Qwen3OmniCaptioner.py index 38f2ce5..c4ad548 100644 --- a/extensions_built_in/captioner/Qwen3OmniCaptioner.py +++ b/extensions_built_in/captioner/Qwen3OmniCaptioner.py @@ -29,6 +29,8 @@ VIDEO_FPS = 2 # still-image files caption through the image pipeline (no audio, no frames) IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"} +# audio-only files caption through the audio pipeline (no frames) +AUDIO_EXTENSIONS = {"mp3", "wav", "flac", "ogg", "m4a", "aac"} # fixed generation ceiling under compiled decode: a constant max_length keeps # the static kv cache (and so the compiled decode graph) at one shape for @@ -660,9 +662,15 @@ class Qwen3OmniCaptioner(BaseCaptioner): def _is_image_file(file_path: str) -> bool: return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS + @staticmethod + def _is_audio_file(file_path: str) -> bool: + return os.path.splitext(file_path)[1].lower().lstrip(".") in AUDIO_EXTENSIONS + def _build_messages(self, _file_path: str): if self._is_image_file(_file_path): media = {"type": "image", "image": _file_path} + elif self._is_audio_file(_file_path): + media = {"type": "audio", "audio": _file_path} else: media = {"type": "video", "video": _file_path} return [ @@ -695,6 +703,10 @@ class Qwen3OmniCaptioner(BaseCaptioner): image = Image.open(file_path).convert("RGB") item = {"file": file_path, "kind": "image", "image": image, "audio": None} + elif self._is_audio_file(file_path): + from transformers.audio_utils import load_audio + + item = {"file": file_path, "kind": "audio", "audio": load_audio(file_path, sampling_rate=16000)} else: from transformers.video_utils import load_video from transformers.audio_utils import load_audio @@ -738,6 +750,13 @@ class Qwen3OmniCaptioner(BaseCaptioner): padding=True, size=self._size_kwargs(), ) + if kind == "audio": + return self.processor( + text=[it["text"] for it in items], + audio=[it["audio"] for it in items], + return_tensors="pt", + padding=True, + ) use_audio = kind == "video_audio" return self.processor( text=[it["text"] for it in items], @@ -770,10 +789,14 @@ class Qwen3OmniCaptioner(BaseCaptioner): # under static cache, generate hands the forward a prepared 4D mask; # the true 2D padding mask is needed for the prefill rope index self.model._pad_mask_2d = inputs.get("attention_mask", None) + gen_kwargs = self._gen_kwargs(inputs["input_ids"].shape[1]) + if items[0]["kind"] == "audio": + # greedy lyric transcription loops on syllables; mild penalty (the model's own talker default) + gen_kwargs["repetition_penalty"] = 1.05 generated_ids = self.model.generate( **inputs, use_audio_in_video=use_audio, - **self._gen_kwargs(inputs["input_ids"].shape[1]), + **gen_kwargs, ) trimmed = generated_ids[:, inputs["input_ids"].shape[1] :] captions = self.processor.batch_decode( @@ -876,7 +899,7 @@ class Qwen3OmniCaptioner(BaseCaptioner): futures.append((path, executor.submit(self._prep_media, path))) # batches must be homogeneous: the processor call differs per kind - buckets = {"image": [], "video_audio": [], "video_silent": []} + buckets = {"image": [], "audio": [], "video_audio": [], "video_silent": []} while futures: if self.is_ui_captioner: self.maybe_stop() diff --git a/ui/src/helpers/captionJobConfig.ts b/ui/src/helpers/captionJobConfig.ts index 2c2a3af..3cd180a 100644 --- a/ui/src/helpers/captionJobConfig.ts +++ b/ui/src/helpers/captionJobConfig.ts @@ -17,7 +17,7 @@ export const defaultCaptionJobConfig: CaptionJobConfig = { quantize: true, qtype: 'convrot8', low_vram: false, - extensions: ['mp4', 'mov', 'webm', 'mkv', 'avi', 'jpg', 'jpeg', 'png', 'bmp', 'webp'], + extensions: ['mp4', 'mov', 'webm', 'mkv', 'avi', 'jpg', 'jpeg', 'png', 'bmp', 'webp', 'mp3', 'wav', 'flac', 'ogg'], path_to_caption: '', recaption: false, compile: true, diff --git a/ui/src/helpers/captionOptions.ts b/ui/src/helpers/captionOptions.ts index a2141bf..a61493e 100644 --- a/ui/src/helpers/captionOptions.ts +++ b/ui/src/helpers/captionOptions.ts @@ -67,6 +67,17 @@ overall_soundscape and non_diegetic_music are always exactly N/A for a still ima Describe only what is actually visible. Be decisive. No preamble and no extra text - output only the three fields.`; +// Captions a song as a YuE2 training caption: one line of comma-separated style +// tags, then a [Lyrics] line and the verbatim lyrics with bracketed section +// headers - the exact prefix layout the yue2 arch's parse_caption expects. +const yue2CaptionPrompt = `Listen to this song and write a YuE2 training caption for it. Output exactly two parts and nothing else. + +Part 1, the first line only: comma-separated style tags describing the music. Cover, in this order where applicable: genre and sub-genre, mood, vocal type (male vocal, female vocal, duet, choir, rap, or instrumental), vocal delivery (breathy, belted, falsetto, spoken, whispered, harmonized), lead instruments and production (acoustic guitar, piano, synth pads, 808 bass, live drums, drum machine, strings, lo-fi, reverb-heavy), tempo feel (slow ballad, mid-tempo, upbeat, fast), and era or scene (90s alt rock, modern trap, bedroom pop). Use concrete lowercase tags, no sentences, no hedging, no artist or song names. Example: alternative rock, melancholic, male vocal, falsetto, electric guitar, piano, live drums, slow ballad, 2000s + +Part 2: on the next line write [Lyrics] and then the complete lyrics transcribed verbatim, one sung line per line. Split the song into sections with a bracketed header on its own line before each one, using these names: [Intro], [Verse 1], [Verse 2], [Pre-Chorus], [Chorus], [Bridge], [Instrumental], [Solo], [Outro]. A short descriptor may follow the name inside the brackets, e.g. [Intro choir] or [Chorus harmonized]. Leave one blank line between sections. Repeat a chorus each time it is sung, but write each repeated line once per time it is sung, never more. Wordless vocalizing (na, la, oh, ah, hums, vocal chops) is never transcribed syllable by syllable: describe it once in the section header instead, e.g. [Intro wordless vocals] or [Bridge oohs], and if you cannot make out real words in a passage, treat it as wordless. Do not add timestamps, speaker labels, quotation marks, translations, or commentary, and do not annotate ad-libs beyond the section header. If the song has no vocals at all, Part 2 is [Lyrics] followed by a single [Instrumental] line. + +Transcribe only what is actually sung. Be decisive. No preamble, no explanations, no markdown - output only the tag line and the lyrics block.`; + // Editable ADDITIONAL INSTRUCTIONS block injected into the Ideogram system prompt. // Users can tweak this for dataset-specific guidance without altering the fixed // output contract, element/background rules, or bbox format. @@ -127,7 +138,7 @@ export const captionerTypes: CaptionOption[] = [ group: 'image/video/sound', defaults: { 'config.process[0].caption.model_name_or_path': ['ai-toolkit/Qwen3-Omni-30B-A3B-Thinking', defaultNameOrPath], - 'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage], defaultExtensions], + 'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage, ...extensionsAudio], defaultExtensions], 'config.process[0].caption.caption_prompt': [defaultVideoCaptionPrompt, undefined], 'config.process[0].caption.max_res': [512, undefined], 'config.process[0].caption.max_new_tokens': [512, undefined], @@ -143,6 +154,7 @@ export const captionerTypes: CaptionOption[] = [ 'General': defaultVideoCaptionPrompt, 'MiniMax H4 T2V': minimaxT2VCaptionPrompt, 'MiniMax H4 Image': minimaxImageCaptionPrompt, + 'YuE2': yue2CaptionPrompt, }, additionalSections: [ 'caption.caption_prompt',