fix: fixed crash in the server

features/image_classifier/model_loader.py

@@ -1,32 +1,43 @@
 import tensorflow as tf
 from tensorflow.keras.models import load_model as keras_load_model
 import os
-from huggingface_hub import snapshot_download  # fix import syntax here
+from huggingface_hub import snapshot_download
 import shutil
 
 # Constants
 REPO_ID = "can-org/AI-VS-HUMAN-IMAGE-classifier"
 MODEL_DIR = "./IMG_models"
-
 MODEL_PATH = os.path.join(MODEL_DIR, 'latest-my_cnn_model.h5')  # adjust path as needed
+
+_model_img = None  # global model variable
+
 def warmup():
     global _model_img
     if not os.path.exists(MODEL_DIR):
         download_model_Repo()
     _model_img = load_model()
+
 def download_model_Repo():
-    # fix typo: os.path.exists (not os.path.exist)
     if os.path.exists(MODEL_DIR):
         return
-    # download the repo snapshot from HF hub
     snapshot_path = snapshot_download(repo_id=REPO_ID)
     os.makedirs(MODEL_DIR, exist_ok=True)
-    # copy contents from snapshot_path to MODEL_DIR, allow existing dirs
     shutil.copytree(snapshot_path, MODEL_DIR, dirs_exist_ok=True)
 
 def load_model():
     if not os.path.exists(MODEL_DIR):
         download_model_Repo()
-    model = keras_load_model(MODEL_PATH)
+
+    # Check for GPU availability
+    gpus = tf.config.list_physical_devices('GPU')
+    if gpus:
+        # GPU is available, load model normally
+        print("GPU detected, loading model on GPU.")
+        model = keras_load_model(MODEL_PATH)
+    else:
+        # No GPU, force CPU usage
+        print("No GPU detected, forcing model loading on CPU.")
+        with tf.device('/CPU:0'):
+            model = keras_load_model(MODEL_PATH)
     return model