pytorch_model.bin

python
import os
import torch
from safetensors import safe_open
from safetensors.torch import save_file
import logging
import shutil
from datetime import datetime

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def backup_existing_file(file_path):
    if os.path.exists(file_path):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_path = f"{file_path}.backup_{timestamp}"
        shutil.copy2(file_path, backup_path)
        logging.info(f"Created backup of existing file: {backup_path}")

def convert_safetensors_to_pytorch(input_file, output_file):
    try:
        # Check if input file exists
        if not os.path.exists(input_file):
            raise FileNotFoundError(f"Input file {input_file} not found.")

        # Backup existing pytorch_model.bin if it exists
        backup_existing_file(output_file)

        # Load the safetensors file
        logging.info(f"Loading safetensors file: {input_file}")
        with safe_open(input_file, framework="pt", device="cpu") as f:
            state_dict = {key: f.get_tensor(key) for key in f.keys()}

        # Save as PyTorch bin file
        logging.info(f"Saving as PyTorch bin file: {output_file}")
        torch.save(state_dict, output_file)

        logging.info("Conversion complete.")
        logging.info(f"Created: {output_file}")

    except Exception as e:
        logging.error(f"An error occurred during conversion: {str(e)}")
        raise

if __name__ == "__main__":
    input_file = "maxcushion.safetensors"
    output_file = "pytorch_model.bin"

    try:
        convert_safetensors_to_pytorch(input_file, output_file)
    except Exception as e:
        logging.error(f"Conversion failed: {str(e)}")
    else:
        logging.info("Script executed successfully.")