From cdf43061f2474e54ec91bd0b4903eb3e5d15a229 Mon Sep 17 00:00:00 2001 From: zonnart17-collab Date: Mon, 24 Aug 2026 13:43:31 +0300 Subject: [PATCH] Fix VAE-aligned chunking in full-song inference --- fullsong_chunking.py | 114 +++++++++++++++++++++++ infer_single.py | 96 ++++++++++++-------- inference_fullsong.py | 97 ++++++++++++-------- tests/test_fullsong_chunking.py | 155 ++++++++++++++++++++++++++++++++ 4 files changed, 386 insertions(+), 76 deletions(-) create mode 100644 fullsong_chunking.py create mode 100644 tests/test_fullsong_chunking.py diff --git a/fullsong_chunking.py b/fullsong_chunking.py new file mode 100644 index 0000000..4a084c0 --- /dev/null +++ b/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] diff --git a/infer_single.py b/infer_single.py index 5a6c4fe..9fa3574 100644 --- a/infer_single.py +++ b/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) diff --git a/inference_fullsong.py b/inference_fullsong.py index afa0175..0197fbb 100644 --- a/inference_fullsong.py +++ b/inference_fullsong.py @@ -18,6 +18,15 @@ from datasets import load_dataset from torch.utils.data import Dataset, DataLoader from tqdm.auto import tqdm +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 from datasets import load_dataset, Audio from utils import Text2AudioDataset, read_wav_file, pad_wav @@ -149,6 +158,12 @@ def load_config(config_path): vae.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), + ) + ## Freeze text encoder param for param in model.text_encoder.parameters(): @@ -167,13 +182,7 @@ def load_config(config_path): output_dir = "/outputs/fullsongs/full10sec40g1" os.makedirs(output_dir, exist_ok=True) - # Parameters - fs = 44100 - chunk_duration = 30 - overlap_duration = 10 - chunk_size = chunk_duration * fs - overlap_size = overlap_duration * fs - stride_size = chunk_size - overlap_size + fs = SONICMASTER_SAMPLE_RATE # Load JSONL with open(jsonl_path, "r") as f: @@ -196,17 +205,10 @@ def load_config(config_path): audio = audio.to(device) # [2, T] # Chunking degraded audio - chunks = [] - start = 0 - while start < audio.shape[1]: - end = min(start + chunk_size, audio.shape[1]) - chunk = audio[:, start:end] - - # Pad last chunk - if chunk.shape[1] < chunk_size: - chunk = torch.nn.functional.pad(chunk, (0, chunk_size - chunk.shape[1])) - chunks.append(chunk) - start += stride_size + target_length = audio.shape[1] + chunks = make_overlapping_chunks(audio, geometry) + if not chunks: + raise RuntimeError(f"No audio content to process: {input_path}") # Pre-encode degraded chunks # with torch.no_grad(): @@ -221,7 +223,18 @@ def load_config(config_path): for b in range(0, num_chunks, batch_size): batch = chunk_tensor[b:b+batch_size] # [B, 2, T] + if batch.shape[-1] != geometry.chunk_size: + raise RuntimeError( + "Chunk batch has " + f"{batch.shape[-1]} samples; expected {geometry.chunk_size}." + ) latent = vae.encode(batch).latent_dist.mode() # [B, C, T'] + if latent.shape[-1] != MAIN_LATENT_FRAMES: + raise RuntimeError( + "VAE main encode length mismatch: " + f"{batch.shape[-1]} samples produced {latent.shape[-1]} " + f"frames; expected {MAIN_LATENT_FRAMES}." + ) degraded_latents_list.append(latent) degraded_latents = torch.cat(degraded_latents_list, dim=0) # [N, C, T'] @@ -243,7 +256,7 @@ def load_config(config_path): num_inference_steps=10, timesteps=None, guidance_scale=1, - duration=chunk_duration, + duration=TRAINED_DURATION_SECONDS, seed=0, disable_progress=False, num_samples_per_prompt=1, @@ -253,27 +266,35 @@ def load_config(config_path): # Decode latent to waveform decoded_wave = vae.decode(result_latent.transpose(2, 1)).sample.cpu() # [1, 2, T] + if decoded_wave.shape[-1] != geometry.chunk_size: + raise RuntimeError( + "VAE decoded chunk length mismatch; refusing to stitch: " + f"chunk {i} produced {decoded_wave.shape[-1]} samples, " + f"expected {geometry.chunk_size}." + ) decoded_chunks.append(decoded_wave) - # Get last 10 seconds of waveform → re-encode as latent - last_10_sec = decoded_wave[:, :, -overlap_size:].to(device) - prev_cond_latent = vae.encode(last_10_sec).latent_dist.mode().transpose(1,2) # [1, C, T'] ->[1, T', C] - - # Stitch decoded chunks with crossfade - final_output = decoded_chunks[0] # [1, 2, T] - for i in range(1, len(decoded_chunks)): - prev = final_output[:, :, -overlap_size:] - curr = decoded_chunks[i][:, :, :overlap_size] - - alpha = torch.linspace(1, 0, steps=overlap_size).view(1, 1, -1) - beta = 1 - alpha - blended = prev * alpha + curr * beta - - final_output = torch.cat([ - final_output[:, :, :-overlap_size], - blended, - decoded_chunks[i][:, :, overlap_size:] - ], dim=2) + # Re-encode the VAE-native carry as conditioning for the next chunk. + carry = decoded_wave[:, :, -geometry.overlap:].to(device) + if carry.shape[-1] != geometry.overlap: + raise RuntimeError( + "Conditioning carry has " + f"{carry.shape[-1]} samples; expected {geometry.overlap}." + ) + carry_latent = vae.encode(carry).latent_dist.mode() + if carry_latent.shape[-1] != CARRY_LATENT_FRAMES: + raise RuntimeError( + "VAE carry encode length mismatch: " + f"{carry.shape[-1]} samples produced " + f"{carry_latent.shape[-1]} frames; " + f"expected {CARRY_LATENT_FRAMES}." + ) + prev_cond_latent = carry_latent.transpose(1,2) # [1, C, T'] ->[1, T', C] + + # Remove only the padded stitched tail; underflow is an explicit error. + final_output = crossfade_and_trim( + decoded_chunks, geometry, target_length + ) # Save to file output_path = os.path.join(output_dir, f"{song_id}_reconstructed.flac") diff --git a/tests/test_fullsong_chunking.py b/tests/test_fullsong_chunking.py new file mode 100644 index 0000000..d4fd611 --- /dev/null +++ b/tests/test_fullsong_chunking.py @@ -0,0 +1,155 @@ +import ast +import sys +import unittest +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from fullsong_chunking import ( + CARRY_LATENT_FRAMES, + MAIN_LATENT_FRAMES, + SONICMASTER_SAMPLE_RATE, + TRAINED_DURATION_SECONDS, + FullSongGeometry, + crossfade_and_trim, + make_overlapping_chunks, + make_vae_native_geometry, +) + + +class FullSongGeometryTests(unittest.TestCase): + def test_published_checkpoint_contract(self): + self.assertEqual(SONICMASTER_SAMPLE_RATE, 44_100) + self.assertEqual(MAIN_LATENT_FRAMES, 645) + self.assertEqual(CARRY_LATENT_FRAMES, MAIN_LATENT_FRAMES // 3) + self.assertEqual(CARRY_LATENT_FRAMES, 215) + self.assertEqual(TRAINED_DURATION_SECONDS, 30) + + def test_vae_native_waveform_geometry(self): + geometry = make_vae_native_geometry(2048, 645, 44_100) + self.assertEqual(geometry.chunk_size, 1_320_960) + self.assertEqual(geometry.overlap, 440_320) + self.assertEqual(geometry.stride, 880_640) + self.assertEqual(geometry.chunk_size % geometry.hop_length, 0) + self.assertEqual(geometry.overlap % geometry.hop_length, 0) + + def test_invalid_checkpoint_contract_is_rejected(self): + with self.assertRaisesRegex(RuntimeError, "hop length"): + make_vae_native_geometry(0, 645, 44_100) + with self.assertRaisesRegex(RuntimeError, "sequence length"): + make_vae_native_geometry(2048, 646, 44_100) + with self.assertRaisesRegex(RuntimeError, "44100 Hz VAE"): + make_vae_native_geometry(2048, 645, 48_000) + + +class FullSongStitchingTests(unittest.TestCase): + geometry = FullSongGeometry( + hop_length=1, + chunk_size=30_000, + overlap=10_000, + stride=20_000, + ) + + def test_empty_audio_produces_no_chunks(self): + audio = torch.empty(2, 0) + self.assertEqual(make_overlapping_chunks(audio, self.geometry), []) + + def test_short_audio_is_padded_then_trimmed_exactly(self): + source = torch.randn(2, 1_000) + chunks = make_overlapping_chunks(source, self.geometry) + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].shape, (2, self.geometry.chunk_size)) + result = crossfade_and_trim( + [chunks[0].unsqueeze(0)], self.geometry, source.shape[1] + ) + self.assertEqual(result.shape, (1, 2, 1_000)) + torch.testing.assert_close(result.squeeze(0), source) + + def test_exact_stride_boundary_is_preserved(self): + source = torch.randn(2, self.geometry.stride) + chunks = make_overlapping_chunks(source, self.geometry) + self.assertEqual(len(chunks), 1) + result = crossfade_and_trim( + [chunks[0].unsqueeze(0)], self.geometry, source.shape[1] + ) + torch.testing.assert_close(result.squeeze(0), source) + + def test_multi_chunk_track_preserves_stereo_and_boundaries(self): + samples = 65_000 + t = torch.arange(samples, dtype=torch.float32) + source = torch.stack((torch.sin(t * 0.013), torch.cos(t * 0.017))) + chunks = make_overlapping_chunks(source, self.geometry) + decoded = [chunk.unsqueeze(0) for chunk in chunks] + result = crossfade_and_trim(decoded, self.geometry, samples) + self.assertEqual(result.shape, (1, 2, samples)) + torch.testing.assert_close(result.squeeze(0), source, rtol=1e-6, atol=1e-6) + self.assertTrue(torch.equal(result[0, :, 0], source[:, 0])) + self.assertTrue(torch.equal(result[0, :, -1], source[:, -1])) + + def test_final_trim_cannot_hide_stitched_underflow(self): + decoded = [torch.zeros(1, 2, 999)] + with self.assertRaisesRegex(RuntimeError, "shorter than the source"): + crossfade_and_trim(decoded, self.geometry, 1_000) + + def test_crossfade_rejects_chunk_shorter_than_overlap(self): + decoded = [ + torch.zeros(1, 2, self.geometry.chunk_size), + torch.zeros(1, 2, self.geometry.overlap - 1), + ] + with self.assertRaisesRegex(RuntimeError, "shorter than the VAE-native overlap"): + crossfade_and_trim(decoded, self.geometry, self.geometry.chunk_size) + + +class FullSongEntryPointTests(unittest.TestCase): + required_imports = { + "CARRY_LATENT_FRAMES", + "MAIN_LATENT_FRAMES", + "SONICMASTER_SAMPLE_RATE", + "TRAINED_DURATION_SECONDS", + "crossfade_and_trim", + "make_overlapping_chunks", + "make_vae_native_geometry", + } + + def entry_point_source(self, filename): + return (ROOT / filename).read_text(encoding="utf-8") + + def imported_helper_names(self, source): + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module == "fullsong_chunking": + return {alias.name for alias in node.names} + return set() + + def assert_shared_contract(self, filename): + source = self.entry_point_source(filename) + self.assertTrue(self.required_imports <= self.imported_helper_names(source)) + self.assertIn("make_vae_native_geometry(", source) + self.assertIn("make_overlapping_chunks(", source) + self.assertIn("crossfade_and_trim(", source) + self.assertIn("duration=TRAINED_DURATION_SECONDS", source) + self.assertIn("VAE main encode length mismatch", source) + self.assertIn("VAE decoded chunk length mismatch", source) + self.assertIn("VAE carry encode length mismatch", source) + + def test_infer_single_uses_shared_contract_and_has_no_false_cli_knobs(self): + source = self.entry_point_source("infer_single.py") + self.assert_shared_contract("infer_single.py") + self.assertNotIn('"--fs"', source) + self.assertNotIn('"--chunk_duration"', source) + self.assertNotIn('"--overlap_duration"', source) + self.assertIn("torchaudio.functional.resample", source) + self.assertIn("SONICMASTER_SAMPLE_RATE", source) + + def test_dataset_fullsong_path_uses_shared_contract(self): + source = self.entry_point_source("inference_fullsong.py") + self.assert_shared_contract("inference_fullsong.py") + self.assertIn("fs = SONICMASTER_SAMPLE_RATE", source) + self.assertIn("Expected {fs} Hz", source) + + +if __name__ == "__main__": + unittest.main()