src/model_config.py

# src/model_config.py

from src.train import Autoencoder
import torch
from huggingface_hub import hf_hub_download
import logging

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

# Device configuration
def set_device():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    return device

# Load model
def load_model(device, h: int = 32):
    # model_path = "models/autoencoder_mnist.pth" # Use local model
    model_path = hf_hub_download(repo_id="dmtschulz/anomaly-detection-model", filename="autoencoder_mnist.pth")

    model = Autoencoder(h).to(device)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()
    log.info("Model loaded from %s", model_path)
    return model