262 lines
8.0 KiB
Python
262 lines
8.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Export AutoencoderOobleck VAE decoder to ONNX format.
|
|
|
|
Exports the decoder half of the VAE for use with TensorRT or ONNX Runtime.
|
|
The encoder is not needed for inference (we only decode latents → audio).
|
|
|
|
Tensor spec:
|
|
Input: "latents" [B, 64, T] (latent channels, latent frames @ 25Hz)
|
|
Output: "audio" [B, 2, samples] (stereo, samples = T * 1920 @ 48kHz)
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class VAEDecoderWrapper(nn.Module):
|
|
"""Wraps AutoencoderOobleck.decoder + post_quant_conv to extract .sample.
|
|
|
|
The raw decoder returns a DecoderOutput namedtuple, which torch.onnx.export
|
|
can't trace cleanly. This wrapper calls the decoder and returns the raw
|
|
tensor directly.
|
|
"""
|
|
|
|
def __init__(self, vae):
|
|
super().__init__()
|
|
self.decoder = vae.decoder
|
|
# post_quant_conv maps from latent space back to decoder input space
|
|
if hasattr(vae, "post_quant_conv") and vae.post_quant_conv is not None:
|
|
self.post_quant_conv = vae.post_quant_conv
|
|
else:
|
|
self.post_quant_conv = None
|
|
|
|
def forward(self, latents: torch.Tensor) -> torch.Tensor:
|
|
if self.post_quant_conv is not None:
|
|
latents = self.post_quant_conv(latents)
|
|
decoded = self.decoder(latents)
|
|
# decoder returns DecoderOutput with .sample attribute
|
|
if hasattr(decoded, "sample"):
|
|
return decoded.sample
|
|
return decoded
|
|
|
|
|
|
def export_vae(vae_path: str, output_path: str, opset: int = 18) -> str:
|
|
"""Export VAE decoder 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 = VAEDecoderWrapper(vae)
|
|
wrapper.eval()
|
|
|
|
# Move to CPU for export (fp32)
|
|
wrapper = wrapper.cpu()
|
|
|
|
# Create dummy input: [batch=1, latent_channels=64, latent_frames=250]
|
|
# 250 frames @ 25Hz = 10 seconds of audio
|
|
dummy_latents = torch.randn(1, 64, 250, dtype=torch.float32)
|
|
|
|
print(f"Dummy input shape: {dummy_latents.shape}")
|
|
print(f"Expected output shape: [1, 2, {250 * 1920}] = [1, 2, {250 * 1920}]")
|
|
|
|
# Test forward pass
|
|
with torch.no_grad():
|
|
test_out = wrapper(dummy_latents)
|
|
print(f"Test forward pass output shape: {test_out.shape}")
|
|
|
|
# Ensure output directory exists
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
|
|
# Export
|
|
print(f"\nExporting to ONNX (opset {opset})...")
|
|
t0 = time.time()
|
|
|
|
dynamic_axes = {
|
|
"latents": {0: "batch", 2: "latent_frames"},
|
|
"audio": {0: "batch", 2: "samples"},
|
|
}
|
|
|
|
torch.onnx.export(
|
|
wrapper,
|
|
(dummy_latents,),
|
|
output_path,
|
|
opset_version=opset,
|
|
input_names=["latents"],
|
|
output_names=["audio"],
|
|
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 = VAEDecoderWrapper(vae)
|
|
wrapper.eval()
|
|
wrapper = wrapper.cpu()
|
|
|
|
# Create test input (shorter for speed: 50 frames = 2 seconds)
|
|
test_latents = torch.randn(1, 64, 50, dtype=torch.float32)
|
|
print(f"Test input shape: {test_latents.shape}")
|
|
|
|
# PyTorch inference
|
|
with torch.no_grad():
|
|
pt_output = wrapper(test_latents).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 = {"latents": test_latents.numpy()}
|
|
ort_output = sess.run(["audio"], 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 decoder 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(args.vae_path, args.output, args.opset)
|
|
|
|
# Validate
|
|
if args.validate:
|
|
validate_onnx(onnx_path, args.vae_path)
|
|
|
|
print("\nDone!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|