update model path

app.py

@@ -9,7 +9,7 @@ import matplotlib.pyplot as plt
 from src.train import load_model, loss_fn
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-model = load_model("./models/autoencoder_mnist.pth")
+model = load_model()
 model.eval()
 
 transform = T.Compose([

requirements.txt

@@ -11,4 +11,5 @@ plotly
 scikit-learn
 tqdm
 python-multipart
-torchvision
+torchvision
+huggingface_hub

src/train.py

@@ -15,6 +15,8 @@ import logging
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger(__name__)
 
+from huggingface_hub import hf_hub_download
+
 # Device configuration
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 log.info("Using device: %s", DEVICE)
@@ -85,10 +87,11 @@ def save_model(model, path):
 
 
 # Load model
-def load_model(path, h=32):
+def load_model(path=None, h=32):
+    if path is None or not os.path.exists(path):
+        path = hf_hub_download(repo_id="dmtschulz/anomaly-detection-model", filename="autoencoder_mnist.pth")
     model = Autoencoder(h).to(DEVICE)
     model.load_state_dict(torch.load(path, map_location=DEVICE))
-    model.to(DEVICE)
     model.eval()
     log.info("Model loaded from %s", path)
     return model