Fix VAE-aligned chunking in full-song inference

fullsong_chunking.py

@@ -0,0 +1,114 @@
+"""Shared VAE-native geometry and stitching for full-song inference."""
+
+from dataclasses import dataclass
+
+import torch
+
+
+SONICMASTER_SAMPLE_RATE = 44_100
+MAIN_LATENT_FRAMES = 645
+CARRY_LATENT_FRAMES = MAIN_LATENT_FRAMES // 3
+TRAINED_DURATION_SECONDS = 30
+
+
+@dataclass(frozen=True)
+class FullSongGeometry:
+    """Waveform geometry corresponding to the published checkpoint latents."""
+
+    hop_length: int
+    chunk_size: int
+    overlap: int
+    stride: int
+
+
+def make_vae_native_geometry(
+    vae_hop_length, model_audio_seq_len, vae_sample_rate
+):
+    """Validate the published checkpoint contract and return sample geometry."""
+    hop_length = int(vae_hop_length)
+    if hop_length <= 0:
+        raise RuntimeError(f"Invalid VAE hop length: {hop_length}")
+    if int(model_audio_seq_len) != MAIN_LATENT_FRAMES:
+        raise RuntimeError(
+            "VAE-native chunking requires the trained SonicMaster sequence "
+            f"length {MAIN_LATENT_FRAMES}, got {model_audio_seq_len}."
+        )
+    if int(vae_sample_rate) != SONICMASTER_SAMPLE_RATE:
+        raise RuntimeError(
+            "The published SonicMaster checkpoint requires a "
+            f"{SONICMASTER_SAMPLE_RATE} Hz VAE, got {vae_sample_rate}."
+        )
+
+    chunk_size = MAIN_LATENT_FRAMES * hop_length
+    overlap = CARRY_LATENT_FRAMES * hop_length
+    if overlap <= 0 or overlap >= chunk_size:
+        raise RuntimeError(
+            "VAE-native overlap must be positive and smaller than the chunk."
+        )
+    return FullSongGeometry(
+        hop_length=hop_length,
+        chunk_size=chunk_size,
+        overlap=overlap,
+        stride=chunk_size - overlap,
+    )
+
+
+def make_overlapping_chunks(audio, geometry):
+    """Return right-padded ``[channels, chunk_size]`` waveform chunks."""
+    chunks = []
+    start = 0
+    total = audio.shape[1]
+    while start < total:
+        end = min(start + geometry.chunk_size, total)
+        chunk = audio[:, start:end]
+        if chunk.shape[1] < geometry.chunk_size:
+            chunk = torch.nn.functional.pad(
+                chunk, (0, geometry.chunk_size - chunk.shape[1])
+            )
+        chunks.append(chunk)
+        start += geometry.stride
+    return chunks
+
+
+def crossfade_and_trim(decoded_chunks, geometry, target_length):
+    """Linear-crossfade chunks and trim padding without hiding underflow."""
+    if not decoded_chunks:
+        raise ValueError("decoded_chunks must not be empty")
+    if target_length < 0:
+        raise ValueError("target_length must not be negative")
+
+    final = decoded_chunks[0]
+    for current in decoded_chunks[1:]:
+        if final.shape[-1] < geometry.overlap:
+            raise RuntimeError(
+                "Previous stitched waveform is shorter than the VAE-native overlap."
+            )
+        if current.shape[-1] < geometry.overlap:
+            raise RuntimeError(
+                "Decoded chunk is shorter than the VAE-native overlap."
+            )
+        previous_overlap = final[:, :, -geometry.overlap :]
+        current_overlap = current[:, :, : geometry.overlap]
+        alpha = torch.linspace(
+            1.0,
+            0.0,
+            steps=geometry.overlap,
+            dtype=previous_overlap.dtype,
+            device=previous_overlap.device,
+        ).view(1, 1, -1)
+        blended = previous_overlap * alpha + current_overlap * (1.0 - alpha)
+        final = torch.cat(
+            [
+                final[:, :, : -geometry.overlap],
+                blended,
+                current[:, :, geometry.overlap :],
+            ],
+            dim=2,
+        )
+
+    if target_length > final.shape[-1]:
+        raise RuntimeError(
+            "Stitched waveform is shorter than the source; refusing to hide a "
+            f"length error with final trim ({final.shape[-1]} < {target_length})."
+        )
+    return final[:, :, :target_length]

