convert.py

#!/usr/bin/env python3
"""
Universal Model Converter for DeepDream-MLX.
Converts PyTorch (.pth) and Torch7 (.t7) models to MLX (.npz).
Also supports auto-downloading standard Places365 models.
Defaults to float16 for optimal performance on Apple Silicon.
"""

import os
import argparse
import glob
import numpy as np
import torch
import torchvision.models as models
from torch.hub import download_url_to_file

# Optional Torchfile for .t7 support
try:
    import torchfile
except ImportError:
    torchfile = None

# --- Configuration ---
PLACES365_URLS = {
    "alexnet": "http://places2.csail.mit.edu/models_places365/alexnet_places365.pth.tar",
    "resnet50": "http://places2.csail.mit.edu/models_places365/resnet50_places365.pth.tar",
    "vgg16": "http://places2.csail.mit.edu/models_places365/vgg16_places365.pth.tar",
    "googlenet": "http://places2.csail.mit.edu/models_places365/googlenet_places365.pth.tar"
}

# --- Helper Functions ---

def convert_tensor(tensor, target_dtype=np.float16):
    """Converts a tensor/array to the target numpy dtype."""
    if isinstance(tensor, torch.Tensor):
        return tensor.cpu().detach().numpy().astype(target_dtype)
    elif isinstance(tensor, np.ndarray):
        return tensor.astype(target_dtype)
    else:
        return np.array(tensor).astype(target_dtype)

def clean_state_dict(state_dict):
    """
    Flattens the state dictionary and removes common prefix artifacts 
    like 'module.' from DataParallel wrapping.
    """
    new_dict = {}
    for k, v in state_dict.items():
        # Remove 'module.' anywhere in the key
        name = k.replace("module.", "")
        new_dict[name] = convert_tensor(v)
    return new_dict

def get_places365_model_skeleton(arch):
    """Returns a standard PyTorch model structure for Places365."""
    if arch == "alexnet":
        return models.alexnet(num_classes=365)
    elif arch == "resnet50":
        return models.resnet50(num_classes=365)
    elif arch == "vgg16":
        return models.vgg16(num_classes=365)
    elif arch == "googlenet":
        return models.googlenet(num_classes=365, aux_logits=False)
    else:
        raise ValueError(f"Unknown architecture: {arch}")

# --- Conversion Logic ---

def convert_torch7(filepath, target_dir):
    if torchfile is None:
        print(f"⚠️  Skipping {filepath}: 'torchfile' not installed. Run `pip install torchfile`.")
        return

    print(f"Processing Torch7 file: {filepath}")
    try:
        model_obj = torchfile.load(filepath)
        converted_state = {}

        def extract_layers(layer, prefix=""):
            if hasattr(layer, 'weight') and layer.weight is not None:
                converted_state[f"{prefix}.weight"] = convert_tensor(layer.weight)
            if hasattr(layer, 'bias') and layer.bias is not None:
                converted_state[f"{prefix}.bias"] = convert_tensor(layer.bias)
            
            if hasattr(layer, 'modules') and layer.modules:
                for i, sublayer in enumerate(layer.modules):
                    # 0-based indexing for compatibility
                    next_prefix = f"{prefix}.{i}" if prefix else f"{i}"
                    extract_layers(sublayer, next_prefix)
        
        extract_layers(model_obj)
        
        if not converted_state:
            print(f"❌ No weights found in {filepath}.")
            return

        name_base = os.path.splitext(os.path.basename(filepath))[0]
        out_path = os.path.join(target_dir, f"{name_base}_t7_mlx.npz")
        np.savez(out_path, **converted_state)
        print(f"✅ Saved {out_path} ({len(converted_state)} tensors)")

    except Exception as e:
        print(f"❌ Failed to convert {filepath}: {e}")

