// UI entries (training form + Generate page) for the models this package // registers in AI_TOOLKIT_MODELS. Loaded at runtime by the UI, not bundled: // see ui/src/extensions/README.md for the convention and the allowed imports. import type { ModelArch } from "@/app/jobs/new/options"; import { defaultSampleConfig, defaultAudioSampleConfig, defaultYue2SampleConfig, } from "@/helpers/defaultSamples"; const defaultNameOrPath = ""; export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ { name: "ace_step_15_xl", label: "ACE-Step 1.5 XL", group: "audio", defaults: { // default updates when [selected, unselected] in the UI "config.process[0].model.name_or_path": [ "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_xl_base_aio.safetensors", defaultNameOrPath, ], "config.process[0].model.quantize": [true, false], "config.process[0].model.quantize_te": [true, false], "config.process[0].model.low_vram": [true, false], "config.process[0].train.unload_text_encoder": [false, false], "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], "config.process[0].train.timestep_type": ["linear", "sigmoid"], "config.process[0].model.qtype": ["qfloat8", "qfloat8"], "config.process[0].sample": [ defaultAudioSampleConfig, defaultSampleConfig, ], }, sampleTags: { CAPTION: { title: "Audio Prompt", type: "text", full: true, }, LYRICS: { title: "Lyrics", type: "multiline", full: true, }, BPM: { title: "BPM", type: "number", }, KEYSCALE: { title: "Key Scale", type: "text", }, TIMESIGNATURE: { title: "Time Signature", type: "text", }, DURATION: { title: "Duration (sec)", type: "number", }, LANGUAGE: { title: "Language", type: "text", }, }, disableSections: ["network.conv"], additionalSections: [ "sample.multi_ctrl_imgs", "model.low_vram", "model.layer_offloading", ], }, { name: "yue2", label: "YuE2", group: "audio", defaults: { // default updates when [selected, unselected] in the UI "config.process[0].model.name_or_path": [ "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", defaultNameOrPath, ], "config.process[0].model.quantize": [true, false], "config.process[0].model.quantize_te": [false, false], "config.process[0].model.low_vram": [false, false], "config.process[0].train.unload_text_encoder": [false, false], "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], "config.process[0].train.timestep_type": ["sigmoid", "sigmoid"], // the int8 repack ships convrot8 layers; requesting convrot8 keeps them as-is (no requantization) "config.process[0].model.qtype": ["convrot8", "qfloat8"], "config.process[0].sample": [ defaultYue2SampleConfig, defaultSampleConfig, ], "config.process[0].datasets[x].cache_latents_to_disk": [true, true], // audio has no resolution; every bucket would duplicate the whole dataset "config.process[0].datasets[x].resolution": [[512], [512, 768, 1024]], // blank captions break lyric following; the AR must always see the prefix "config.process[0].datasets[x].caption_dropout_rate": [0, 0.05], "config.process[0].model.model_kwargs": [ { cot: "full", abc_dropout: 0.5, sample_ar_repetition_penalty: 1.2, ar_kl_weight: 0.2, }, {}, ], }, // native YuE2 prompt: style text, a [Lyrics] line, the lyrics hasMultiLinePrompts: true, modelNotes: (

Experimental. The AR (composition) model memorizes quickly and does not work well on small datasets. A handful of songs is enough for it to learn the exact token sequence of each song; after that it stops generalizing and free-running samples drift away from the training material. Expect to need a large, varied dataset for the AR side to learn a style rather than the songs themselves.

YuE2 is two experts on one backbone. The AR expert reads the style line and lyrics and writes the song as a sequence of semantic codec tokens (25 per second). The NAR expert then renders those tokens into audio latents with flow matching. Training a LoRA here trains both: next-token loss on the AR over the whole song from its start, flow loss on the NAR over a random window.

The AR loss is the one to watch (loss/ar_ce). It starts near 5 and, on a small dataset, falls toward 0 within a few hundred steps, which is memorization. ar_kl_weight in model kwargs anchors the AR to the base model so it cannot collapse onto the training songs; ar_lr_multiplier and a smaller AR rank slow it further. Style comes mostly from the NAR, lyric following from the AR. Do not use caption dropout: a blank prompt breaks lyric following.

The official audio-to-token encoder is unreleased. Training uses the community tokenizer by Kytra ( @sin_ceriously ), a MERT-v2-FullSong head that maps real audio to YuE2 codec tokens:{" "} Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4 . It is downloaded on first use. Its tokens are an approximation of the model's own, so rendered samples will not be bit-faithful to the training audio even when the AR replays a song exactly. Though it is extremely close.

ABC generation. Generation is two stages: the AR first writes a lead sheet of the whole song in ABC notation (sections, chords, vocal and instrumental melody), then writes the codec tokens conditioned on that sheet. It can also run without a sheet ("off" mode), which uses a different instruction line. Samples here do the two stages with{" "} cot: full (chords) or cot: melody (melody only), or the single stage with cot: off.

Training for both. Every song is transcribed to an ABC sheet with SheetSage2 and the AR is trained on lyrics to sheet and sheet to tokens. abc_dropout (default 0.5) is the fraction of training items fed without the sheet instead, as an off-mode prompt, so one LoRA works in both modes. Set it to 0 to train the sheet path only, or 1 to train off mode only.

Caching is required. The sheet is produced at latent-cache time and stored with the latents and codec tokens, so Cache Latents to Disk must stay on: without it, SheetSage2 (about 12 s per song), the MERT tokenizer and the VAE would run again on every training step. The cache records which mode built it; changing cot means deleting the dataset's _latent_cache folder so the sheets are rebuilt. Audio has no resolution, so keep a single resolution bucket per dataset or every bucket duplicates the songs.

Prompt format: a style line, then [Lyrics], then the lyrics with bracketed section headers such as [Verse 1]{" "} and [Chorus]. The Qwen3-Omni captioner has a YuE2 preset that writes captions in this layout. Sample length is set by the Duration field in the sample section.

), // no separate text encoder: the prompt side is the AR expert, covered by the transformer quantization disableSections: ["network.conv", "model.quantize_te"], additionalSections: ["model.low_vram", "sample.duration"], }, { name: "ace_step_15", label: "ACE-Step 1.5", group: "audio", defaults: { // default updates when [selected, unselected] in the UI "config.process[0].model.name_or_path": [ "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", defaultNameOrPath, ], "config.process[0].model.quantize": [true, false], "config.process[0].model.quantize_te": [true, false], "config.process[0].model.low_vram": [true, false], "config.process[0].train.unload_text_encoder": [false, false], "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], "config.process[0].train.timestep_type": ["linear", "sigmoid"], "config.process[0].model.qtype": ["qfloat8", "qfloat8"], "config.process[0].sample": [ defaultAudioSampleConfig, defaultSampleConfig, ], }, sampleTags: { CAPTION: { title: "Audio Prompt", type: "text", full: true, }, LYRICS: { title: "Lyrics", type: "multiline", full: true, }, BPM: { title: "BPM", type: "number", }, KEYSCALE: { title: "Key Scale", type: "text", }, TIMESIGNATURE: { title: "Time Signature", type: "text", }, DURATION: { title: "Duration (sec)", type: "number", }, LANGUAGE: { title: "Language", type: "text", }, }, disableSections: ["network.conv"], additionalSections: [ "sample.multi_ctrl_imgs", "model.low_vram", "model.layer_offloading", ], }, ];