export_vae.py

import torch
from diffusers import AutoencoderKL
import argparse
import os

def export_vae(model_path_or_id, output_path):
    print(f"Loading VAE from: {model_path_or_id}")
    if os.path.exists(model_path_or_id) and model_path_or_id.endswith(".safetensors"):
        # Load from local safetensors
        vae = AutoencoderKL.from_single_file(model_path_or_id)
    else:
        # Load from HuggingFace model ID or local directory
        vae = AutoencoderKL.from_pretrained(model_path_or_id)
    
    vae.eval()
    
    class VaeEncoderWrapper(torch.nn.Module):
        def __init__(self, vae):
            super().__init__()
            self.encoder = vae.encoder
            self.quant_conv = vae.quant_conv
            # SDXL scaling factor
            self.scaling_factor = 0.13025

        def forward(self, sample):
            # Preprocessed image sample is shape [1, 3, 1024, 1024]
            # VAE encoder output is [1, 8, 128, 128] representing mean and logvar
            h = self.encoder(sample)
            moments = self.quant_conv(h)
            # Take the mean (first 4 channels) as the latent representation
            mean = moments[:, :4, :, :]
            # Scale the latent
            latent = mean * self.scaling_factor
            return latent

    wrapper = VaeEncoderWrapper(vae)
    wrapper.eval()
    
    # Dummy input: 1024x1024 RGB image normalized to [-1, 1]
    dummy_input = torch.zeros(1, 3, 1024, 1024, dtype=torch.float32)
    
    print(f"Exporting VAE encoder wrapper to ONNX: {output_path}")
    torch.onnx.export(
        wrapper,
        dummy_input,
        output_path,
        export_params=True,
        opset_version=18,
        do_constant_folding=True,
        input_names=["image"],
        output_names=["latent"],
        dynamic_axes=None  # Static shapes for max TensorRT optimization
    )
    print("Export complete!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Export SDXL VAE Encoder to ONNX")
    parser.add_argument("--model", type=str, default="stabilityai/sdxl-vae", 
                        help="Hugging Face model ID, local folder, or path to sdxl_vae.safetensors")
    parser.add_argument("--output", type=str, default="vae_encoder.onnx", 
                        help="Path to output ONNX file")
    args = parser.parse_args()
    
    export_vae(args.model, args.output)