Train 50% of the time with and without sheet music.

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:
+        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_sum = ce_sum + torch.nn.functional.cross_entropy(logits, tgt, reduction="sum")
-            if base_hidden is not None:
+            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)
+                    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[0, MUSIC_END].item()
-                    p_target.append(p.gather(1, tgt[:, None]).squeeze(1))
-            del logits
+                    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]):

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.
         </p>
         <p>
-          Training matches Comfy's Generate ABC path (<code>cot: full</code> in model kwargs): at cache time every song is
-          transcribed to an ABC lead sheet with SheetSage2 (the Comfy-Org/YuE2 <code>audio_encoders</code> weights, no
-          extra packages), the sheet is stored with the latents, and the AR learns
-          lyrics to sheet and sheet to tokens. <code>abc_dropout</code> (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. <code>cot: melody</code> uses sheets without chords; <code>cot: off</code> skips the sheet, which is what
-          Comfy runs when its ABC input is left empty. Changing the mode requires re-caching the dataset.
+          <b>ABC generation.</b> 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 <code>cot: full</code> (chords) or <code>cot: melody</code> (melody only), or the single
+          stage with <code>cot: off</code>.
+        </p>
+        <p>
+          <b>Training for both.</b> Every song is transcribed to an ABC sheet with SheetSage2 and the AR is trained on
+          lyrics to sheet and sheet to tokens. <code>abc_dropout</code> (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.
+        </p>
+        <p>
+          <b>Caching is required.</b> 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{' '}
+          <code>cot</code> means deleting the dataset's <code>_latent_cache</code> folder so the sheets are rebuilt.
+          Audio has no resolution, so keep a single resolution bucket per dataset or every bucket duplicates the songs.
         </p>
         <p>
           Prompt format: a style line, then <code>[Lyrics]</code>, then the lyrics with bracketed section headers such