fix(ai): support both old and new mediapipe APIs for airgapped Docker (#69)

MediaPipe >= 0.10.30 removed the mp.solutions namespace. This broke
face blur, face enhance, red-eye removal, and photo restoration for
users running newer mediapipe versions (closes #43).

All 5 Python scripts that use mediapipe now try the legacy mp.solutions
API first and fall back to the new mp.tasks API on AttributeError.
Model files (blaze_face_short_range.task, face_landmarker.task) are
pre-downloaded during Docker build into /opt/models/mediapipe/ so the
image works fully airgapped. Local dev auto-downloads to .models/.

Co-authored-by: stirling-image <stirling-image@users.noreply.github.com>

docker/download_models.py

@@ -65,6 +65,14 @@ NAFNET_MODEL_URL = (
 NAFNET_MODEL_PATH = os.path.join(NAFNET_MODEL_DIR, "NAFNet-SIDD-width64.pth")
 NAFNET_MIN_SIZE = 60_000_000  # ~67 MB
 
+MEDIAPIPE_MODEL_DIR = "/opt/models/mediapipe"
+FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task"
+FACE_DETECT_MODEL_PATH = os.path.join(MEDIAPIPE_MODEL_DIR, "blaze_face_short_range.task")
+FACE_DETECT_MIN_SIZE = 100_000  # ~200 KB
+FACE_LANDMARKER_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
+FACE_LANDMARKER_MODEL_PATH = os.path.join(MEDIAPIPE_MODEL_DIR, "face_landmarker.task")
+FACE_LANDMARKER_MIN_SIZE = 1_000_000  # ~7 MB
+
 REMBG_MODELS = [
     "u2net",
     "isnet-general-use",
@@ -310,19 +318,60 @@ def download_nafnet_model():
     print(f"  NAFNet model downloaded: {size:,} bytes")
 
 
+def download_mediapipe_task_models():
+    """Download MediaPipe tasks API model files for face detection and landmarks.
+
+    These models are used by the mp.tasks fallback when mp.solutions is
+    unavailable (mediapipe >= 0.10.30). Pre-downloading ensures the Docker
+    image works fully airgapped.
+    """
+    print("=== Downloading MediaPipe task models ===")
+    os.makedirs(MEDIAPIPE_MODEL_DIR, exist_ok=True)
+
+    for url, path, name, min_size in [
+        (FACE_DETECT_MODEL_URL, FACE_DETECT_MODEL_PATH,
+         "blaze_face_short_range", FACE_DETECT_MIN_SIZE),
+        (FACE_LANDMARKER_MODEL_URL, FACE_LANDMARKER_MODEL_PATH,
+         "face_landmarker", FACE_LANDMARKER_MIN_SIZE),
+    ]:
+        print(f"  Downloading {name}...")
+        urllib.request.urlretrieve(url, path)
+        size = os.path.getsize(path)
+        assert size > min_size, (
+            f"{name} model too small: {size} bytes (expected > {min_size})"
+        )
+        print(f"  {name} downloaded ({size / 1_000_000:.1f} MB)")
+    print("MediaPipe task models downloaded.\n")
+
+
 def verify_mediapipe():
     """Verify MediaPipe face detection models are bundled in the wheel."""
     print("=== Verifying MediaPipe models ===")
     import mediapipe as mp
 
-    for selection in [0, 1]:
+    try:
-        label = "short-range" if selection == 0 else "full-range"
+        for selection in [0, 1]:
-        print(f"  Verifying {label} model (selection={selection})...")
+            label = "short-range" if selection == 0 else "full-range"
-        detector = mp.solutions.face_detection.FaceDetection(
+            print(f"  Verifying {label} model (selection={selection})...")
-            model_selection=selection, min_detection_confidence=0.5
+            detector = mp.solutions.face_detection.FaceDetection(
+                model_selection=selection, min_detection_confidence=0.5
+            )
+            detector.close()
+            print(f"  {label} model OK")
+    except AttributeError:
+        # mediapipe >= 0.10.30 removed mp.solutions; verify tasks API instead
+        print("  mp.solutions unavailable, verifying mp.tasks API...")
+        options = mp.tasks.vision.FaceDetectorOptions(
+            base_options=mp.tasks.BaseOptions(
+                model_asset_path=FACE_DETECT_MODEL_PATH
+            ),
+            running_mode=mp.tasks.vision.RunningMode.IMAGE,
+            min_detection_confidence=0.5,
         )
+        detector = mp.tasks.vision.FaceDetector.create_from_options(options)
         detector.close()
-        print(f"  {label} model OK")
+        print("  mp.tasks FaceDetector OK")
+
     print("MediaPipe models verified.\n")
 
 
@@ -423,6 +472,19 @@ def smoke_test():
     assert os.path.isdir(vl_dir), f"PaddleOCR-VL model missing: {vl_dir}"
     print("  PaddleOCR-VL model verified")
 
+    # MediaPipe task models must exist (for mp.tasks fallback)
+    assert os.path.exists(FACE_DETECT_MODEL_PATH), (
+        f"MediaPipe face detector model missing: {FACE_DETECT_MODEL_PATH}"
+    )
+    assert os.path.getsize(FACE_DETECT_MODEL_PATH) > FACE_DETECT_MIN_SIZE
+    print("  MediaPipe face detector model verified")
+
+    assert os.path.exists(FACE_LANDMARKER_MODEL_PATH), (
+        f"MediaPipe face landmarker model missing: {FACE_LANDMARKER_MODEL_PATH}"
+    )
+    assert os.path.getsize(FACE_LANDMARKER_MODEL_PATH) > FACE_LANDMARKER_MIN_SIZE
+    print("  MediaPipe face landmarker model verified")
+
     print("Smoke test passed.\n")
 
 
@@ -439,6 +501,7 @@ def main():
     download_paddleocr_vl_model()
     download_scunet_model()
     download_nafnet_model()
+    download_mediapipe_task_models()
     verify_mediapipe()
     smoke_test()
     print("All models downloaded and verified.")

packages/ai/python/detect_faces.py

@@ -12,19 +12,22 @@ def emit_progress(percent, stage):
 # ── Model path for new mp.tasks API ─────────────────────────────────
 
 _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task"
-_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.task"
-_FACE_DETECT_MODEL_PATH = os.path.join(_MODEL_DIR, "blaze_face_short_range.task")
+_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.task")
 
 
 def _ensure_face_detect_model():
-    """Download the face detector model if not present."""
+    """Resolve face detector model. Docker path first, then local dev."""
-    if os.path.exists(_FACE_DETECT_MODEL_PATH):
+    if os.path.exists(_DOCKER_MODEL_PATH):
-        return _FACE_DETECT_MODEL_PATH
+        return _DOCKER_MODEL_PATH
-    os.makedirs(_MODEL_DIR, exist_ok=True)
+    if os.path.exists(_LOCAL_MODEL_PATH):
+        return _LOCAL_MODEL_PATH
+    os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True)
     import urllib.request
     emit_progress(15, "Downloading face detection model")
-    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_MODEL_PATH)
+    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _LOCAL_MODEL_PATH)
-    return _FACE_DETECT_MODEL_PATH
+    return _LOCAL_MODEL_PATH
 
 
 def _detect_with_solutions(img_array, min_confidence):

packages/ai/python/enhance_faces.py

@@ -40,19 +40,22 @@ CODEFORMER_MODEL_PATH = os.environ.get(
 # ── Model path for new mp.tasks API ─────────────────────────────────
 
 _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task"
-_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.task"
-_FACE_DETECT_MODEL_PATH = os.path.join(_MODEL_DIR, "blaze_face_short_range.task")
+_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.task")
 
 
 def _ensure_face_detect_model():
-    """Download the face detector model if not present."""
+    """Resolve face detector model. Docker path first, then local dev."""
-    if os.path.exists(_FACE_DETECT_MODEL_PATH):
+    if os.path.exists(_DOCKER_MODEL_PATH):
-        return _FACE_DETECT_MODEL_PATH
+        return _DOCKER_MODEL_PATH
-    os.makedirs(_MODEL_DIR, exist_ok=True)
+    if os.path.exists(_LOCAL_MODEL_PATH):
+        return _LOCAL_MODEL_PATH
+    os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True)
     import urllib.request
     emit_progress(15, "Downloading face detection model")
-    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_MODEL_PATH)
+    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _LOCAL_MODEL_PATH)
-    return _FACE_DETECT_MODEL_PATH
+    return _LOCAL_MODEL_PATH
 
 
 def detect_faces_mediapipe(img_array, sensitivity):

