dtoro/mule-image
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: GPU acceleration support for ONNX Runtime inference

Browse Source 
Centralize execution provider selection in providers.py with
auto-detection and graceful fallback. All ONNX sessions (embedder,
detector, face processor, recognizer) now use the configured providers.

- New VISION_EXECUTION_PROVIDERS env var: "auto" for GPU auto-detect,
  or explicit "CUDAExecutionProvider,CPUExecutionProvider"
- Provider priority: CUDA > ROCm > OpenVINO > CPU (when set to "auto")
- docker-compose.yml includes commented-out NVIDIA GPU deploy section
- Supports onnxruntime-gpu as a drop-in replacement for onnxruntime

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
authentik Default Admin
2026-04-12 22:12:21 +02:00
parent 94c07b1d0d
commit bbb8e4850c
7 changed files with 138 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
backend/app/config.py
View File

@@ -66,6 +66,11 @@ class VisionSettings(BaseModel):
enabled: bool = True
backend: str = "onnx" # "onnx" | "rocm" (future)
models_dir: str = "/data/models"
# ONNX Runtime execution providers in priority order.
# Auto-detected at startup; falls back to CPU if GPU is unavailable.
# Options: "CUDAExecutionProvider", "ROCMExecutionProvider",
# "OpenVINOExecutionProvider", "CPUExecutionProvider"
execution_providers: list[str] = ["CPUExecutionProvider"]
embedder: EmbedderSettings = EmbedderSettings()
ocr: OCRSettings = OCRSettings()
detector: DetectorSettings = DetectorSettings()
@@ -183,9 +188,22 @@ class Settings(BaseSettings):
def performance(self) -> PerformanceSettings:
return self.config.performance
# ONNX Runtime execution providers, overridable via env var.
# Comma-separated: "CUDAExecutionProvider,CPUExecutionProvider"
# or "auto" for GPU auto-detection.
vision_execution_providers: str = Field(
default="CPUExecutionProvider",
env="VISION_EXECUTION_PROVIDERS",
)
@property
def vision(self) -> VisionSettings:
return self.config.vision
v = self.config.vision
# Override execution_providers from env if set.
providers = [p.strip() for p in self.vision_execution_providers.split(",") if p.strip()]
if providers:
v.execution_providers = providers
return v
class Config:
env_file = ".env"

6
backend/app/services/vision/detect.py
View File

@@ -120,12 +120,10 @@ class YOLOv8Detector(ObjectDetector):
def __init__(self, settings: VisionSettings):
model_path = Path(settings.models_dir) / "detect" / "yolov8n.onnx"
opts = ort.SessionOptions()
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
from app.services.vision.providers import create_session
logger.info("Loading YOLOv8n from %s", model_path)
self._session = ort.InferenceSession(str(model_path), opts, providers=["CPUExecutionProvider"])
self._session = create_session(str(model_path), configured_providers=settings.execution_providers)
self._conf_threshold = settings.detector.min_confidence
self._max_detections = settings.detector.max_detections

20
backend/app/services/vision/embed.py
View File

@@ -67,15 +67,15 @@ class OpenCLIPEmbedder(Embedder):
visual_path = model_dir / "visual.onnx"
textual_path = model_dir / "textual.onnx"
opts = ort.SessionOptions()
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
from app.services.vision.providers import create_session
from app.config import settings as app_settings
providers = app_settings.vision.execution_providers
logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])
self._visual = create_session(str(visual_path), configured_providers=providers)
logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])
self._textual = create_session(str(textual_path), configured_providers=providers)
def embed_image(self, image: np.ndarray) -> np.ndarray:
inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD)
@@ -110,15 +110,15 @@ class SigLIP2Embedder(Embedder):
visual_path = model_dir / "visual.onnx"
textual_path = model_dir / "textual.onnx"
opts = ort.SessionOptions()
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
from app.services.vision.providers import create_session
from app.config import settings as app_settings
providers = app_settings.vision.execution_providers
logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])
self._visual = create_session(str(visual_path), configured_providers=providers)
logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])
self._textual = create_session(str(textual_path), configured_providers=providers)
def embed_image(self, image: np.ndarray) -> np.ndarray:
inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD)

