VideoX-Fun/slim_all_onnx.py

#!/usr/bin/env python3
"""
Script to run onnxslim on all ONNX files in transformers_body_only_1728_992_split_onnx directory
"""
import os
import shutil
import subprocess
from pathlib import Path

def main():
    # Directory containing ONNX files
    onnx_dir = Path("transformers_body_only_1728_992_split_onnx")
    
    if not onnx_dir.exists():
        print(f"Error: Directory {onnx_dir} not found!")
        return
    
    # Get all .onnx files
    onnx_files = sorted(onnx_dir.glob("*.onnx"))
    
    if not onnx_files:
        print(f"No ONNX files found in {onnx_dir}")
        return
    
    print(f"Found {len(onnx_files)} ONNX files")
    print("=" * 80)
    
    success_count = 0
    failed_files = []
    
    for idx, onnx_file in enumerate(onnx_files, 1):
        print(f"\n[{idx}/{len(onnx_files)}] Processing: {onnx_file.name}")
        print(f"File size: {onnx_file.stat().st_size / (1024**3):.2f} GB")
        
        # Create backup before processing
        backup_file = onnx_file.with_suffix(".onnx.backup")
        shutil.copy2(onnx_file, backup_file)
        
        try:
            # Run onnxslim: onnxslim input.onnx output.onnx (in-place)
            cmd = ["onnxslim", str(onnx_file), str(onnx_file)]
            print(f"Running: {' '.join(cmd)}")
            
            result = subprocess.run(
                cmd,
                check=True,
                capture_output=True,
                text=True
            )
            
            print(f"✓ Success: {onnx_file.name}")
            if result.stdout:
                print(f"  Output: {result.stdout.strip()}")
            success_count += 1
            
            # Remove backup on success
            backup_file.unlink()
            
        except subprocess.CalledProcessError as e:
            print(f"✗ Failed: {onnx_file.name}")
            print(f"  Error: {e.stderr}")
            failed_files.append(onnx_file.name)
            
            # Restore from backup
            if backup_file.exists():
                shutil.copy2(backup_file, onnx_file)
                backup_file.unlink()
                print(f"  Restored from backup")
                
        except Exception as e:
            print(f"✗ Failed: {onnx_file.name}")
            print(f"  Error: {str(e)}")
            failed_files.append(onnx_file.name)
            
            # Restore from backup
            if backup_file.exists():
                shutil.copy2(backup_file, onnx_file)
                backup_file.unlink()
                print(f"  Restored from backup")
    
    # Summary
    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    print(f"Total files: {len(onnx_files)}")
    print(f"Successful: {success_count}")
    print(f"Failed: {len(failed_files)}")
    
    if failed_files:
        print("\nFailed files:")
        for file in failed_files:
            print(f"  - {file}")

if __name__ == "__main__":
    main()