test_e2e.py

"""End-to-end smoke test: load engine, generate audio, validate output.

Confirms:
  1. Package imports cleanly
  2. Models download (or load from cache)
  3. AWQ swap works
  4. Single-speaker generation produces valid audio
  5. Multi-speaker generation produces valid audio
"""
import os, sys, time, logging
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# Use a known voice file
VOICES_DIR = Path("/home/nishant/Downloads/voices")
SS_VOICE = VOICES_DIR / "Alan_Rickman.wav"
MS_VOICES = [
    VOICES_DIR / "David_Attenborough CC3.wav",
    VOICES_DIR / "Alan_Rickman.wav",
    VOICES_DIR / "FreemanM_1.wav",
]

OUT_DIR = Path(__file__).parent / "test_output"
OUT_DIR.mkdir(exist_ok=True)

print("=" * 60)
print("vibevoice-engine end-to-end smoke test")
print("=" * 60)

# Step 1: import
print("\n[1/4] Import package")
from vibevoice_engine import VibeVoiceEngine, EngineConfig
print("  ✓ imports OK")

# Step 2: instantiate + load (downloads if not cached)
print("\n[2/4] Load engine (unified AWQ-INT4 model)")
t0 = time.time()
engine = VibeVoiceEngine(
    model_path="ncoder-ai/VibeVoice-Large-AWQ",
    inference_steps=7,
    dtype="float16",
    attn_implementation="sdpa",
)
engine.load()
print(f"  ✓ loaded in {time.time()-t0:.1f}s")
import torch
if torch.cuda.is_available():
    print(f"  ✓ VRAM: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

# Step 3: single-speaker generation
print("\n[3/4] Single-speaker generation")
t0 = time.time()
audio = engine.generate(
    "Speaker 0: This is a test of the standalone engine package.",
    voices=[str(SS_VOICE)],
)
elapsed = time.time() - t0
duration = len(audio) / engine.config.sample_rate
print(f"  ✓ gen={elapsed:.2f}s, audio={duration:.2f}s, RTF={elapsed/duration:.3f}")
import soundfile as sf
sf.write(OUT_DIR / "ss_test.wav", audio, engine.config.sample_rate)
print(f"  ✓ wrote {OUT_DIR}/ss_test.wav ({(OUT_DIR/'ss_test.wav').stat().st_size/1024:.0f} KB)")
# Sanity: audio is non-trivial
import numpy as np
assert audio.shape[0] > engine.config.sample_rate, f"audio too short: {audio.shape}"
assert np.abs(audio).max() > 0.01, f"audio is silent: max={np.abs(audio).max()}"
print(f"  ✓ audio sanity OK (shape={audio.shape}, max={np.abs(audio).max():.3f})")

# Step 4: multi-speaker generation
print("\n[4/4] Multi-speaker generation (3 speakers)")
script = """\
Speaker 0: Welcome to the council. We've gathered to discuss the situation.
Speaker 1: I appreciate the invitation. The mountains have been quiet too long.
Speaker 2: Quiet doesn't mean safe. The shadows return whether we look or not."""
t0 = time.time()
audio = engine.generate(script, voices=[str(p) for p in MS_VOICES])
elapsed = time.time() - t0
duration = len(audio) / engine.config.sample_rate
print(f"  ✓ gen={elapsed:.2f}s, audio={duration:.2f}s, RTF={elapsed/duration:.3f}")
sf.write(OUT_DIR / "ms_test.wav", audio, engine.config.sample_rate)
print(f"  ✓ wrote {OUT_DIR}/ms_test.wav ({(OUT_DIR/'ms_test.wav').stat().st_size/1024:.0f} KB)")
assert audio.shape[0] > engine.config.sample_rate * 5, f"multi-speaker audio too short: {audio.shape}"
print(f"  ✓ audio sanity OK (shape={audio.shape}, max={np.abs(audio).max():.3f})")

print("\n" + "=" * 60)
print("ALL TESTS PASSED ✓")
print("=" * 60)