Initial release

tools/onnx-export/export_vae_encoder.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+"""Export AutoencoderOobleck VAE encoder to ONNX format.
+
+Exports the encoder half of the VAE for use with TensorRT or ONNX Runtime.
+The encoder converts audio → latent space for timbre/cover VAE encoding.
+
+Tensor spec:
+  Input:  "audio"    [B, 2, samples]    (stereo, samples @ 48kHz)
+  Output: "latents"  [B, 64, T]         (latent channels, latent frames @ 25Hz)
+
+Note: The encoder output is 128ch (64 mean + 64 scale). We only need the mean
+for deterministic encoding, so we slice to the first 64 channels.
+"""
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class VAEEncoderWrapper(nn.Module):
+    """Wraps AutoencoderOobleck.encoder + quant_conv to extract mean latents.
+
+    The raw encoder returns 128ch (mean + scale). For deterministic encoding
+    we only need the first 64 channels (mean). This wrapper handles the
+    quant_conv and slicing.
+    """
+
+    def __init__(self, vae):
+        super().__init__()
+        self.encoder = vae.encoder
+        # quant_conv maps from encoder output space to latent space
+        if hasattr(vae, "quant_conv") and vae.quant_conv is not None:
+            self.quant_conv = vae.quant_conv
+        else:
+            self.quant_conv = None
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            audio: [B, 2, samples] stereo audio at 48kHz
+
+        Returns:
+            latents: [B, 64, T] mean latents (deterministic)
+        """
+        encoded = self.encoder(audio)
+        # encoder returns EncoderOutput with .latent_dist or raw tensor
+        if hasattr(encoded, "latent_dist"):
+            h = encoded.latent_dist.mean
+        elif hasattr(encoded, "sample"):
+            h = encoded.sample
+        else:
+            h = encoded
+
+        if self.quant_conv is not None:
+            h = self.quant_conv(h)
+
+        # h is [B, 128, T] — first 64 = mean, last 64 = log_var
+        # Only return mean for deterministic encoding
+        return h[:, :64, :]
+
+
+def export_vae_encoder(vae_path: str, output_path: str, opset: int = 18) -> str:
+    """Export VAE encoder to ONNX.
+
+    Args:
+        vae_path: Path to the VAE checkpoint directory (config.json + safetensors).
+        output_path: Path to write the ONNX file.
+        opset: ONNX opset version.
+
+    Returns:
+        The output path of the exported ONNX file.
+    """
+    from diffusers import AutoencoderOobleck
+
+    print(f"Loading VAE from: {vae_path}")
+    vae = AutoencoderOobleck.from_pretrained(vae_path)
+    vae.eval()
+
+    wrapper = VAEEncoderWrapper(vae)
+    wrapper.eval()
+
+    # Move to CPU for export (fp32)
+    wrapper = wrapper.cpu()
+
+    # Create dummy input: [batch=1, channels=2, samples=480000]
+    # 480000 samples @ 48kHz = 10 seconds of audio
+    dummy_audio = torch.randn(1, 2, 480000, dtype=torch.float32)
+
+    print(f"Dummy input shape: {dummy_audio.shape}")
+    print(f"Expected output shape: [1, 64, {480000 // 1920}] = [1, 64, {480000 // 1920}]")
+
+    # Test forward pass
+    with torch.no_grad():
+        test_out = wrapper(dummy_audio)
+    print(f"Test forward pass output shape: {test_out.shape}")
+
+    # Ensure output directory exists
+    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+
+    # Export
+    print(f"\nExporting to ONNX (opset {opset})...")
+    t0 = time.time()
+
+    dynamic_axes = {
+        "audio": {0: "batch", 2: "samples"},
+        "latents": {0: "batch", 2: "latent_frames"},
+    }
+
+    torch.onnx.export(
+        wrapper,
+        (dummy_audio,),
+        output_path,
+        opset_version=opset,
+        input_names=["audio"],
+        output_names=["latents"],
+        dynamic_axes=dynamic_axes,
+        do_constant_folding=True,
+        dynamo=False,  # Force legacy TorchScript exporter (dynamo hits cp1252 UnicodeError on Windows)
+    )
+
+    export_time = time.time() - t0
+    file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
+    print(f"Export complete in {export_time:.1f}s")
+    print(f"Output file: {output_path}")
+    print(f"File size: {file_size_mb:.1f} MB")
+
+    # Validate with onnx checker
+    import onnx
+
+    print("\nValidating ONNX model...")
+    model = onnx.load(output_path)
+    onnx.checker.check_model(model, full_check=True)
+    print("ONNX checker: PASSED")
+
+    # Print model info
+    graph = model.graph
+    print(f"\nModel inputs:")
+    for inp in graph.input:
+        shape = [d.dim_param or d.dim_value for d in inp.type.tensor_type.shape.dim]
+        print(f"  {inp.name}: {shape}")
+    print(f"Model outputs:")
+    for out in graph.output:
+        shape = [d.dim_param or d.dim_value for d in out.type.tensor_type.shape.dim]
+        print(f"  {out.name}: {shape}")
+
+    return output_path
+
+
+def validate_onnx(onnx_path: str, vae_path: str):
+    """Compare ONNX Runtime output against PyTorch output.
+
+    Args:
+        onnx_path: Path to the exported ONNX file.
+        vae_path: Path to the VAE checkpoint directory.
+    """
+    import onnxruntime as ort
+    from diffusers import AutoencoderOobleck
+
+    print("\n" + "=" * 60)
+    print("VALIDATION: Comparing ONNX vs PyTorch outputs")
+    print("=" * 60)
+
+    # Load PyTorch model
+    print("Loading PyTorch VAE...")
+    vae = AutoencoderOobleck.from_pretrained(vae_path)
+    vae.eval()
+    wrapper = VAEEncoderWrapper(vae)
+    wrapper.eval()
+    wrapper = wrapper.cpu()
+
+    # Create test input (shorter for speed: 96000 samples = 2 seconds)
+    test_audio = torch.randn(1, 2, 96000, dtype=torch.float32)
+    print(f"Test input shape: {test_audio.shape}")
+
+    # PyTorch inference
+    with torch.no_grad():
+        pt_output = wrapper(test_audio).numpy()
+    print(f"PyTorch output shape: {pt_output.shape}")
+
+    # ONNX Runtime inference
+    print("Loading ONNX model in onnxruntime...")
+    available_providers = ort.get_available_providers()
+    print(f"Available providers: {available_providers}")
+
+    # Use CUDA if available, else CPU
+    if "CUDAExecutionProvider" in available_providers:
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        print("Using: CUDAExecutionProvider")
+    else:
+        providers = ["CPUExecutionProvider"]
+        print("Using: CPUExecutionProvider (CUDA not available)")
+
+    sess = ort.InferenceSession(onnx_path, providers=providers)
+
+    ort_input = {"audio": test_audio.numpy()}
+    ort_output = sess.run(["latents"], ort_input)[0]
+    print(f"ONNX output shape: {ort_output.shape}")
+
+    # Compare
+    abs_diff = np.abs(pt_output - ort_output)
+    max_diff = abs_diff.max()
+    mean_diff = abs_diff.mean()
+    rel_diff = abs_diff / (np.abs(pt_output) + 1e-8)
+
+    print(f"\nDiff statistics:")
+    print(f"  Max absolute diff:  {max_diff:.6e}")
+    print(f"  Mean absolute diff: {mean_diff:.6e}")
+    print(f"  Max relative diff:  {rel_diff.max():.6e}")
+    print(f"  Mean relative diff: {rel_diff.mean():.6e}")
+
+    # Threshold check
+    if max_diff < 1e-4:
+        print("\n[PASS] VALIDATION PASSED: Outputs match within tolerance (1e-4)")
+    elif max_diff < 1e-3:
+        print("\n[WARN] VALIDATION WARNING: Small differences detected (< 1e-3)")
+        print("  This is acceptable for fp32 export, TRT fp16 will diverge more.")
+    else:
+        print(f"\n[FAIL] VALIDATION FAILED: Max diff {max_diff:.6e} exceeds tolerance")
+        sys.exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Export AutoencoderOobleck VAE encoder to ONNX"
+    )
+    parser.add_argument(
+        "--vae-path",
+        type=str,
+        required=True,
+        help="Path to VAE checkpoint directory (config.json + safetensors)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="Output path for the ONNX file",
+    )
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=18,
+        help="ONNX opset version (default: 18)",
+    )
+    parser.add_argument(
+        "--validate",
+        action="store_true",
+        help="Validate ONNX output against PyTorch output using onnxruntime",
+    )
+
+    args = parser.parse_args()
+
+    # Verify input exists
+    if not os.path.isdir(args.vae_path):
+        print(f"ERROR: VAE path not found: {args.vae_path}")
+        sys.exit(1)
+
+    config_path = os.path.join(args.vae_path, "config.json")
+    if not os.path.isfile(config_path):
+        print(f"ERROR: config.json not found in {args.vae_path}")
+        sys.exit(1)
+
+    # Export
+    onnx_path = export_vae_encoder(args.vae_path, args.output, args.opset)
+
+    # Validate
+    if args.validate:
+        validate_onnx(onnx_path, args.vae_path)
+
+    print("\nDone!")
+
+
+if __name__ == "__main__":
+    main()