Update app/main.py

app/main.py

@@ -68,31 +68,10 @@ logger = logging.getLogger("deepguard_api")
 if not logger.handlers:
     logging.basicConfig(level=logging.INFO)
 
-# Try to load an IMAGE model if present
-# Default to the trained hybrid model file name; can be overridden via MODEL_PATH env var.
-# Always resolve the path relative to the project root so that cwd doesn't matter.
-_raw_model_path = os.environ.get("MODEL_PATH", "hybrid_deepfake_model.h5")
-if os.path.isabs(_raw_model_path):
-    MODEL_PATH = _raw_model_path
-else:
-    MODEL_PATH = os.path.join(ROOT_DIR, _raw_model_path)
-
-_loaded_model = None
-_model_load_error: Optional[str] = None
-try:
-    if os.path.exists(MODEL_PATH):
-        logger.info(f"Loading IMAGE model from: {MODEL_PATH}")
-        _loaded_model = model_module.load_model_from_checkpoint(MODEL_PATH)
-        logger.info(f"Image model loaded successfully: name={getattr(_loaded_model, 'name', 'unknown')}")
-    else:
-        _model_load_error = f"Model file not found at path: {MODEL_PATH}"
-        logger.error(_model_load_error)
-except Exception as e:
-    _loaded_model = None
-    _model_load_error = f"Error loading model from '{MODEL_PATH}': {e}"
-    logger.exception(_model_load_error)
+# ❌ IMAGE MODEL REMOVED - Only Video Model for Memory Optimization
+# Image model loading disabled to save memory
 
-# Optional VIDEO model + feature extractor
+# ✅ VIDEO MODEL ONLY - Optimized Loading
 _raw_video_model_path = os.environ.get("VIDEO_MODEL_PATH", "video_deepfake_model.h5")
 if os.path.isabs(_raw_video_model_path):
     VIDEO_MODEL_PATH = _raw_video_model_path
@@ -106,12 +85,32 @@ _video_feature_extractor = None
 try:
     if os.path.exists(VIDEO_MODEL_PATH):
         logger.info(f"Loading VIDEO model from: {VIDEO_MODEL_PATH}")
+        
+        # Memory optimization before loading
+        import gc
+        import tensorflow as tf
+        gc.collect()
+        tf.keras.backend.clear_session()
+        
+        # Load video model with optimizations
         _video_model = model_module.load_model_from_checkpoint(VIDEO_MODEL_PATH)
+        
+        # Optimize for inference
+        _video_model.trainable = False
+        for layer in _video_model.layers:
+            layer.trainable = False
+        
+        # Build feature extractor
         _video_feature_extractor = model_module.build_video_feature_extractor()
-        logger.info(f"Video model loaded successfully: name={getattr(_video_model, 'name', 'unknown')}")
+        
+        logger.info(f"Video model loaded and optimized successfully!")
+        
+        # Clear memory after loading
+        gc.collect()
+        
     else:
         _video_model_error = f"Video model file not found at path: {VIDEO_MODEL_PATH}"
-        logger.warning(_video_model_error)
+        logger.error(_video_model_error)
 except Exception as e:
     _video_model = None
     _video_feature_extractor = None
@@ -129,30 +128,31 @@ class PredictResponse(BaseModel):
 def health():
     """
     Simple health check.
-    Returns whether the model is loaded and exposes basic debug info.
+    Returns whether the video model is loaded and exposes basic debug info.
     """
     return {
         "status": "ok",
-        "model_loaded": _loaded_model is not None,
-        "model_path": MODEL_PATH,
-        "model_error": _model_load_error,
+        "image_model_enabled": False,  # ❌ Disabled
         "video_model_loaded": _video_model is not None,
         "video_model_path": VIDEO_MODEL_PATH,
         "video_model_error": _video_model_error,
+        "memory_optimized": True,
+        "model_type": "video_only"
     }
 
 @app.post("/predict", response_model=PredictResponse)
 async def predict(file: UploadFile = File(...)):
     """
-    Accepts an uploaded image file, returns prediction.
+    Accepts an uploaded video file, returns prediction.
+    ❌ Image processing disabled - Video only for memory optimization.
     """
