Update app.py

app.py

@@ -21,7 +21,7 @@ import traceback
 import os
 from smolagents import Tool
 from typing import Optional
-from weavemuse.models.notagen.inference import inference_patch, download_model_weights
+from weavemuse.models.notagen.inference import inference_patch, load_model_weights
 
 
 logging.basicConfig(level=logging.INFO)
@@ -142,7 +142,7 @@ class SimpleNotaGenTool(Tool):
         Args:
             repo_id: Repository ID on HuggingFace
         """        
-        download_model_weights()
+        load_model_weights(model_id=repo_id)
         logger.info("✅ NotaGen model weights downloaded successfully!")