Upload app.py with huggingface_hub

app.py

@@ -5,6 +5,7 @@ from flask_socketio import SocketIO
 import joblib
 import torch
 import numpy as np
+from huggingface_hub import hf_hub_download
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -13,69 +14,38 @@ logger = logging.getLogger(__name__)
 app = Flask(__name__)
 socketio = SocketIO(app, cors_allowed_origins="*")
 
-# Create necessary directories
-def create_directories():
-    directories = [
-        'heart/models',
-        'models'
-    ]
-    for directory in directories:
-        try:
-            os.makedirs(directory, exist_ok=True)
-            logger.info(f"Created directory: {directory}")
-        except Exception as e:
-            logger.warning(f"Failed to create directory {directory}: {str(e)}")
-
 # Load models
 def load_models():
     global heart_model, autoencoder
     heart_model = None
     autoencoder = None
     
-    # Create directories first
-    create_directories()
-    
-    # Define possible model paths
-    model_paths = [
-        os.path.join('heart', 'models', 'heart_model.joblib'),
-        os.path.join(os.path.dirname(__file__), 'heart', 'models', 'heart_model.joblib'),
-        os.path.join('/app', 'heart', 'models', 'heart_model.joblib')
-    ]
-    
-    # Try loading heart model
-    for path in model_paths:
-        try:
-            logger.info(f"Attempting to load heart model from {path}")
-            if os.path.exists(path):
-                heart_model = joblib.load(path)
-                logger.info("Heart model loaded successfully")
-                break
-            else:
-                logger.warning(f"Model file not found at {path}")
-        except Exception as e:
-            logger.warning(f"Failed to load heart model from {path}: {str(e)}")
-            continue
-    
-    # Try loading autoencoder
-    autoencoder_paths = [
-        os.path.join('models', 'best_model.pth'),
-        os.path.join(os.path.dirname(__file__), 'models', 'best_model.pth'),
-        os.path.join('/app', 'models', 'best_model.pth')
-    ]
+    try:
+        # Download and load heart model
+        logger.info("Downloading heart model from Hugging Face Hub...")
+        heart_model_path = hf_hub_download(
+            repo_id="leo861/app",
+            filename="heart/models/heart_model.joblib",
+            cache_dir="models"
+        )
+        heart_model = joblib.load(heart_model_path)
+        logger.info("Heart model loaded successfully")
+    except Exception as e:
+        logger.error(f"Failed to load heart model: {str(e)}")
     
-    for path in autoencoder_paths:
-        try:
-            logger.info(f"Attempting to load autoencoder from {path}")
-            if os.path.exists(path):
-                autoencoder = torch.load(path)
-                autoencoder.eval()
-                logger.info("Autoencoder model loaded successfully")
-                break
-            else:
-                logger.warning(f"Model file not found at {path}")
-        except Exception as e:
-            logger.warning(f"Failed to load autoencoder from {path}: {str(e)}")
-            continue
+    try:
+        # Download and load autoencoder
+        logger.info("Downloading autoencoder from Hugging Face Hub...")
+        autoencoder_path = hf_hub_download(
+            repo_id="leo861/app",
+            filename="models/best_model.pth",
+            cache_dir="models"
+        )
+        autoencoder = torch.load(autoencoder_path)
+        autoencoder.eval()
+        logger.info("Autoencoder model loaded successfully")
+    except Exception as e:
+        logger.error(f"Failed to load autoencoder: {str(e)}")
 
 # Load models on startup
 logger.info("Loading trained models...")