autoregressive

backend.py

@@ -77,6 +77,16 @@ N_CHUNKS = 24            # fixed sequence length (measures) -- matches training
 # positions the model never trained on. Continuation therefore SLIDES a fixed
 # 24-measure window instead of growing it.
 CONTINUE_MEASURES = 8
+
+# Noise level the already-written measures are held at while the new ones are
+# sampled. 0.0 == hand the past over perfectly clean, which is what makes the
+# re-decoded context bit-match the audio already in the mix. Diffusion Forcing's
+# stabilization result is that a SMALL nonzero level makes long autoregressive
+# rollouts more robust (the model stops trusting its own artifacts), so if
+# repeated Continues harden into a loop or drift, try 0.02-0.1. See
+# real_backend.fm_df_sample.
+CONTINUE_CTX_NOISE = 0.0
+
 # Style embedding width. The DiT is now conditioned on CLAP embeddings
 # (clap.CLAP proj_dim=512), NOT the old 128-d contrast.py vectors. Both stage
 # checkpoints carry style_dim=512 in their model_args; keep this in sync.
@@ -131,26 +141,53 @@ def shift_window(c: "Conditioning", n_new: int) -> "Conditioning":
     """Slide a conditioning window forward by `n_new` measures.
 
     The first T-n_new slots become the OLD window's tail -- those are the
-    measures the model will be shown as fixed context -- and the n_new slots
-    exposed at the end inherit the old tail's conditioning, so a continuation
-    carries the character of the bars it is following instead of jumping. The
-    per-measure `active` flags travel with their signals, so a section you had
-    running free keeps running free.
-
-    That inheritance is a starting point, not a verdict: the new measures are
-    ordinary editable slots, so shape them and hit Regenerate if the
-    continuation should go somewhere the previous eight bars didn't."""
+    measures handed to the model as fixed context, so their conditioning has to
+    travel with their latents or the model is told one thing and shown another.
+    Values and per-measure `active` flags both shift.
+
+    The n_new slots exposed at the end are the ones being WRITTEN, and what goes
+    in them is a real decision:
+
+      style, bpm  -- carried forward from the last context measure, active flag
+                     included. These are properties of the piece, not
+                     descriptions of a particular bar; a continuation should
+                     stay in the same style and the same tempo unless asked
+                     otherwise, and bpm additionally has to hold a real number
+                     because it sets the decoded measure duration.
+
+      chroma, rms, density, zcr, flatness  -- turned OFF (free). These describe
+                     what a bar actually contains, and nobody knows that for a
+                     bar that has not been written yet. Copying the previous
+                     eight bars' values forward would assert the new bars have
+                     the same chords, the same energy and the same texture as
+                     the ones they follow, which is both a claim we cannot make
+                     and a direct instruction to repeat. Freeing them swaps in
+                     the null token and clears the stage-2 presence channel,
+                     which is in-distribution (training dropped signals per
+                     token, p_ind_uncond=0.5 at stage 2) and lets the frozen
+                     context latents decide the continuation instead.
+
+    The freed values still get back-filled from the render by
+    `_measure_render`, so after a Continue the curves show what the model
+    actually played and are immediately editable -- shape them and hit
+    Regenerate to push the new bars somewhere else."""
     T = c.T
     n_new = max(1, min(int(n_new), T - 1))
     n_ctx = T - n_new
+    last = n_ctx - 1                     # last context measure, post-shift
+    carry = ("style", "bpm")
     new = c.copy()
     for s in Conditioning._NPZ_SIGNALS:
         arr, old = getattr(new, s), getattr(c, s)
         arr[:n_ctx] = old[n_new:]
-        arr[n_ctx:] = old[n_ctx:]
+        # seed every new slot from the last context measure. For style/bpm that
+        # IS the conditioning; for the freed signals it is just a placeholder so
+        # the arrays hold something plausible until the render overwrites them.
+        arr[n_ctx:] = arr[last]
     for s in ALL_SIGNALS:
-        new.active[s][:n_ctx] = c.active[s][n_new:]
-        new.active[s][n_ctx:] = c.active[s][n_ctx:]
+        a = new.active[s]
+        a[:n_ctx] = c.active[s][n_new:]
+        a[n_ctx:] = a[last] if s in carry else False
     return new
 
 
@@ -294,7 +331,7 @@ class Track:
     # (guidance 3.0 matches training's `{i}_cfg.wav`). MockBackend ignores it.
     guidance: float = 3.0
     # Cached model latents (B=1, T, N, C) from the last render. Used by
-    # RealBackend for RePaint-style region inpainting; ignored by MockBackend.
+    # RealBackend to freeze measures during region regen / continuation.
     latents: Optional[np.ndarray] = None
 
     # ---- continuation ------------------------------------------------------
@@ -778,8 +815,8 @@ class MockBackend(ModelBackend):
 
     def regenerate_region(self, track, start, end, seed, n_steps=100):
         """Mock inpaint: keep audio outside [start,end), re-render inside with a
