import torch
import os
from torchvision.models import resnet50, ResNet50_Weights

def download_pretrained_model():
    try:
        # Load ResNet50 model with the best available weights
        print("Downloading ResNet50 model with ImageNet-1K weights...")
        model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        model.eval()
        
        # Save the model with safe loading
        print("Saving model to best_model.pth...")
        torch.save(model.state_dict(), 'best_model.pth', _use_new_zipfile_serialization=True)
        
        # Verify the file exists
        if os.path.exists('best_model.pth'):
            model_size = os.path.getsize('best_model.pth') / (1024 * 1024)  # Size in MB
            print(f"Model saved successfully! Size: {model_size:.2f} MB")
        else:
            print("Error: Model file was not created")
            
    except Exception as e:
        print(f"An error occurred: {str(e)}")

if __name__ == "__main__":
    download_pretrained_model()