infer_single.py

@@ -12,6 +12,15 @@
 from diffusers import AutoencoderOobleck
 
 # Local imports (repo root is on sys.path when this file is executed)
+from fullsong_chunking import (
+    CARRY_LATENT_FRAMES,
+    MAIN_LATENT_FRAMES,
+    SONICMASTER_SAMPLE_RATE,
+    TRAINED_DURATION_SECONDS,
+    crossfade_and_trim,
+    make_overlapping_chunks,
+    make_vae_native_geometry,
+)
 from model import TangoFlux
 
 hf_token = (
@@ -20,7 +29,6 @@
     or os.getenv("HUGGINGFACEHUB_API_TOKEN")
 )
 
-
 def parse_args():
     p = argparse.ArgumentParser("Single-sample inference for SonicMaster")
     p.add_argument("--ckpt", type=str, required=True,
@@ -35,9 +43,6 @@ def parse_args():
     # Optional knobs (safe defaults)
     p.add_argument("--config", type=str, default=str(Path(__file__).parent / "configs" / "tangoflux_config.yaml"),
                    help="YAML config defining model sizes/hparams.")
-    p.add_argument("--fs", type=int, default=44100, help="Target sample rate.")
-    p.add_argument("--chunk_duration", type=int, default=30, help="Chunk length in seconds.")
-    p.add_argument("--overlap_duration", type=int, default=10, help="Overlap (and carry) in seconds.")
     p.add_argument("--vae_batch_size", type=int, default=10, help="Batch size for VAE encoding over chunks.")
     p.add_argument("--num_inference_steps", type=int, default=10)
     p.add_argument("--guidance_scale", type=float, default=1.0)
@@ -85,6 +90,12 @@ def main():
     ).to(device)
     vae.eval()
 
+    geometry = make_vae_native_geometry(
+        vae_hop_length=getattr(vae, "hop_length", 0),
+        model_audio_seq_len=model.audio_seq_len,
+        vae_sample_rate=getattr(vae, "sampling_rate", 0),
+    )
+
     # --------- Read & standardize input ----------
     in_path = Path(args.input)
     if not in_path.exists():
@@ -97,31 +108,19 @@ def main():
     elif audio.shape[0] > 2:
         audio = audio[:2, :]
 
-    # Resample to target fs
-    if sr != args.fs:
-        audio = torchaudio.functional.resample(audio, sr, args.fs)
-        sr = args.fs
+    # The published VAE and SonicMaster checkpoint are trained at 44.1 kHz.
+    if sr != SONICMASTER_SAMPLE_RATE:
+        audio = torchaudio.functional.resample(
+            audio, sr, SONICMASTER_SAMPLE_RATE
+        )
+        sr = SONICMASTER_SAMPLE_RATE
 
     audio = audio.to(device)
 
     # --------- Chunking ----------
-    fs = args.fs
-    chunk_size = args.chunk_duration * fs
-    overlap = args.overlap_duration * fs
-    if overlap <= 0 or overlap >= chunk_size:
-        raise ValueError("overlap_duration must be >0 and smaller than chunk_duration.")
-    stride = chunk_size - overlap
-
-    chunks = []
-    start = 0
+    fs = SONICMASTER_SAMPLE_RATE
     T = audio.shape[1]
-    while start < T:
-        end = min(start + chunk_size, T)
-        ch = audio[:, start:end]
-        if ch.shape[1] < chunk_size:
-            ch = torch.nn.functional.pad(ch, (0, chunk_size - ch.shape[1]))
-        chunks.append(ch)
-        start += stride
+    chunks = make_overlapping_chunks(audio, geometry)
 
     if not chunks:
         raise RuntimeError("No audio content to process.")
