import base64 import gc import io import logging import os import platform import sys import threading import time from contextlib import asynccontextmanager from typing import Any, Dict, Optional import torch import uvicorn from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from PIL import Image, ImageFile from pydantic import BaseModel, Field # Allow loading of truncated images ImageFile.LOAD_TRUNCATED_IMAGES = True # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", handlers=[logging.StreamHandler(sys.stdout)], ) logger = logging.getLogger(__name__) # Model-specific imports from the cloned repository # PYTHONPATH is set in Dockerfile to find 'model_code.src' or 'src' directly try: from lightning.fabric import Fabric from src.config import Config from src.model.dfdet import DeepfakeDetectionModel except ImportError as e: logger.error(f"Error importing model-specific modules: {e}. Check PYTHONPATH.") # Attempt to add path manually if Docker's PYTHONPATH isn't picked up as expected in some envs current_dir_for_app = os.path.dirname(os.path.abspath(__file__)) model_code_path = os.path.join(current_dir_for_app, "model_code") if model_code_path not in sys.path: sys.path.insert(0, model_code_path) logger.info(f"Added {model_code_path} to sys.path") # Retry imports from lightning.fabric import Fabric from src.config import Config from src.model.dfdet import DeepfakeDetectionModel @asynccontextmanager async def lifespan(app: FastAPI): # Code to run on startup if PRELOAD_MODEL: logger.info("Preloading model at startup (PRELOAD_MODEL=true)") try: load_model_internal() except Exception as e: logger.error( f"Fatal error during model preloading: {e}. Service might not function." ) else: logger.info("Model will be loaded on first request (PRELOAD_MODEL=false).") # Start a background timer to check for model unloading if not preloading timer_thread = None if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: def periodic_unload_check(): unload_model_if_idle() # Reschedule the check only if the model is still loaded or the app is running # A more robust way would be to manage the timer cancellation in the 'yield' part if model is not None and not PRELOAD_MODEL: # Check if model still exists timer_thread = threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check) timer_thread.start() # Initial call after a short delay timer_thread = threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check) timer_thread.start() logger.info(f"Model idle check timer started (interval: {MODEL_TIMEOUT / 2}s).") yield # Application runs here # Code to run on shutdown logger.info("Shutting down Yermandy CLIP detection service.") if timer_thread and timer_thread.is_alive(): timer_thread.cancel() logger.info("Cancelled model idle check timer.") # Ensure model is unloaded on shutdown if it exists global model, preprocessing_fn, fabric if model is not None: logger.info("Unloading model on shutdown.") del model, preprocessing_fn, fabric model, preprocessing_fn, fabric = None, None, None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Initialize FastAPI app with lifespan app = FastAPI( title="Yermandy CLIP Deepfake Detection Model Service", description="Service for detecting deepfake images using the Yermandy CLIP-based model.", version="1.0.0", lifespan=lifespan, # Add the lifespan manager here ) # CORS setup app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --- Model Configuration & Globals --- MODEL_NAME = "yermandy_clip_detection" MODEL_PATH = os.environ.get("MODEL_PATH", "model_code/weights/model.ckpt") def _get_device(): """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() if override == "cpu": return torch.device("cpu") if override == "cuda" and torch.cuda.is_available(): return torch.device("cuda") if ( override == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): return torch.device("mps") if override: pass # Invalid override, fall through to auto-detect if ( platform.system() == "Darwin" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): return torch.device("mps") if torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") DEVICE = _get_device() if DEVICE.type == "cuda": torch.backends.cudnn.benchmark = True torch.set_float32_matmul_precision("high") if DEVICE.type == "cuda": logger.info( "Device: cuda (%s, %.1f GB VRAM)", torch.cuda.get_device_name(0), torch.cuda.get_device_properties(0).total_memory / 1024**3, ) else: logger.warning( "Device: %s (no CUDA available -- check nvidia-container-toolkit)", DEVICE, ) PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) # Default 10 minutes # Global variables for the model and related components model: Optional[DeepfakeDetectionModel] = None preprocessing_fn: Optional[callable] = None fabric: Optional[Fabric] = None model_lock = threading.Lock() last_used_time = 0 from pydantic import BaseModel, Field # ConfigDict not needed here if no extra config class ImageInput(BaseModel): image_data: str = Field( ..., description="Base64 encoded image string" ) # Renamed field threshold: Optional[float] = Field( 0.35, ge=0.0, le=1.0, description="Classification threshold" ) def load_model_internal(): """Loads the deepfake detection model and its components.""" global model, preprocessing_fn, fabric, last_used_time with model_lock: if model is not None: # Check again after acquiring lock last_used_time = time.time() return logger.info(f"Loading model from {MODEL_PATH} onto {DEVICE}...") if not os.path.exists(MODEL_PATH): logger.error(f"Model weights not found at {MODEL_PATH}") raise FileNotFoundError(f"Model weights not found at {MODEL_PATH}") try: ckpt = torch.load(MODEL_PATH, map_location="cpu") # Load to CPU first # Ensure hyper_parameters are available if "hyper_parameters" not in ckpt: logger.error( "Checkpoint does not contain 'hyper_parameters'. Cannot initialize model." ) raise ValueError("Invalid checkpoint: missing 'hyper_parameters'") model_config = Config(**ckpt["hyper_parameters"]) _model = DeepfakeDetectionModel(model_config) _model.load_state_dict(ckpt["state_dict"]) _model.eval() # Set to evaluation mode _model.to(DEVICE) # Move to CPU _preprocessing_fn = _model.get_preprocessing() # For CPU, Fabric precision should be "32-true" or similar, not from checkpoint if it's like "bf16-mixed" # The original inference.py loads precision from ckpt["hyper_parameters"]["precision"] # For CPU, it's safer to override this. _accelerator = "cuda" if DEVICE.type == "cuda" else "cpu" _fabric = Fabric(accelerator=_accelerator, devices=1, precision="32-true") # fabric.launch() is not needed here as we are not in a distributed script _model = _fabric.setup_module(_model) # Prepare model with Fabric # Assign to global variables model = _model preprocessing_fn = _preprocessing_fn fabric = _fabric last_used_time = time.time() logger.info("Model loaded successfully.") except Exception as e: logger.exception(f"Failed to load model: {e}") # Reset globals if loading failed model = None preprocessing_fn = None fabric = None raise finally: gc.collect() def ensure_model_loaded(): """Ensures the model is loaded, loading it if necessary.""" global last_used_time if model is None: load_model_internal() else: last_used_time = time.time() # Update last used time if already loaded def unload_model_if_idle(): """Unloads the model if it has been idle for longer than MODEL_TIMEOUT.""" global model, preprocessing_fn, fabric if model is None or PRELOAD_MODEL: # Don't unload if preloaded or already unloaded return with model_lock: if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): logger.info( f"Unloading model due to inactivity (timeout: {MODEL_TIMEOUT}s)." ) del model del preprocessing_fn del fabric model = None preprocessing_fn = None fabric = None gc.collect() logger.info("Model unloaded and memory cleared.") # --- FastAPI Endpoints --- @app.on_event("startup") async def startup_event(): if PRELOAD_MODEL: logger.info("Preloading model at startup (PRELOAD_MODEL=true)") try: load_model_internal() except Exception as e: logger.error( f"Fatal error during model preloading: {e}. Service might not function." ) else: logger.info("Model will be loaded on first request (PRELOAD_MODEL=false).") # Start a background timer to check for model unloading if not preloading if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: def periodic_unload_check(): unload_model_if_idle() if ( model is None and not PRELOAD_MODEL ): # if model got unloaded and we are not preloading, stop timer return # Reschedule the check threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check).start() # Initial call after a short delay threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check).start() logger.info(f"Model idle check timer started (interval: {MODEL_TIMEOUT / 2}s).") @app.get("/") async def root(): return { "model_name": MODEL_NAME, "description": "Deepfake detection model based on Yermandy's CLIP work.", "model_path": MODEL_PATH, "device": str(DEVICE), "model_loaded": model is not None, "lazy_loading_enabled": not PRELOAD_MODEL, "model_timeout_seconds": ( MODEL_TIMEOUT if not PRELOAD_MODEL else "N/A (preloaded)" ), } def _gpu_health_info() -> dict: """Return GPU metrics for the health endpoint.""" if torch.cuda.is_available() and DEVICE.type == "cuda": return { "gpu_name": torch.cuda.get_device_name(0), "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), "vram_total_mb": round( torch.cuda.get_device_properties(0).total_memory / 1024**2 ), } return {} @app.get("/health") async def health(): model_file_exists = os.path.exists(MODEL_PATH) status_message = "healthy" if not model_file_exists: status_message = "error_missing_weights" # Could add a quick inference test here if model is loaded for a more thorough check return { "status": status_message, "model_name": MODEL_NAME, "device": str(DEVICE), "model_weights_found": model_file_exists, "model_loaded": model is not None, **_gpu_health_info(), } @app.post("/unload") async def unload_model(): """Endpoint to manually unload the model""" global model if model is None: return {"status": "not_loaded", "message": "Model is not currently loaded"} logger.info("Manually unloading model") # Delete model and clear memory del model model = None # Check if cuda is available instead of using undefined USE_GPU if torch.cuda.is_available(): # Modified line torch.cuda.empty_cache() gc.collect() logger.info("Model unloaded and memory cleared") return {"status": "success", "message": "Model unloaded successfully"} @app.post("/predict", response_model=Dict[str, Any]) async def predict(image_input: ImageInput): try: ensure_model_loaded() # This will load the model if it's not already loaded if model is None or preprocessing_fn is None or fabric is None: logger.error("Model components are not available for prediction.") raise HTTPException( status_code=503, detail="Model is not loaded or failed to load." ) start_time = time.time() # Decode base64 image try: image_bytes = base64.b64decode(image_input.image_data) pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") except Exception as e: logger.error(f"Invalid image data: {e}") raise HTTPException(status_code=400, detail=f"Invalid image data: {e}") # Preprocess the image # The preprocessing_fn from the model expects a PIL image and returns a tensor. # The model's forward pass expects a batch, so unsqueeze(0). image_tensor = preprocessing_fn(pil_image).unsqueeze(0) # Move tensor to the device Fabric prepared the model for (CPU in this case) # The dtype should be handled by Fabric setup_module or .to(DEVICE) image_tensor = image_tensor.to( DEVICE ) # fabric.to_device(image_tensor) could also be used if fabric object is accessible # Perform inference with torch.no_grad(): output = model(image_tensor) # model is already setup by fabric # Process output # output.logits_labels is what the original inference.py uses # It's a tensor of shape [batch_size, num_classes], e.g., [1, 2] for [prob_real, prob_fake] probabilities_tensor = output.logits_labels.softmax(dim=1) # Probability of being FAKE is the second element (index 1) probability_fake = probabilities_tensor[0, 1].item() prediction = 1 if probability_fake >= image_input.threshold else 0 class_label = "fake" if prediction == 1 else "real" inference_time_seconds = time.time() - start_time logger.info( f"Prediction for {MODEL_NAME} completed in {inference_time_seconds:.4f}s. Prob Fake: {probability_fake:.4f}" ) # Schedule model unload if not preloaded if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: threading.Timer( MODEL_TIMEOUT + 5.0, unload_model_if_idle ).start() # Check slightly after timeout return { "model": MODEL_NAME, "probability": probability_fake, "prediction": prediction, "class": class_label, "inference_time": inference_time_seconds, } except FileNotFoundError as e: logger.error(f"Model file not found: {e}") raise HTTPException(status_code=503, detail=f"Model weights missing: {e}") except HTTPException: raise # Re-raise HTTPException directly except Exception as e: logger.exception(f"Error during prediction: {e}") # Log full traceback raise HTTPException( status_code=500, detail=f"Internal server error during prediction: {e}" ) if __name__ == "__main__": port = int(os.environ.get("MODEL_PORT", 5002)) logger.info(f"Starting {MODEL_NAME} server on port {port} with CPU: {DEVICE}") # When running app.py directly, PYTHONPATH might need to be set if model_code is not in the same dir # This is handled by Dockerfile's ENV PYTHONPATH for containerized execution. # For local dev, if model_code is sibling to this app.py's dir (e.g. in models/image/yermandy_clip_detection/) # and model_code is cloned as models/image/yermandy_clip_detection/model_code/ # local_mc_path = os.path.join(os.path.dirname(__file__), 'model_code') # if os.path.isdir(local_mc_path) and local_mc_path not in sys.path: # sys.path.insert(0, local_mc_path) uvicorn.run(app, host="0.0.0.0", port=port)