packages/ai/python/face_landmarks.py

@@ -79,12 +79,15 @@ def detect_with_solutions(img_array):
 # ── New API: mp.tasks (mediapipe >= 0.10.30) ───────────────────────
 
 MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
+_DOCKER_MODEL_PATH = "/opt/models/mediapipe/face_landmarker.task"
 MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
 MODEL_PATH = os.path.join(MODEL_DIR, "face_landmarker.task")
 
 
 def ensure_model():
-    """Download the face landmarker model if not present."""
+    """Resolve face landmarker model. Docker path first, then local dev."""
+    if os.path.exists(_DOCKER_MODEL_PATH):
+        return _DOCKER_MODEL_PATH
     if os.path.exists(MODEL_PATH):
         return MODEL_PATH
     os.makedirs(MODEL_DIR, exist_ok=True)

packages/ai/python/red_eye_removal.py

@@ -12,19 +12,22 @@ def emit_progress(percent, stage):
 # ── Model path for new mp.tasks API ─────────────────────────────────
 
 _FACE_MESH_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
-_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_DOCKER_MODEL_PATH = "/opt/models/mediapipe/face_landmarker.task"
-_FACE_MESH_MODEL_PATH = os.path.join(_MODEL_DIR, "face_landmarker.task")
+_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "face_landmarker.task")
 
 
 def _ensure_face_mesh_model():
