upload_model.py

#!/usr/bin/env python3
"""
SkinSight Model Uploader for Hugging Face Hub.

Uploads the trained LoRA adapters so anyone can use the model easily.
Usage:
    huggingface-cli login
    python upload_model.py --repo-name "your-username/skinsight-vit"
"""

import argparse
import os
import sys
import torch
from pathlib import Path

sys.path.insert(0, "src")
from skinsight.training import TrainingConfig, create_model, load_checkpoint


def main():
    parser = argparse.ArgumentParser(description="Upload trained SkinSight model to Hugging Face")
    parser.add_argument("--repo-name", type=str, required=True, help="Target HF repo (e.g., username/skinsight-vit)")
    parser.add_argument("--checkpoint", type=str, default="checkpoints/best_model.pt", help="Path to best_model.pt")
    parser.add_argument("--commit-message", type=str, default="Upload trained SkinSight LoRA adapters", help="Commit message")
    args = parser.parse_args()

    checkpoint_path = Path(args.checkpoint)
    if not checkpoint_path.exists():
        print(f"Error: Checkpoint not found at {checkpoint_path}")
        print("Make sure training has completed or saved at least one best_model.pt")
        sys.exit(1)

    print(f"Loading checkpoint from {checkpoint_path}...")
    config = TrainingConfig()
    
    # Needs to be on CPU for export
    config.device = "cpu"
    
    # Create base model + LoRA
    model = create_model(config)
    
    # Load weights
    meta = load_checkpoint(model, checkpoint_path, device="cpu")
    print(f"Loaded model from Epoch {meta['epoch']} with Validation Accuracy: {meta['accuracy']:.4f}")

    print(f"\\nPushing LoRA adapters to HF Hub at {args.repo_name}...")
    try:
        # push_to_hub automatically handles creating the repo if it doesn't exist
        model.push_to_hub(
            args.repo_name,
            commit_message=args.commit_message,
            private=False, # Set to True if you want it hidden
        )
        print("\\n✅ Upload successful!")
        print(f"\\nYou and others can now use this model anywhere with:")
        print(f"  from peft import PeftModel, PeftConfig")
        print(f"  from transformers import AutoModelForImageClassification")
        print(f"  ")
        print(f"  # Load base model")
        print(f"  base_model = AutoModelForImageClassification.from_pretrained('google/vit-base-patch16-224')")
        print(f"  # Apply your trained SkinSight LoRA")
        print(f"  model = PeftModel.from_pretrained(base_model, '{args.repo_name}')")
        
    except Exception as e:
        print(f"\\n❌ Upload failed: {e}")
        print("\\nDid you log in? Run: huggingface-cli login")

if __name__ == "__main__":
    main()