YuE2 LoRA Trainer

What this node is for

Music generation in ComfyUI has one obvious trainable option - ACE-Step, which shipped LoRA training with v1.5 and runs on a 4GB card - and one stubbornly interesting one: YuE2, the 3B successor to the "Suno at home" model that got the whole idea of local song generation taken seriously. ACE-Step is the polished, instrumentals-are-great option. YuE2 is the one people reach for when they want vocals and song structure. Neither had a ComfyUI trainer until this pack.

YuE2LoRATrainer is that trainer. Feed it a dataset built from a folder of your own tracks, give it a trigger word like mystyle, queue the prompt, and it writes <lora_name>.safetensors into models/loras in native ComfyUI format - loaded with the stock LoraLoaderModelOnly on the native YuE2 checkpoint. No conversion step, no A1111 detour.

The honest framing: this is experimental software from a small utility-node author (Starnodes, of helper-nodes and Image Manager). Treat your first run as a test, not a production job.

The mechanism, in one paragraph

YuE2 has three parts: a 2.2B AR language model that plans the song and writes semantic tokens, a 1.5B NAR flow-matching branch that renders 64-channel VAE latents into audio, and a 48 kHz stereo VAE. This trainer adapts the NAR branch only. The composer stays frozen - m-a-p never released an audio→token encoder or training code - so what you get is a style/instrumentation/timbre LoRA, not a voice clone. The README is blunt that voice cloning still doesn't work. Training uses the released flow-matching objective, conditioned on the checkpoint's own text prefix (trigger_word, caption).

Inputs you actually touch

Required, all of them - there's nothing optional:

• dataset - the YUE2_TRAIN_DATASET from YuE2TrainingDataset. This is the only wire in.
• checkpoint - a native YuE2 all-in-one file from models/checkpoints. Pick the same file you generate with. Non-YuE2 files and the INT8/quantized variant are rejected outright.
• trigger_word - default mystyle. It goes at the start of your style prompt at generation time.
• steps (3000) and learning_rate (1e-4) - the two dials that decide whether the LoRA does anything. Rank 32 / alpha 32 is the tuned default and you can leave it.
• lora_name - writes models/loras/<lora_name>.safetensors, plus a <lora_name>_raw.safetensors sibling when EMA is on.

The rest have defaults worth not fiddling with: lr_scheduler = cosine, ema_decay = 0.999, t_sampling = logit_normal, max_grad_norm = 1.0, warmup_steps = 50, caption_dropout = 0.1, grad_accum = 1, target_preset = nar_attn_mlp. Two exceptions: flip optimizer to adamw_8bit if you're tight on VRAM, and set save_every if you want intermediate checkpoints to A/B. And leave EMA on - the main file gets the smoothed weights, the _raw sibling gets the noisy ones, and that pairing is what removes most run-to-run variance.

Outputs: lora_path (STRING - the file you just wrote) and training_log, which is the wire into YuE2TrainingCurve.

Installing

ComfyUI Manager, search ComfyUI-YuE2-Trainer, or:

cd ComfyUI/custom_nodes
git clone https://github.com/Starnodes2024/ComfyUI-YuE2-Trainer.git
python_embeded/python.exe -m pip install -r ComfyUI-YuE2-Trainer/requirements.txt

Two small deps and nothing exotic: soundfile (audio fallback when torchaudio can't decode a file) and matplotlib (curve node only). The official m-a-p inference code is bundled, so there's no second pack to install.

The model: one file, yue2_3b_bf16.safetensors (~7.8 GB) from Comfy-Org/YuE2, into ComfyUI/models/checkpoints/. One download covers training and generation. Use bf16, not the INT8 repack - quantized weights can't be trained. Then restart ComfyUI; three nodes appear under YuE2/Training.

Where people get burned

• VRAM. 24 GB recommended. The node unloads your other ComfyUI models before training starts - deliberate, one training run owns the GPU. Out of memory? Lower clip_seconds in the dataset node first, then try adamw_8bit.
• Nonsense checkpoint errors. Pointing checkpoint at a GGUF/INT8 file or a non-YuE2 model gives you "not a native YuE2 all-in-one checkpoint" - a key-signature check, not a bad download. An empty trigger_word and path characters in lora_name also raise immediately; those are guards, not bugs.
• Overfitting looks like success. Muffled, repetitive output that weakly copies your training tracks means too many steps or too high an LR - drop to 5e-5 or cut steps. If the trigger has no audible effect at all, go the other way (2e-4).
• Speed expectations. On a 4090, 10-second clips run ~1–3 s/step, so 1000 steps is 25–50 minutes. A 3000-step run is an evening.

And the one nobody reads: YuE2 weights are CC BY-NC 4.0. Your LoRA is a derivative, and so is anything you generate with it. Non-commercial, attribution required.

Is it the one to reach for?

If you want local, trainable music with a 4 GB floor, ACE-Step is still the answer. Reach for YuE2LoRATrainer when you specifically want YuE2's vocal-forward songwriting and you're willing to babysit an experimental trainer with a reconstructed objective, no caption files, and a 24 GB appetite. Set the trigger word at the start of the style prompt and keep cot = off when you generate - that matches the text-only regime the LoRA was trained in.