-    """Download the face landmarker model if not present."""
+    """Resolve face landmarker model. Docker path first, then local dev."""
-    if os.path.exists(_FACE_MESH_MODEL_PATH):
+    if os.path.exists(_DOCKER_MODEL_PATH):
-        return _FACE_MESH_MODEL_PATH
+        return _DOCKER_MODEL_PATH
-    os.makedirs(_MODEL_DIR, exist_ok=True)
+    if os.path.exists(_LOCAL_MODEL_PATH):
+        return _LOCAL_MODEL_PATH
+    os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True)
     import urllib.request
     emit_progress(15, "Downloading face mesh model")
-    urllib.request.urlretrieve(_FACE_MESH_MODEL_URL, _FACE_MESH_MODEL_PATH)
+    urllib.request.urlretrieve(_FACE_MESH_MODEL_URL, _LOCAL_MODEL_PATH)
-    return _FACE_MESH_MODEL_PATH
+    return _LOCAL_MODEL_PATH
 
 
 def _mesh_with_solutions(img_array, max_faces=10, min_confidence=0.5):

packages/ai/python/restore.py

@@ -220,19 +220,22 @@ def _get_codeformer_path():
 # ── Model path for new mp.tasks API ─────────────────────────────────
 
 _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task"
-_FACE_DETECT_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_FACE_DETECT_DOCKER_PATH = "/opt/models/mediapipe/blaze_face_short_range.task"
-_FACE_DETECT_MODEL_PATH = os.path.join(_FACE_DETECT_MODEL_DIR, "blaze_face_short_range.task")
+_FACE_DETECT_LOCAL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
+_FACE_DETECT_LOCAL_PATH = os.path.join(_FACE_DETECT_LOCAL_DIR, "blaze_face_short_range.task")
 
 
 def _ensure_face_detect_model():
-    """Download the face detector model if not present."""
+    """Resolve face detector model. Docker path first, then local dev."""
-    if os.path.exists(_FACE_DETECT_MODEL_PATH):
+    if os.path.exists(_FACE_DETECT_DOCKER_PATH):
-        return _FACE_DETECT_MODEL_PATH
+        return _FACE_DETECT_DOCKER_PATH
-    os.makedirs(_FACE_DETECT_MODEL_DIR, exist_ok=True)
+    if os.path.exists(_FACE_DETECT_LOCAL_PATH):
+        return _FACE_DETECT_LOCAL_PATH
+    os.makedirs(_FACE_DETECT_LOCAL_DIR, exist_ok=True)
     import urllib.request
     emit_progress(15, "Downloading face detection model")
-    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_MODEL_PATH)
+    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_LOCAL_PATH)
-    return _FACE_DETECT_MODEL_PATH
+    return _FACE_DETECT_LOCAL_PATH
 
 
 def enhance_faces(img_bgr, fidelity=0.7):