-        new seed. The real backend resamples only those latents (RePaint-style)
-        and re-injects known latents each Euler step -- see RealBackend notes."""
+        new seed. The real backend resamples only those latents, holding the
+        kept measures clean at noise level ~0 -- see RealBackend notes."""
         tr = track.copy()
         s, e = self._range(start, end, tr.cond.T)
         full, bounds = self._synthesize(tr.cond, seed=seed)   # fresh noise render

real_backend.py

@@ -90,7 +90,8 @@ import numpy as np
 
 from backend import (
     ModelBackend, Track, Conditioning,
-    ALL_SIGNALS, CONTINUE_MEASURES, EDITABLE_SIGNALS, N_CHUNKS, STYLE_DIM,
+    ALL_SIGNALS, CONTINUE_CTX_NOISE, CONTINUE_MEASURES, EDITABLE_SIGNALS,
+    N_CHUNKS, STYLE_DIM,
     TARGET_SIG, SR,
     bpm_curve, eval_shape, shift_window,
 )
@@ -127,13 +128,13 @@ def _crossfade_segments(segment_a, segment_b, sample_rate, crossfade_ms=15):
 
 
 # =============================================================================
-# Inpainting sampler (RePaint-style), standalone so diffusion_forcing.py is
-# untouched. Works with the same FM diffusion + get_prediction/CFG machinery.
+# Masked sampler, DIFFUSION-FORCING style. Standalone so diffusion_forcing.py
+# stays untouched; uses the same FM diffusion + get_prediction/CFG machinery.
 # =============================================================================
