rename_onnx_files.py

#!/usr/bin/env python3
"""
Script to rename ONNX model files and their external data files together.
Handles both the .onnx file and the .onnx.data file automatically.
"""

import argparse
import shutil
import sys
from pathlib import Path


def rename_onnx_files(old_name, new_name, dry_run=False):
    """
    Rename ONNX model file and its external data file.

    Args:
        old_name: Current model filename (without or with .onnx extension)
        new_name: New model filename (without or with .onnx extension)
        dry_run: If True, show what would be renamed without actually renaming
    """
    print("📝 ONNX FILE RENAMING TOOL")
    print("=" * 80)
    print()

    # Ensure proper extensions
    old_path = Path(old_name)
    new_path = Path(new_name)

    # Add .onnx extension if not present
    if old_path.suffix != ".onnx":
        old_path = old_path.with_suffix(".onnx")
    if new_path.suffix != ".onnx":
        new_path = new_path.with_suffix(".onnx")

    # Construct data file paths
    old_data_path = Path(str(old_path) + ".data")
    new_data_path = Path(str(new_path) + ".data")

    print(f"Source model:      {old_path}")
    print(f"Destination model: {new_path}")
    print()

    # Check if source files exist
    if not old_path.exists():
        print(f"❌ Error: Model file not found: {old_path}")
        return False

    has_data_file = old_data_path.exists()
    if has_data_file:
        print(f"✓ Found model file:      {old_path}")
        print(f"✓ Found data file:       {old_data_path}")
    else:
        print(f"✓ Found model file:      {old_path}")
        print("ℹ️  No external data file found (this is OK for small models)")

    print()

    # Check if destination files already exist
    if new_path.exists():
        print(f"⚠️  Warning: Destination file already exists: {new_path}")
        response = input("   Overwrite? (y/N): ").strip().lower()
        if response != "y":
            print("❌ Cancelled by user")
            return False

    if has_data_file and new_data_path.exists():
        print(f"⚠️  Warning: Destination data file already exists: {new_data_path}")
        response = input("   Overwrite? (y/N): ").strip().lower()
        if response != "y":
            print("❌ Cancelled by user")
            return False

    print()

    # Perform renaming
    if dry_run:
        print("🔍 DRY RUN - No files will be modified")
        print()
        print("Would rename:")
        print(f"  {old_path} → {new_path}")
        if has_data_file:
            print(f"  {old_data_path} → {new_data_path}")
        print()
        return True

    print("=" * 80)
    print("📦 RENAMING FILES")
    print("=" * 80)
    print()

    try:
        # Rename model file
        print(f"Renaming: {old_path.name} → {new_path.name}")
        shutil.move(str(old_path), str(new_path))
        print("✅ Model file renamed")

        # Rename data file if it exists
        if has_data_file:
            print(f"Renaming: {old_data_path.name} → {new_data_path.name}")
            shutil.move(str(old_data_path), str(new_data_path))
            print("✅ Data file renamed")

        print()

        # Update external data references in the ONNX model
        if has_data_file:
            print("Updating external data references in model...")
            try:
                import onnx

                model = onnx.load(str(new_path))

                # Update external data location references
                updates_made = 0
                for tensor in model.graph.initializer:
                    if (
                        tensor.HasField("data_location")
                        and tensor.data_location == onnx.TensorProto.EXTERNAL
                    ):
                        for entry in tensor.external_data:
                            if entry.key == "location":
                                old_location = entry.value
                                # Update any .onnx.data reference to use the new filename
                                if old_location.endswith(".onnx.data"):
                                    entry.value = new_data_path.name
                                    if updates_made == 0:  # Only print first update
                                        print(
                                            f"  Updated reference: {old_location} → {new_data_path.name}"
                                        )
                                    updates_made += 1

                if updates_made > 0:
                    print(f"  Total updates: {updates_made} tensor(s)")

                # Save the updated model
                onnx.save(model, str(new_path))
                print("✅ External data references updated")
            except ImportError:
                print(
                    "⚠️  Warning: onnx package not installed, skipping reference updates"
                )
                print("   The model may need manual reference updates")
            except Exception as e:
                print(f"⚠️  Warning: Could not update references: {e}")

        print()
        print("=" * 80)
        print("✨ RENAMING COMPLETE")
        print("=" * 80)
        print()
        print("Renamed files:")
        print(f"  ✓ {new_path}")
        if has_data_file:
            print(f"  ✓ {new_data_path}")

        return True

    except Exception as e:
        print(f"❌ Error during renaming: {e}")
        print()
        print("Attempting to rollback changes...")

        # Try to rollback
        try:
            if new_path.exists() and not old_path.exists():
                shutil.move(str(new_path), str(old_path))
                print("✓ Rolled back model file")

            if has_data_file and new_data_path.exists() and not old_data_path.exists():
                shutil.move(str(new_data_path), str(old_data_path))
                print("✓ Rolled back data file")

            print("✅ Rollback successful")
        except Exception as rollback_error:
            print(f"❌ Rollback failed: {rollback_error}")
            print("⚠️  WARNING: Files may be in an inconsistent state!")

        return False


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Rename ONNX model files and their external data files",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Rename model.onnx to new_model.onnx
  python rename_onnx_files.py model.onnx new_model.onnx

  # Works without .onnx extension too
  python rename_onnx_files.py model new_model

  # Dry run to see what would be renamed
  python rename_onnx_files.py model.onnx new_model.onnx --dry-run

This script will automatically rename both:
  - model.onnx → new_model.onnx
  - model.onnx.data → new_model.onnx.data (if it exists)

And update internal references to the data file.
        """,
    )

    parser.add_argument("old_name", type=str, help="Current model filename")

    parser.add_argument("new_name", type=str, help="New model filename")

    parser.add_argument(
        "--dry-run",
        "-n",
        action="store_true",
        help="Show what would be renamed without actually renaming",
    )

    args = parser.parse_args()

    success = rename_onnx_files(
        old_name=args.old_name, new_name=args.new_name, dry_run=args.dry_run
    )

    sys.exit(0 if success else 1)