Initial release
This commit is contained in:
civ
@@ -0,0 +1,261 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user