Upload test_e2e.py with huggingface_hub

test_e2e.py

@@ -0,0 +1,375 @@
+#!/usr/bin/env python3
+"""
+End-to-end test comparing PyTorch SAM Audio with ONNX Runtime.
+
+This script:
+1. Loads a real audio sample from AudioCaps
+2. Runs PyTorch inference using the original SAMAudio model
+3. Runs ONNX inference using the exported models
+4. Compares the output waveforms
+"""
+
+import torch
+import torchaudio
+import numpy as np
+import os
+from datasets import load_dataset
+
+
+def load_audiocaps_sample():
+    """Load a sample from AudioCaps dataset."""
+    print("Loading AudioCaps sample...")
+    dset = load_dataset(
+        "parquet",
+        data_files="hf://datasets/OpenSound/AudioCaps/data/test-00000-of-00041.parquet",
+    )
+    sample = dset["train"][8]["audio"].get_all_samples()
+    print(f"  Sample rate: {sample.sample_rate}")
+    print(f"  Duration: {sample.data.shape[-1] / sample.sample_rate:.2f}s")
+    return sample
+
+
+def run_pytorch_inference(sample, device="cpu"):
+    """Run inference using PyTorch SAMAudio model."""
+    print("\n=== PyTorch Inference ===")
+    
+    from sam_audio import SAMAudio, SAMAudioProcessor
+    
+    # Load model and processor
+    print("Loading SAMAudio model...")
+    model = SAMAudio.from_pretrained("facebook/sam-audio-small").to(device).eval()
+    processor = SAMAudioProcessor.from_pretrained("facebook/sam-audio-small")
+    
+    # Resample and prepare input
+    wav = torchaudio.functional.resample(
+        sample.data, sample.sample_rate, processor.audio_sampling_rate
+    )
+    wav = wav.mean(0, keepdim=True)  # Convert to mono
+    
+    print(f"  Input audio shape: {wav.shape}")
+    print(f"  Sample rate: {processor.audio_sampling_rate}")
+    
+    # Prepare inputs with explicit anchor
+    inputs = processor(
+        audios=[wav],
+        descriptions=["A horn honking"],
+        anchors=[[["+", 6.3, 7.0]]]
+    ).to(device)
+    
+    # Run separation
+    print("Running separation...")
+    with torch.inference_mode():
+        result = model.separate(inputs)
+    
+    separated_audio = result.target[0].cpu().numpy()
+    print(f"  Output shape: {separated_audio.shape}")
+    
+    return separated_audio, processor.audio_sampling_rate, wav.numpy()
+
+
+def run_onnx_inference(sample, model_dir="."):
+    """Run inference using ONNX models."""
+    print("\n=== ONNX Runtime Inference ===")
+    
+    import onnxruntime as ort
+    from transformers import AutoTokenizer
+    import json
+    
+    # Load models
+    print("Loading ONNX models...")
+    providers = ["CPUExecutionProvider"]
+    
+    dacvae_encoder = ort.InferenceSession(
+        os.path.join(model_dir, "dacvae_encoder.onnx"),
+        providers=providers,
+    )
+    dacvae_decoder = ort.InferenceSession(
+        os.path.join(model_dir, "dacvae_decoder.onnx"),
+        providers=providers,
+    )
+    t5_encoder = ort.InferenceSession(
+        os.path.join(model_dir, "t5_encoder.onnx"),
+        providers=providers,
+    )
+    dit = ort.InferenceSession(
+        os.path.join(model_dir, "dit_single_step.onnx"),
+        providers=providers,
+    )
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_dir, "tokenizer"))
+    print("  All models loaded")
+    
+    # Prepare audio (resample to 44.1kHz for DACVAE)
+    wav = torchaudio.functional.resample(
+        sample.data, sample.sample_rate, 44100
+    )
+    wav = wav.mean(0, keepdim=True)  # Convert to mono
+    audio = wav.numpy().reshape(1, 1, -1).astype(np.float32)
+    
+    print(f"  Input audio shape: {audio.shape}")
+    
+    # 1. Encode audio
+    print("Encoding audio...")
+    latent = dacvae_encoder.run(
+        ["latent_features"],
+        {"audio": audio}
+    )[0]
+    print(f"  Audio latent shape: {latent.shape}")
+    
+    # 2. Encode text
+    print("Encoding text...")
+    tokens = tokenizer(
+        "A horn honking",
+        return_tensors="np",
+        padding=True,
+        truncation=True,
+        max_length=77,
+    )
+    text_features = t5_encoder.run(
+        ["hidden_states"],
+        {
+            "input_ids": tokens["input_ids"].astype(np.int64),
+            "attention_mask": tokens["attention_mask"].astype(np.int64),
+        }
+    )[0]
+    print(f"  Text features shape: {text_features.shape}")
+    
+    # 3. Run ODE solving (simplified - just one step for testing)
+    print("Running DiT (simplified test - 1 step)...")
+    batch_size = 1
+    latent_dim = latent.shape[1]  # 128
+    time_steps = latent.shape[2]
+    
+    # Prepare inputs
+    # SAMAudio._get_audio_features: returns torch.cat([audio_features, audio_features], dim=2)
+    # So audio_features is the mixture DUPLICATED, not mixture + zeros!
+    mixture_features = latent.transpose(0, 2, 1)  # (B, T, 128) - from DACVAE
+    
+    # Duplicate mixture features (this is what SAMAudio actually does)
+    audio_features = np.concatenate([
+        mixture_features,  # Mixture latent
+        mixture_features   # Mixture latent (DUPLICATE - not zeros!)
+    ], axis=-1)  # -> (B, T, 256)
+    
+    # noisy_audio starts from random noise for ODE solving from t=0 to t=1
+    # SAMAudio uses: noise = torch.randn_like(audio_features)
+    initial = np.random.randn(batch_size, time_steps, 256).astype(np.float32)
+    
+    # Just run one step to verify the model works
+    velocity = dit.run(
+        ["velocity"],
+        {
+            "noisy_audio": initial,
+            "time": np.array([0.0], dtype=np.float32),
+            "audio_features": audio_features,
+            "text_features": text_features,
+            "text_mask": tokens["attention_mask"].astype(bool),
+            "masked_video_features": np.zeros((batch_size, 1024, time_steps), dtype=np.float32),
+            "anchor_ids": np.zeros((batch_size, time_steps), dtype=np.int64),
+            "anchor_alignment": np.zeros((batch_size, time_steps), dtype=np.int64),
+            "audio_pad_mask": np.ones((batch_size, time_steps), dtype=bool),
+        }
+    )[0]
+    print(f"  DiT velocity shape: {velocity.shape}")
+
+    
+    # 4. Run full ODE solve (16 steps midpoint method)
+    print("Running full ODE solve (16 steps)...")
+    num_steps = 16
+    dt = 1.0 / num_steps
+    x = initial.copy()
+    
+    for i in range(num_steps):
+        t = np.array([i * dt], dtype=np.float32)
+        t_mid = np.array([t[0] + dt / 2], dtype=np.float32)
+        
+        # k1 = f(t, x)
+        k1 = dit.run(
+            ["velocity"],
+            {
+                "noisy_audio": x,
+                "time": t,
+                "audio_features": audio_features,
+                "text_features": text_features,
+                "text_mask": tokens["attention_mask"].astype(bool),
+                "masked_video_features": np.zeros((batch_size, 1024, time_steps), dtype=np.float32),
+                "anchor_ids": np.zeros((batch_size, time_steps), dtype=np.int64),
+                "anchor_alignment": np.zeros((batch_size, time_steps), dtype=np.int64),
+                "audio_pad_mask": np.ones((batch_size, time_steps), dtype=bool),
+            }
+        )[0]
+        
+        # Midpoint
+        x_mid = x + (dt / 2) * k1
+        
+        # k2 = f(t_mid, x_mid)
+        k2 = dit.run(
+            ["velocity"],
+            {
+                "noisy_audio": x_mid,
+                "time": t_mid,
+                "audio_features": audio_features,
+                "text_features": text_features,
+                "text_mask": tokens["attention_mask"].astype(bool),
+                "masked_video_features": np.zeros((batch_size, 1024, time_steps), dtype=np.float32),
+                "anchor_ids": np.zeros((batch_size, time_steps), dtype=np.int64),
+                "anchor_alignment": np.zeros((batch_size, time_steps), dtype=np.int64),
+                "audio_pad_mask": np.ones((batch_size, time_steps), dtype=bool),
+            }
+        )[0]
+        
+        # Update
+        x = x + dt * k2
+        print(f"  Step {i+1}/{num_steps}")
+    
+    # 5. Extract separated latent and decode in chunks
+    # (The DACVAE decoder was exported with fixed time=25, so we decode in chunks)
+    print("Decoding audio...")
+    # SAMAudio: target is first 128 dims, residual is second 128 dims
+    # generated_features.reshape(2*B, C, T) -> first B = channels 0:128 (target)
+    target_latent = x[:, :, :latent_dim].transpose(0, 2, 1)  # (B, 128, T) - TARGET
+    separated_latent = target_latent
+    
+    # The decoder expects chunks of 25 time steps
+    chunk_size = 25
+    T = separated_latent.shape[2]
+    
+    # Process in chunks and concatenate
+    audio_chunks = []
+    for start_idx in range(0, T, chunk_size):
+        end_idx = min(start_idx + chunk_size, T)
+        chunk = separated_latent[:, :, start_idx:end_idx]
+        
+        # Pad last chunk if needed
+        actual_size = chunk.shape[2]
+        if actual_size < chunk_size:
+            pad_size = chunk_size - actual_size
+            chunk = np.pad(chunk, ((0, 0), (0, 0), (0, pad_size)), mode='constant')
+        
+        chunk_audio = dacvae_decoder.run(
+            ["waveform"],
+            {"latent_features": chunk.astype(np.float32)}
+        )[0]
+        
+        # For padded chunks, trim the output
+        if actual_size < chunk_size:
+            # Each time step produces hop_length (1920) samples at 48kHz
+            samples_per_step = 1920
+            trim_samples = actual_size * samples_per_step
+            chunk_audio = chunk_audio[:, :, :trim_samples]
+        
+        audio_chunks.append(chunk_audio)
+        print(f"    Decoded chunk {start_idx//chunk_size + 1}/{(T + chunk_size - 1)//chunk_size}")
+    
+    # Concatenate all chunks
+    separated_audio = np.concatenate(audio_chunks, axis=2)
+    
+    print(f"  Output audio shape: {separated_audio.shape}")
+    
+    return separated_audio.squeeze(), 44100
+
+
+
+def compare_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr):
+    """Compare PyTorch and ONNX outputs."""
+    print("\n=== Comparison ===")
+    
+    import scipy.signal
+    
+    # Resample to same rate if needed
+    if pytorch_sr != onnx_sr:
+        print(f"Resampling PyTorch output from {pytorch_sr} to {onnx_sr}...")
+        # Use scipy for resampling
+        num_samples = int(len(pytorch_audio) * onnx_sr / pytorch_sr)
+        pytorch_audio_resampled = scipy.signal.resample(pytorch_audio, num_samples)
+    else:
+        pytorch_audio_resampled = pytorch_audio
+    
+    # Trim to same length
+    min_len = min(len(pytorch_audio_resampled), len(onnx_audio))
+    pytorch_trimmed = pytorch_audio_resampled[:min_len]
+    onnx_trimmed = onnx_audio[:min_len]
+    
+    # Compute differences
+    diff = np.abs(pytorch_trimmed - onnx_trimmed)
+    max_diff = diff.max()
+    mean_diff = diff.mean()
+    
+    # Compute correlation
+    correlation = np.corrcoef(pytorch_trimmed, onnx_trimmed)[0, 1]
+    
+    print(f"  PyTorch audio length: {len(pytorch_audio)} samples")
+    print(f"  ONNX audio length: {len(onnx_audio)} samples")
+    print(f"  Max difference: {max_diff:.6f}")
+    print(f"  Mean difference: {mean_diff:.6f}")
+    print(f"  Correlation: {correlation:.6f}")
+    
+    return max_diff, mean_diff, correlation
+
+
+def save_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr, input_audio, input_sr):
+    """Save audio outputs for listening comparison."""
+    import soundfile as sf
+    
+    output_dir = "test_outputs"
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Save input
+    sf.write(os.path.join(output_dir, "input.wav"), input_audio.squeeze(), input_sr)
+    print(f"Saved input to {output_dir}/input.wav")
+    
+    # Save PyTorch output
+    sf.write(os.path.join(output_dir, "pytorch_output.wav"), pytorch_audio, pytorch_sr)
+    print(f"Saved PyTorch output to {output_dir}/pytorch_output.wav")
+    
+    # Save ONNX output
+    sf.write(os.path.join(output_dir, "onnx_output.wav"), onnx_audio, onnx_sr)
+    print(f"Saved ONNX output to {output_dir}/onnx_output.wav")
+
+
+def main():
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="End-to-end SAM Audio test")
+    parser.add_argument("--model-dir", default=".", help="ONNX model directory")
+    parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+    parser.add_argument("--save-outputs", action="store_true", help="Save audio files")
+    parser.add_argument("--skip-pytorch", action="store_true", help="Skip PyTorch inference")
+    args = parser.parse_args()
+    
+    # Load sample
+    sample = load_audiocaps_sample()
+    
+    # Run PyTorch inference
+    if not args.skip_pytorch:
+        pytorch_audio, pytorch_sr, input_audio = run_pytorch_inference(sample, args.device)
+    else:
+        print("\nSkipping PyTorch inference")
+        pytorch_audio, pytorch_sr = None, None
+        input_audio = sample.data.mean(0).numpy()
+    
+    # Run ONNX inference
+    onnx_audio, onnx_sr = run_onnx_inference(sample, args.model_dir)
+    
+    # Compare outputs
+    if pytorch_audio is not None:
+        compare_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr)
+    
+    # Save outputs
+    if args.save_outputs:
+        print("\n=== Saving Outputs ===")
+        if pytorch_audio is not None:
+            save_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr, 
+                        input_audio, sample.sample_rate)
+        else:
+            import soundfile as sf
+            os.makedirs("test_outputs", exist_ok=True)
+            sf.write("test_outputs/onnx_output.wav", onnx_audio, onnx_sr)
+            print("Saved ONNX output to test_outputs/onnx_output.wav")
+    
+    print("\n✓ End-to-end test complete!")
+
+
+if __name__ == "__main__":
+    main()