From 213d32cad3d6d3c89af2d088a92abdc7a7f11485 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 14 Sep 2026 17:50:15 -0600 Subject: [PATCH] Train 50% of the time with and without sheet music. --- .../audio_models/yue2/yue2_model.py | 68 +++++++++++++------ ui/src/app/jobs/new/options.tsx | 26 +++++-- 2 files changed, 68 insertions(+), 26 deletions(-) diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py index 14b3f36..bdad93b 100644 --- a/extensions_built_in/audio_models/yue2/yue2_model.py +++ b/extensions_built_in/audio_models/yue2/yue2_model.py @@ -195,6 +195,9 @@ class YuE2AudioModel(BaseAudioModel): self.abc_dropout = float(kw.get("abc_dropout", 0.5)) self.sheetsage_path = kw.get("sheetsage_path", SHEETSAGE_FILE) self.transcriber: Optional[SheetSage2Transcriber] = None + # encoders are only needed while caching; once training steps run without encode calls they are released + self._train_calls = 0 + self._last_encode_call = 0 self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0)) self.sample_ar_temperature = float(kw.get("sample_ar_temperature", 1.0)) # reference sampler value; 1.0 lets a memorized song replay (the penalty knocks it off path) @@ -292,6 +295,7 @@ class YuE2AudioModel(BaseAudioModel): "off", the SheetSage2 sheet as ``abc_ids`` [B, N] (text ids, -1 padded).""" if device is None: device = self.vae_device_torch + self._last_encode_call = self._train_calls if self.vae.device == torch.device("cpu"): self.vae.to(device) latents = self.vae.encode(audio_tensor.to(device=device, dtype=self.vae.dtype)) # [B, 64, T] @@ -386,35 +390,52 @@ class YuE2AudioModel(BaseAudioModel): net.is_active = was_active want_info = self.debug and self._loss_log_every > 0 and (self._loss_log_step + 1) % self._loss_log_every == 0 lm_head = model.ar.model.lm_head - ce_sum = hidden.new_zeros((), dtype=torch.float32) - kl_sum = hidden.new_zeros((), dtype=torch.float32) if base_hidden is not None else None - p_target, p_end_first = [], None - chunk = 2048 - for s0 in range(0, n, chunk): - h = hidden[s0 : s0 + chunk] - tgt = ids[s0 : s0 + chunk] - if torch.is_grad_enabled(): - logits = torch.utils.checkpoint.checkpoint(lambda x: lm_head(x).float(), h, use_reentrant=False) - else: - logits = lm_head(h).float() - ce_sum = ce_sum + torch.nn.functional.cross_entropy(logits, tgt, reduction="sum") - if base_hidden is not None: + want_kl = base_hidden is not None + + def chunk_losses(h, tgt, bh): + # everything vocab-sized lives inside this checkpointed function, so autograd keeps only + # the [chunk, hidden] inputs and recomputes the fp32 logits/log-probs in backward + logits = lm_head(h).float() + ce = torch.nn.functional.cross_entropy(logits, tgt, reduction="sum") + if bh is not None: with torch.no_grad(): - base_logp = torch.log_softmax(lm_head(base_hidden[s0 : s0 + chunk]).float(), -1) - kl_sum = kl_sum + torch.nn.functional.kl_div(torch.log_softmax(logits, -1), base_logp, log_target=True, reduction="sum") + base_logp = torch.log_softmax(lm_head(bh).float(), -1) + kl = torch.nn.functional.kl_div(torch.log_softmax(logits, -1), base_logp, log_target=True, reduction="sum") + else: + kl = ce.new_zeros(()) if want_info: with torch.no_grad(): p = logits.softmax(-1) - if s0 == 0: - p_end_first = p[0, MUSIC_END].item() - p_target.append(p.gather(1, tgt[:, None]).squeeze(1)) - del logits + p_first_end = p[0, MUSIC_END] + p_tgt = p.gather(1, tgt[:, None]).squeeze(1) + else: + p_first_end, p_tgt = ce.new_zeros(()), ce.new_zeros((0,)) + return ce, kl, p_first_end, p_tgt + + ce_sum = hidden.new_zeros((), dtype=torch.float32) + kl_sum = hidden.new_zeros((), dtype=torch.float32) + p_target, p_end_first = [], None + chunk = 512 # vocab-sized fp32 transients scale with this (~1.9 GB per 512 positions incl. KL) + for s0 in range(0, n, chunk): + h = hidden[s0 : s0 + chunk] + tgt = ids[s0 : s0 + chunk] + bh = base_hidden[s0 : s0 + chunk] if want_kl else None + if torch.is_grad_enabled(): + ce, kl, p_first, p_tgt = torch.utils.checkpoint.checkpoint(chunk_losses, h, tgt, bh, use_reentrant=False) + else: + ce, kl, p_first, p_tgt = chunk_losses(h, tgt, bh) + ce_sum = ce_sum + ce + kl_sum = kl_sum + kl + if want_info: + if s0 == 0: + p_end_first = p_first.item() + p_target.append(p_tgt) info = "" if want_info: p_target = torch.cat(p_target) info = (f"[ar targets {n} unique {ids.unique().numel()} (song {total}, whole {whole_song}) " f"prefix {prefix.shape[0]} p(END|prefix) {p_end_first:.3f} p(target) median {p_target.median().item():.3f}]") - return ce_sum / n, (kl_sum / n if kl_sum is not None else None), info + return ce_sum / n, (kl_sum / n if want_kl else None), info def _ar_inputs(self, prefix: torch.Tensor, abc_ids: torch.Tensor, tokens: torch.Tensor, end_token: bool = True): """One item: prefix head (unpadded [L, H]) + abc + [ABC_END, MUSIC_START] + codec tokens [+ MUSIC_END] @@ -494,6 +515,13 @@ class YuE2AudioModel(BaseAudioModel): prefix_embeds = text_embeddings.text_embeds.to(device, self.torch_dtype) prefix_mask = text_embeddings.attention_mask.to(device) train_ar = self.ar_loss_weight > 0 and torch.is_grad_enabled() + if torch.is_grad_enabled(): + self._train_calls += 1 + if self._train_calls - self._last_encode_call > 2 and (self.semantic_tokenizer is not None or self.transcriber is not None): + # latents are cached: MERT + SheetSage2 (~4 GB) are dead weight for the rest of training + self.semantic_tokenizer = None + self.transcriber = None + flush() captions = batch.get_caption_list() if hasattr(batch, "get_caption_list") else None preds, lm_losses, kl_losses = [], [], [] for i in range(latent_model_input.shape[0]): diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index 5a97ddc..d7cd351 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/ui/src/app/jobs/new/options.tsx @@ -1398,6 +1398,8 @@ export const modelArchs: ModelArch[] = [ '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 }, {}], @@ -1449,12 +1451,24 @@ export const modelArchs: ModelArch[] = [ close.

- Training matches Comfy's Generate ABC path (cot: full in model kwargs): at cache time every song is - transcribed to an ABC lead sheet with SheetSage2 (the Comfy-Org/YuE2 audio_encoders weights, no - extra packages), the sheet is stored with the latents, and the AR learns - lyrics to sheet and sheet to tokens. abc_dropout (default 0.5) is the fraction of training steps fed - without the sheet as an off-mode prompt, so the same LoRA also works when Comfy's ABC input is empty. cot: melody uses sheets without chords; cot: off skips the sheet, which is what - Comfy runs when its ABC input is left empty. Changing the mode requires re-caching the dataset. + 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