def convert_pytorch(filepath, target_dir):
    print(f"Processing PyTorch file: {filepath}")
    try:
        checkpoint = torch.load(filepath, map_location="cpu")
        
        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
            state_dict = checkpoint['state_dict']
        elif isinstance(checkpoint, dict):
            state_dict = checkpoint
        else:
            print(f"❌ Unknown checkpoint format in {filepath}")
            return

        clean_dict = clean_state_dict(state_dict)
        
        name_base = os.path.splitext(os.path.basename(filepath))[0]
        # Avoid double extension if file was .pth.tar
        if name_base.endswith(".pth"):
            name_base = os.path.splitext(name_base)[0]
            
        out_path = os.path.join(target_dir, f"{name_base}_mlx.npz")
        np.savez(out_path, **clean_dict)
        
        size_mb = os.path.getsize(out_path) / (1024*1024)
        print(f"✅ Saved {out_path} ({size_mb:.1f} MB)")

    except Exception as e:
        print(f"❌ Failed to convert {filepath}: {e}")

def download_and_convert_places365(arch, download_dir, target_dir):
    url = PLACES365_URLS.get(arch)
    if not url:
        print(f"No URL for {arch}")
        return

    filename = os.path.join(download_dir, os.path.basename(url))
    
    # 1. Download
    if not os.path.exists(filename):
        print(f"Downloading {arch} from {url}...")
        try:
            download_url_to_file(url, filename)
        except Exception as e:
            print(f"Download failed: {e}")
            return
    else:
        print(f"Found cached {filename}")

    # 2. Load into standard Skeleton (ensures structural correctness)
    print(f"Loading {arch} into PyTorch structure...")
    try:
        model = get_places365_model_skeleton(arch)
        checkpoint = torch.load(filename, map_location="cpu")
        state_dict = checkpoint['state_dict'] if 'state_dict' in checkpoint else checkpoint
        
        # Robust Load
        new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
        try:
            model.load_state_dict(new_state_dict, strict=True)
        except:
            model.load_state_dict(new_state_dict, strict=False)
            
        # 3. Export
        model.eval()
        final_dict = clean_state_dict(model.state_dict())
        out_path = os.path.join(target_dir, f"{arch}_places365_mlx.npz")
        np.savez(out_path, **final_dict)
        print(f"✅ Saved {out_path}")

    except Exception as e:
        print(f"Failed to process {arch}: {e}")

# --- Main CLI ---

def main():
    parser = argparse.ArgumentParser(description="DeepDream-MLX Model Converter")
    parser.add_argument("--scan", default="toConvert", help="Directory to scan for local files")
    parser.add_argument("--download", choices=["alexnet", "resnet50", "vgg16", "googlenet", "all"], 
                        help="Download and convert specific Places365 models")
    parser.add_argument("--dest", default=".", help="Output directory for .npz files")
    args = parser.parse_args()

    if not os.path.exists(args.dest):
        os.makedirs(args.dest)

    # 1. Handle Downloads
    if args.download:
        if not os.path.exists(args.scan):
            os.makedirs(args.scan)
            
        targets = ["alexnet", "resnet50", "vgg16", "googlenet"] if args.download == "all" else [args.download]
        for t in targets:
            download_and_convert_places365(t, args.scan, args.dest)

    # 2. Handle Local Scan
    if os.path.exists(args.scan):
        print(f"\nScanning '{args.scan}' for local models...")
        files = glob.glob(os.path.join(args.scan, "*"))
        for f in files:
            if os.path.isdir(f): continue
            ext = os.path.splitext(f)[1].lower()
            
            if ext == ".t7":
                convert_torch7(f, args.dest)
            elif ext in [".pth", ".pt", ".tar", ".pkl"]:
                # If it looks like a downloaded places file we already processed, skip to avoid duplication
                # heuristic: if we just downloaded it. 
                convert_pytorch(f, args.dest)
            elif ext in [".caffemodel"]:
                print(f"⚠️  Skipping Caffe model {os.path.basename(f)} (Not supported)")

if __name__ == "__main__":
    main()