-def fm_inpaint_sample(
+def fm_df_sample(
     diffusion,          # FM instance (net wrapper's .diffusion)
     net,                # the raw DiT (wrapper.net)
-    known_x0,           # (B, T, N, C) clean latents to preserve outside the mask
+    known_x0,           # (B, T, N, C) clean latents to preserve where kept
     keep_mask,          # (B, T, 1, 1) bool: True where latents are FROZEN (kept)
     n_steps,
     net_kwargs=None,
@@ -144,15 +145,45 @@ def fm_inpaint_sample(
     rescale_phi=0.0,
     cfg_mode="independent",
     t_dist="uniform",
-    n_resample=1,       # RePaint resampling steps per timestep (1 == plain)
+    ctx_noise=0.0,      # noise level to hold the context at; see below
 ):
-    """RePaint-style masked FM sampling.
-
-    At every Euler step the KNOWN region is overwritten with the forward-noised
-    ground-truth latents at the current noise level, so only the masked region
-    is actually generated while staying coherent with the fixed context. Setting
-    n_resample > 1 does RePaint's harmonization (renoise/redenoise) which helps
-    the boundary blend.
+    """Masked FM sampling with a PER-MEASURE noise level.
+
+    This model is a diffusion-forcing model, and that is the whole point of it.
+    FM.loss draws
+
+        t = torch.sigmoid(torch.randn(B, T))        # (B, T) -- one per measure
+
+    so every measure was trained at its OWN independently-sampled noise level,
+    and MetaConditionalModernDiTV2 modulates per measure to match (`t = t +
+    style + bpm`, then `t[:, :, None]` broadcast across that measure's latent
+    positions). Attention is fully bidirectional, so a clean measure and a
+    pure-noise measure can sit in the same sequence and the clean one is read as
+    clean. Conditioning on the past is therefore NATIVE: hand the context
+    measures over at t~0 with their real latents and denoise only the rest.
+
+    That is what this does. The context is pinned to `known_x0` UNCHANGED at
+    every Euler step and reported at noise level `ctx_noise`, while the free
+    measures walk the schedule from t=1 down to 0.
+
+    Why not RePaint (which is what this function used to do): RePaint exists
+    because ordinary diffusion models accept a single scalar t, so the only way
+    to express "this part is known" is to forward-noise the known part to the
+    same level as the part being generated and splice it in. Here that is both
+    unnecessary and actively destructive -- at the first step (t=1) the context
+    gets noised to t=1, i.e. handed over as pure noise, so the model has no past
+    to condition on exactly when the global structure is being decided. It also
+    redraws that noise every step, so what context does survive flickers. Uniform
+    t is in-distribution for training (all-equal is one draw of (B,T) among
+    many), but it throws away the one capability that makes rollout work.
+
+    ctx_noise: 0.0 hands the context over clean. Diffusion Forcing's stabilization
+    result is that a SMALL nonzero level on the history makes long autoregressive
+    rollouts more robust -- the model treats its own past output as uncertain
+    instead of trusting artifacts -- so if repeated Continues drift or harden,
+    try ~0.02-0.1 (CONTINUE_CTX_NOISE in backend.py). Held clean at 0.0 the
+    context latents are bit-exact, which is what lets the re-decoded context
+    match the audio already in the mix.
 
     Returns x0-space latents (B, T, N, C)."""
     import torch
@@ -170,42 +201,41 @@ def fm_inpaint_sample(
         t_steps[0] = 1.0
         t_steps[-1] = 0.0
 
-    keep = keep_mask.to(dev)
+    keep = keep_mask.to(dev)                       # (B, T, 1, 1)
+    keep_row = keep[:, :, 0, 0]                    # (B, T) -- for the t vector
+    ctx_noise = float(ctx_noise)
 
-    def _noised_known(t_scalar):
-        # FM forward process: x_t = alpha(t)*x0 + sigma(t)*noise  (per add_noise)
-        tt = t_scalar.view(1, 1, 1, 1).expand(B, T, N, 1)
-        eps = torch.randn_like(known_x0)
-        return diffusion.alpha(tt) * known_x0 + diffusion.sigma(tt) * eps
+    # The context at level `ctx_noise`. Drawn ONCE, not per step: the context is
+    # a fixed thing the model is reading, not a quantity being resampled. At the
+    # default 0.0 this is exactly known_x0 (alpha(0)=1, sigma(0)=sigma_min=1e-5).
+    if ctx_noise <= 0.0:
+        ctx_x = known_x0
+    else:
+        tt = torch.full((B, T, N, 1), ctx_noise, device=dev, dtype=known_x0.dtype)
+        ctx_x = (diffusion.alpha(tt) * known_x0
+                 + diffusion.sigma(tt) * torch.randn_like(known_x0))
 
     with torch.no_grad():
-        # clamp the known region at the initial (fully-noised) level too
-        x_t = torch.where(keep, _noised_known(t_steps[0]), x_t)
+        x_t = torch.where(keep, ctx_x, x_t)
         for i in range(n_steps):
             t_cur, t_next = t_steps[i], t_steps[i + 1]
-            for r in range(n_resample):
-                t = t_cur.repeat(x_t.shape[0], x_t.shape[1])
-                neg_v = diffusion.get_prediction(
-                    net, x_t, t,
-                    net_kwargs=net_kwargs,
-                    uncond_net_kwargs=uncond_net_kwargs,
-                    guidance=guidance,
-                    memory_efficient=memory_efficient,
-                    rescale_phi=rescale_phi,
-                    cfg_mode=cfg_mode,
-                )
-                x_next = x_t + neg_v * (t_next - t_cur)
-                # force the known region to the correctly-noised GT at t_next
-                x_next = torch.where(keep, _noised_known(t_next), x_next)
-
-                if r < n_resample - 1 and t_next > 0:
-                    # RePaint harmonize: jump back up to t_cur and redo
-                    step_up = (diffusion.sigma(t_cur) / diffusion.sigma(t_next)).clamp(max=5.0)
-                    x_t = x_next + torch.randn_like(x_next) * (
-                        diffusion.sigma(t_cur) - diffusion.sigma(t_next))
-                    x_t = torch.where(keep, _noised_known(t_cur), x_t)
-                else:
-                    x_t = x_next
+            # per-measure noise level: context sits at ctx_noise, free measures
+            # at wherever the schedule currently is.
+            t = torch.where(keep_row, torch.full_like(keep_row, ctx_noise,
+                                                      dtype=x_t.dtype),
+                            t_cur.to(x_t.dtype))
+            neg_v = diffusion.get_prediction(
+                net, x_t, t,
+                net_kwargs=net_kwargs,
+                uncond_net_kwargs=uncond_net_kwargs,
+                guidance=guidance,
+                memory_efficient=memory_efficient,
+                rescale_phi=rescale_phi,
+                cfg_mode=cfg_mode,
+            )
+            x_t = x_t + neg_v * (t_next - t_cur)
+            # the Euler step moved the context too; put it back untouched.
+            x_t = torch.where(keep, ctx_x, x_t)
     return x_t
 
 
@@ -920,7 +950,7 @@ class RealBackend(ModelBackend):
         return (out / (np.linalg.norm(out) + 1e-9)).astype(np.float32)
 
     # =====================================================================
-    # region regenerate == RePaint inpainting over measure rows [start, end)
+    # region regenerate == masked resample of measure rows [start, end)
     # =====================================================================
     def regenerate_region(self, track, start, end, seed, n_steps=SAMPLE_STEPS):
         import torch
@@ -945,11 +975,12 @@ class RealBackend(ModelBackend):
             self._sampling_args(c, getattr(tr, "guidance", None))
 
         with self._ctx():
-            y = fm_inpaint_sample(
+            y = fm_df_sample(
                 self.diffusion, self.net, known_x0, keep, n_steps,
                 net_kwargs=cfg_net_kwargs, uncond_net_kwargs=uncond_net_kwargs,
                 guidance=guidances, noise=noise, memory_efficient=False,
-                cfg_mode=cfg_mode, t_dist=self.t_dist, n_resample=1)
+                cfg_mode=cfg_mode, t_dist=self.t_dist,
+                ctx_noise=CONTINUE_CTX_NOISE)
 
         bpm_row = torch.from_numpy(
             np.asarray(c.bpm, dtype=np.float32))[None].to(self.device)
@@ -978,11 +1009,12 @@ class RealBackend(ModelBackend):
                               FROZEN latents         sampled
 
         The 24-n_new context measures keep their exact latents from the previous
-        render and are held fixed through every Euler step -- the same
-        fm_inpaint_sample machinery region-regen uses, just with the free region
-        pinned to the end of the window. So the new bars are conditioned on the
-        actual audio that precedes them, not merely on a similar-looking
-        conditioning curve.
+        render and are handed to the model CLEAN, at noise level ~0, while the
+        new measures denoise from t=1 -- diffusion forcing's per-measure noise
+        level, which is what this model was trained for (see fm_df_sample). Same
+        machinery region-regen uses, just with the free region pinned to the end
+        of the window. So the new bars are conditioned on the actual audio that
+        precedes them, not merely on a similar-looking conditioning curve.
 
         Only the newly sampled tail is appended to the song; the context
         measures are already in the mix and get re-decoded solely because the
@@ -1020,11 +1052,12 @@ class RealBackend(ModelBackend):
             self._sampling_args(c, getattr(tr, "guidance", None))
 
         with self._ctx():
-            y = fm_inpaint_sample(
+            y = fm_df_sample(
                 self.diffusion, self.net, known_x0, keep, n_steps,
                 net_kwargs=cfg_net_kwargs, uncond_net_kwargs=uncond_net_kwargs,
                 guidance=guidances, noise=noise, memory_efficient=False,
-                cfg_mode=cfg_mode, t_dist=self.t_dist, n_resample=1)
+                cfg_mode=cfg_mode, t_dist=self.t_dist,
+                ctx_noise=CONTINUE_CTX_NOISE)
 
         bpm_row = torch.from_numpy(
             np.asarray(c.bpm, dtype=np.float32))[None].to(self.device)