-    global _loaded_model, _video_model, _video_feature_extractor
-    if _loaded_model is None:
-        # Do NOT silently fall back to an untrained model.
-        # This would give meaningless predictions (often always one class).
-        detail = _model_load_error or (
-            "Model is not loaded. Ensure MODEL_PATH points to a valid trained model "
-            "(.h5 file) and restart the API server."
+    global _video_model, _video_feature_extractor
+    
+    # ❌ IMAGE PROCESSING DISABLED
+    if _video_model is None:
+        detail = _video_model_error or (
+            "Video model is not loaded. Ensure VIDEO_MODEL_PATH points to a valid "
+            "trained video model (.h5) and restart the API server."
         )
         raise HTTPException(status_code=500, detail=detail)
 
@@ -162,54 +162,50 @@ async def predict(file: UploadFile = File(...)):
 
     video_exts = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".webm"}
 
-    # ----------------------- VIDEO INPUT -----------------------
-    if ext in video_exts:
-        if _video_model is None or _video_feature_extractor is None:
-            detail = _video_model_error or (
-                "Video model is not loaded. Ensure VIDEO_MODEL_PATH points to a valid "
-                "trained video model (.h5) and that TensorFlow/OpenCV are installed."
-            )
-            raise HTTPException(status_code=500, detail=detail)
-
-        contents = await file.read()
-        tmp_path = None
-        try:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
-                tmp_file.write(contents)
-                tmp_path = tmp_file.name
-
-            # Use unified helper that supports both images & videos
-            result = model_module.predict_from_input_unified(
-                _loaded_model,
-                tmp_path,
-                input_type="video",
-                video_model=_video_model,
-                feature_extractor=_video_feature_extractor,
-            )
-            return JSONResponse(content=result)
-        except HTTPException:
-            raise
-        except Exception as e:
-            raise HTTPException(status_code=400, detail=f"Could not process video: {e}")
-        finally:
-            if tmp_path and os.path.exists(tmp_path):
-                try:
-                    os.unlink(tmp_path)
-                except Exception:
-                    pass
-
-    # ----------------------- IMAGE INPUT -----------------------
+    # ❌ ONLY VIDEO PROCESSING ALLOWED
+    if ext not in video_exts:
+        raise HTTPException(
+            status_code=400, 
+            detail=f"❌ Image processing disabled. Please upload a video file. Supported formats: {', '.join(video_exts)}"
+        )
+
+    # ----------------------- VIDEO PROCESSING ONLY -----------------------
     contents = await file.read()
+    tmp_path = None
     try:
-        from PIL import Image
-
-        img = Image.open(io.BytesIO(contents)).convert("RGB")
-        arr = np.asarray(img.resize((224, 224)))
+        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
+            tmp_file.write(contents)
+            tmp_path = tmp_file.name
+
+        # Memory optimization before prediction
+        import gc
+        import tensorflow as tf
+        gc.collect()
+        
+        # Use video model for prediction
+        result = model_module.predict_from_input_unified(
+            None,  # No image model
+            tmp_path,
+            input_type="video",
+            video_model=_video_model,
+            feature_extractor=_video_feature_extractor,
+        )
+        
+        # Clear memory after prediction
+        gc.collect()
+        
+        return JSONResponse(content=result)
+        
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Could not read image: {e}")
-
-    result = model_module.predict_from_input(_loaded_model, arr)
-    return JSONResponse(content=result)
+        raise HTTPException(status_code=400, detail=f"Could not process video: {e}")
+    finally:
+        if tmp_path and os.path.exists(tmp_path):
+            try:
+                os.unlink(tmp_path)
+            except Exception:
+                pass
 
 if __name__ == "__main__":
     # Run using the already imported `app` instance instead of a string path