Add ViT detector api

app/core/detector/config.py

@@ -14,7 +14,8 @@ if torch.cuda.is_available():
     print(f"GPU: {torch.cuda.get_device_name(0)}")
 
 # Model configuration
-BASE_MODEL_NAME = "shunda012/siglip-deepfake-detector"
+SIGLIP_MODEL_NAME = "shunda012/siglip-deepfake-detector"
+VIT_MODEL_NAME = "shunda012/vit-deepfake-detector"
 
 # Prediction threshold
 REAL_THRESHOLD = 0.90  # classify as real only when P(real) >= 90%

app/core/detector/model.py

@@ -5,9 +5,14 @@ Model loading for the deepfake detector.
 from dataclasses import dataclass
 from typing import Optional
 
-from transformers import AutoImageProcessor, SiglipForImageClassification
+from transformers import (
+    AutoImageProcessor,
+    SiglipForImageClassification,
+    ViTImageProcessor,
+    ViTForImageClassification,
+)
 
-from app.core.detector.config import BASE_MODEL_NAME, DEVICE
+from app.core.detector.config import SIGLIP_MODEL_NAME, VIT_MODEL_NAME, DEVICE
 
 
 @dataclass(frozen=True)
@@ -18,7 +23,16 @@ class SiglipResources:
     processor: AutoImageProcessor
 
 
+@dataclass(frozen=True)
+class ViTResources:
+    """Container for the ViT model and processor."""
+
+    model: ViTForImageClassification
+    processor: ViTImageProcessor
+
+
 _siglip_resources: Optional[SiglipResources] = None
+_vit_resources: Optional[ViTResources] = None
 
 
 def get_siglip_model() -> SiglipResources:
@@ -34,8 +48,8 @@ def get_siglip_model() -> SiglipResources:
     if _siglip_resources is None:
         print("Loading SigLIP Model...")
 
-        siglip_processor = AutoImageProcessor.from_pretrained(BASE_MODEL_NAME)
-        siglip_model = SiglipForImageClassification.from_pretrained(BASE_MODEL_NAME)
+        siglip_processor = AutoImageProcessor.from_pretrained(SIGLIP_MODEL_NAME)
+        siglip_model = SiglipForImageClassification.from_pretrained(SIGLIP_MODEL_NAME)
         siglip_model = siglip_model.to(DEVICE)
         siglip_model.eval()
 
@@ -45,3 +59,29 @@ def get_siglip_model() -> SiglipResources:
         )
 
     return _siglip_resources
+
+
+def get_vit_model() -> ViTResources:
+    """
+    Get or load the merged ViT detector model.
+
+    Returns:
+        ViTResources: Loaded model and processor (cached singleton).
+    """
+
+    global _vit_resources
+
+    if _vit_resources is None:
+        print("Loading ViT Model...")
+
+        vit_processor = ViTImageProcessor.from_pretrained(VIT_MODEL_NAME)
+        vit_model = ViTForImageClassification.from_pretrained(VIT_MODEL_NAME)
+        vit_model = vit_model.to(DEVICE)
+        vit_model.eval()
+
+        _vit_resources = ViTResources(
+            model=vit_model,
+            processor=vit_processor,
+        )
+
+    return _vit_resources

app/detector.py

@@ -11,7 +11,12 @@ from typing import Callable
 from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
 from PIL import Image, UnidentifiedImageError
 
-from app.core.detector.model import SiglipResources, get_siglip_model
+from app.core.detector.model import (
+    SiglipResources,
+    ViTResources,
+    get_siglip_model,
+    get_vit_model,
+)
 from app.services.detector.prediction import predict_single_image
 from app.services.detector.transforms import get_eval_transforms
 
@@ -27,8 +32,16 @@ def get_siglip_transforms():
     return get_eval_transforms(resources.processor, "siglip")
 
 
-@detector.post("/detect")
-async def detect_deepfake(
+@lru_cache(maxsize=1)
+def get_vit_transforms():
+    """Build and cache ViT evaluation transforms once per process."""
+
+    resources = get_vit_model()
+    return get_eval_transforms(resources.processor, "vit")
+
+
+@detector.post("/siglip-detect")
+async def siglip_detect_deepfake(
     file: UploadFile = File(...),
     resources: SiglipResources = Depends(get_siglip_model),
     siglip_transforms: Callable = Depends(get_siglip_transforms),
@@ -59,3 +72,37 @@ async def detect_deepfake(
     except Exception as exc:  # pragma: no cover - defensive server guard
         logger.exception("Unhandled error during deepfake detection")
         raise HTTPException(status_code=500, detail="Error processing image") from exc
+
+
+@detector.post("/vit-detect")
+async def vit_detect_deepfake(
+    file: UploadFile = File(...),
+    resources: ViTResources = Depends(get_vit_model),
+    vit_transforms: Callable = Depends(get_vit_transforms),
+):
+    """
+    Detect if an image is a deepfake or real using ViT + LoRA model.
+
+    Args:
+        file: Uploaded image file
+
+    Returns:
+        JSON response with prediction results
+    """
+
+    try:
+        image_bytes = await file.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+
+        result = predict_single_image(
+            image, resources.model, vit_transforms, "ViT + LoRA"
+        )
+
+        return result
+    except UnidentifiedImageError:
+        raise HTTPException(status_code=422, detail="Invalid or unsupported image file")
+    except HTTPException:
+        raise
+    except Exception as exc:  # pragma: no cover - defensive server guard
+        logger.exception("Unhandled error during deepfake detection")
+        raise HTTPException(status_code=500, detail="Error processing image") from exc

app/services/detector/transforms.py

@@ -8,11 +8,11 @@ from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normal
 def get_eval_transforms(processor, model_type="vit"):
     """
     Create evaluation transforms based on processor settings.
-    
+
     Args:
         processor: The image processor from the model
         model_type: Type of model ("vit" or "siglip")
-        
+
     Returns:
         Composed transforms for image preprocessing
     """
@@ -20,7 +20,7 @@ def get_eval_transforms(processor, model_type="vit"):
     image_mean = processor.image_mean
     image_std = processor.image_std
     normalize = Normalize(mean=image_mean, std=image_std)
-    
+
     return Compose(
         [
             Resize(size if model_type == "siglip" else 256),