import torch
import torch.nn as nn
from torchvision import models
import huggingface_hub
from huggingface_hub import hf_hub_download
import torchvision

model_repo_id = "CSSE416-final-project/faceRecogModel"
weight_file_id = "modelWeights100.bin"

def load_model(repo_id):
    # Download the model weights from the repo
    weights_path = hf_hub_download(repo_id=model_repo_id, filename=weight_file_id)

    # Initialize the ResNet-18 architecture
    model = torchvision.models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 100)  # Adjust for your task (e.g., 128 classes)

    # Load the model weights
    state_dict = torch.load(weights_path, map_location=torch.device("cpu"))
    model.load_state_dict(state_dict)
    model.eval()  # Set the model to evaluation mode
    return model