280 lines
8.5 KiB
Python
280 lines
8.5 KiB
Python
#!/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()
|