fix readme and config path

README.md

@@ -34,10 +34,9 @@ pip install -r requirements.txt
 python src/train.py
 ```
 By default, the model will be saved to `models/autoencoder_mnist.pth`.
-Then in `main.py` set
-`MODEL_PATH = "models/autoencoder_mnist.pth"`
-and in `train.py` set 
-`path = "models/autoencoder_mnist.pth"`
+
+Then in `model_config.py` set
+`model_path`.
 
 ## 🎛️ Run the Gradio Frontend
 ```

app.py

@@ -18,7 +18,7 @@ def request_pydantic_schema(_: type, handler: GetCoreSchemaHandler):
 Request.__get_pydantic_core_schema__ = request_pydantic_schema
 
 device = set_device()
-model = load_model(use_local=False)
+model = load_model()
 model.eval()
 
 transform = T.Compose([

app/main.py

@@ -22,7 +22,7 @@ app = FastAPI()
 
 device = set_device()
 log.info("Using device: %s", device)
-model = load_model(use_local=False)
+model = load_model()
 log.info("Model loaded. App is ready.")
 
 model.eval()

src/model_config.py

@@ -12,14 +12,12 @@ def set_device():
     return device
 
 # Load model
-def load_model(h: int = 32, use_local: bool = False):
-    if use_local:
-        model_path = "models/autoencoder_mnist.pth"
-    else:
-        model_path = hf_hub_download(
-            repo_id="dmtschulz/anomaly-detection-model",
-            filename="autoencoder_mnist.pth"
-        )
+def load_model(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)

src/train.py

@@ -9,7 +9,7 @@ from torch.utils.data import DataLoader, Subset
 from torchvision import transforms as T
 from torchvision.datasets import MNIST
 from torchvision.utils import save_image
-from model_config import set_device
+from src.model_config import set_device
 import logging
 
 logging.basicConfig(level=logging.INFO)