@@ -131,7 +130,18 @@ def main():
     latents = []
     for b in range(0, chunk_tensor.shape[0], args.vae_batch_size):
         batch = chunk_tensor[b:b + args.vae_batch_size].to(device)
+        if batch.shape[-1] != geometry.chunk_size:
+            raise RuntimeError(
+                "Chunk batch has "
+                f"{batch.shape[-1]} samples; expected {geometry.chunk_size}."
+            )
         z = vae.encode(batch).latent_dist.mode()  # [B, C, T']
+        if z.shape[-1] != MAIN_LATENT_FRAMES:
+            raise RuntimeError(
+                "VAE main encode length mismatch: "
+                f"{batch.shape[-1]} samples produced {z.shape[-1]} frames; "
+                f"expected {MAIN_LATENT_FRAMES}."
+            )
         latents.append(z)
     degraded_latents = torch.cat(latents, dim=0)  # [N, C, T']
 
@@ -151,7 +161,7 @@ def main():
             num_inference_steps=args.num_inference_steps,
             timesteps=None,
             guidance_scale=args.guidance_scale,
-            duration=args.chunk_duration,
+            duration=TRAINED_DURATION_SECONDS,
             seed=args.seed,
             disable_progress=True,
             num_samples_per_prompt=1,
@@ -161,26 +171,36 @@ def main():
 
         # Decode to waveform on CPU for stitching
         wav = vae.decode(result_latent.transpose(2, 1)).sample.cpu()  # [1, 2, T]
+        if wav.shape[-1] != geometry.chunk_size:
+            raise RuntimeError(
+                "VAE decoded chunk length mismatch; refusing to stitch: "
+                f"chunk {i} produced {wav.shape[-1]} samples, "
+                f"expected {geometry.chunk_size}."
+            )
         # Safety clamp to [-1, 1]
         wav = torch.clamp(wav, -1.0, 1.0)
         decoded_chunks.append(wav)
 
         # Carry last overlap as conditioning (back on device)
-        last = wav[:, :, -overlap:].to(device)
-        prev_cond = vae.encode(last).latent_dist.mode().transpose(1, 2)  # [1, T', C]
+        last = wav[:, :, -geometry.overlap:].to(device)
+        if last.shape[-1] != geometry.overlap:
+            raise RuntimeError(
+                "Conditioning carry has "
+                f"{last.shape[-1]} samples; expected {geometry.overlap}."
+            )
+        carry_latent = vae.encode(last).latent_dist.mode()
+        if carry_latent.shape[-1] != CARRY_LATENT_FRAMES:
+            raise RuntimeError(
+                "VAE carry encode length mismatch: "
+                f"{last.shape[-1]} samples produced {carry_latent.shape[-1]} frames; "
+                f"expected {CARRY_LATENT_FRAMES}."
+            )
+        prev_cond = carry_latent.transpose(1, 2)  # [1, T', C]
 
     # --------- Crossfade stitch ----------
-    final = decoded_chunks[0]  # [1, 2, T]
-    for i in range(1, len(decoded_chunks)):
-        prev = final[:, :, -overlap:]
-        curr = decoded_chunks[i][:, :, :overlap]
-        alpha = torch.linspace(1.0, 0.0, steps=overlap).view(1, 1, -1)
-        beta = 1.0 - alpha
-        blended = prev * alpha + curr * beta
-        final = torch.cat(
-            [final[:, :, :-overlap], blended, decoded_chunks[i][:, :, overlap:]],
-            dim=2,
-        )
+    # The last partial VAE-native chunk is right-padded for inference. Remove
+    # only its excess stitched tail by preserving the exact resampled length.
+    final = crossfade_and_trim(decoded_chunks, geometry, T)
 
     # --------- Save (honor extension) ----------
     out_path = Path(args.output)