6
backend/app/services/vision/faces.py
View File

@@ -71,11 +71,9 @@ class YuNetSFaceProcessor(FaceProcessor):
logger.info("YuNet face detector loaded via OpenCV")
# SFace via ONNX Runtime
opts = ort.SessionOptions()
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
from app.services.vision.providers import create_session
ort.set_default_logger_severity(3)
self._recognizer = ort.InferenceSession(sface_path, opts, providers=["CPUExecutionProvider"])
self._recognizer = create_session(sface_path, configured_providers=settings.execution_providers)
logger.info("SFace recognizer loaded via ONNX Runtime")
self._min_face_size = settings.faces.min_face_size

5
backend/app/services/vision/insightface_processor.py
View File

@@ -23,10 +23,13 @@ class InsightFaceProcessor(FaceProcessor):
model_root = str(Path(settings.models_dir) / "face" / "insightface")
logger.info("Loading InsightFace buffalo_l from %s", model_root)
from app.services.vision.providers import get_providers
providers = get_providers(settings.execution_providers)
self._app = FaceAnalysis(
name="buffalo_l",
root=model_root,
providers=["CPUExecutionProvider"],
providers=providers,
)
self._app.prepare(ctx_id=-1, det_size=(640, 640))
self._min_det_score = settings.faces.recognition_threshold

86
backend/app/services/vision/providers.py Normal file
View File

@@ -0,0 +1,86 @@
"""
ONNX Runtime execution provider resolution with GPU auto-detection.
Resolves configured execution providers against what's actually available
in the current ONNX Runtime build. Falls back to CPU if no GPU provider
is available. Logs the selected provider so users can confirm GPU is active.
"""
import logging
import onnxruntime as ort
logger = logging.getLogger(__name__)
_resolved: list[str] | None = None
def get_providers(configured: list[str] | None = None) -> list[str]:
"""Return the best available execution providers.
1. If `configured` is provided, filter to only those that are
actually available in the current ORT build.
2. If none of the configured providers are available, fall back
to CPUExecutionProvider.
3. Auto-detect: if configured is ["auto"], probe for GPU providers.
Results are cached after first call.
"""
global _resolved
if _resolved is not None:
return _resolved
available = set(ort.get_available_providers())
logger.info("ONNX Runtime available providers: %s", sorted(available))
if configured is None or configured == ["CPUExecutionProvider"]:
_resolved = ["CPUExecutionProvider"]
return _resolved
if configured == ["auto"]:
# Auto-detect: prefer CUDA > ROCm > OpenVINO > CPU
priority = [
"CUDAExecutionProvider",
"ROCMExecutionProvider",
"OpenVINOExecutionProvider",
]
for p in priority:
if p in available:
_resolved = [p, "CPUExecutionProvider"]
logger.info("Auto-detected GPU provider: %s", p)
return _resolved
_resolved = ["CPUExecutionProvider"]
logger.info("No GPU provider detected, using CPU")
return _resolved
# Filter configured list to available providers.
resolved = [p for p in configured if p in available]
if not resolved:
logger.warning(
"None of the configured providers %s are available. "
"Falling back to CPU. Available: %s",
configured,
sorted(available),
)
resolved = ["CPUExecutionProvider"]
else:
# Always include CPU as fallback.
if "CPUExecutionProvider" not in resolved:
resolved.append("CPUExecutionProvider")
_resolved = resolved
logger.info("Using ONNX Runtime providers: %s", _resolved)
return _resolved
def create_session(
model_path: str,
opts: ort.SessionOptions | None = None,
configured_providers: list[str] | None = None,
) -> ort.InferenceSession:
"""Create an ONNX InferenceSession with the best available providers."""
providers = get_providers(configured_providers)
if opts is None:
opts = ort.SessionOptions()
opts.inter_op_num_threads = 2
opts.intra_op_num_threads = 2
return ort.InferenceSession(model_path, opts, providers=providers)