fix: model weights setup — export scripts, ORT compat, bootstrap

- Add export_models.py for OpenCLIP ViT-B/32 and YOLOv8n ONNX export - Fix ArgMax(13) ORT ARM64 incompatibility by passing eot_indices as a separate ONNX input (computed outside the graph in embed.py) - Use legacy TorchScript exporter (dynamo=False) for IR version 9 compat - Upgrade onnxruntime to 1.18.1 - Rewrite bootstrap_models.py with clear separation of auto-downloadable models (YuNet, SFace) vs manually-exported ones (OpenCLIP, YOLOv8n) - Wire bootstrap into worker CMD (runs before Celery) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-10 09:41:25 +02:00
parent 29177f0c1a
commit 2a6661f779
5 changed files with 233 additions and 38 deletions
--- a/backend/app/services/vision/bootstrap_models.py
+++ b/backend/app/services/vision/bootstrap_models.py
@@ -4,7 +4,11 @@ Download vision model weights on first worker boot.
 Run as:  python -m app.services.vision.bootstrap_models

 Or called from the vision worker entrypoint before Celery starts.
-Downloads are idempotent — existing files with matching sizes are skipped.
+Downloads are idempotent — existing files are skipped.
+
+For models that require export (OpenCLIP, YOLOv8n), see export_models.py.
+Those must be exported once on any machine with pip, then placed in
+the models volume before the worker starts.
 """
 import logging
 import os
@@ -15,35 +19,30 @@ from app.config import settings

 logger = logging.getLogger(__name__)

-# (relative_path, url, expected_size_bytes_approx)
-# Sizes are approximate — used only for skip-if-exists checks, not integrity.
-MODELS = [
-    # OpenCLIP ViT-B/32 — visual and textual encoders (ONNX)
-    # These must be exported manually via export_openclip.py (see below).
-    # Placeholder entries — bootstrap will warn if missing.
-    ("embed/visual.onnx", None, None),
-    ("embed/textual.onnx", None, None),
-
-    # YOLOv8n — object detection
-    # Export: `yolo export model=yolov8n.pt format=onnx imgsz=640`
-    # Placeholder — must be exported from ultralytics offline.
-    ("detect/yolov8n.onnx", None, None),
-
-    # YuNet — face detection (Apache 2.0, opencv_zoo)
+# (relative_path, url, description)
+# Models with url=None must be pre-exported via export_models.py.
+DOWNLOADS = [
+    # YuNet — face detection (Apache 2.0, opencv_zoo, ~233 KB)
    (
        "face/yunet.onnx",
        "https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
-        233_000,
+        "YuNet face detector",
    ),
-
-    # SFace — face recognition (Apache 2.0, opencv_zoo)
+    # SFace — face recognition (Apache 2.0, opencv_zoo, ~37 MB)
    (
        "face/sface.onnx",
        "https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx",
-        37_000_000,
+        "SFace face recognizer",
    ),
 ]

+# Models that need manual export via export_models.py
+EXPORTS = [
+    ("embed/visual.onnx", "OpenCLIP ViT-B/32 visual encoder"),
+    ("embed/textual.onnx", "OpenCLIP ViT-B/32 textual encoder"),
+    ("detect/yolov8n.onnx", "YOLOv8n object detector"),
+]
+

 def bootstrap(models_dir: str | None = None):
    """Ensure all model files are present. Download what we can, warn about
@@ -51,32 +50,46 @@ def bootstrap(models_dir: str | None = None):
    base = Path(models_dir or settings.vision.models_dir)
    base.mkdir(parents=True, exist_ok=True)

-    for rel_path, url, expected_size in MODELS:
+    # Download auto-downloadable models
+    for rel_path, url, desc in DOWNLOADS:
        dest = base / rel_path
        dest.parent.mkdir(parents=True, exist_ok=True)

        if dest.exists():
-            logger.debug("Model already exists: %s", dest)
+            logger.debug("Already exists: %s (%s)", dest, desc)
            continue

-        if url is None:
-            logger.warning(
-                "Model file %s not found and has no auto-download URL. "
-                "See bootstrap_models.py for export instructions.",
-                dest,
-            )
-            continue
-
-        logger.info("Downloading %s → %s", url, dest)
+        logger.info("Downloading %s → %s", desc, dest)
        try:
            urlretrieve(url, str(dest))
-            actual = dest.stat().st_size
-            logger.info("Downloaded %s (%d bytes)", rel_path, actual)
+            size_kb = dest.stat().st_size / 1024
+            logger.info("Downloaded %s (%.0f KB)", desc, size_kb)
        except Exception as e:
-            logger.error("Failed to download %s: %s", rel_path, e)
+            logger.error("Failed to download %s: %s", desc, e)
            if dest.exists():
                dest.unlink()

+    # Check for manually-exported models
+    missing = []
+    for rel_path, desc in EXPORTS:
+        dest = base / rel_path
+        if not dest.exists():
+            missing.append((rel_path, desc))
+
+    if missing:
+        logger.warning(
+            "Missing %d model(s) that require manual export via export_models.py:",
+            len(missing),
+        )
+        for rel_path, desc in missing:
+            logger.warning("  %s — %s", base / rel_path, desc)
+        logger.warning(
+            "Run: python -m app.services.vision.export_models --models-dir %s",
+            base,
+        )
+    else:
+        logger.info("All model files present in %s", base)
+

 if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
--- a/backend/app/services/vision/embed.py
+++ b/backend/app/services/vision/embed.py
@@ -71,8 +71,13 @@ class OpenCLIPEmbedder(Embedder):
        import open_clip
        tokenizer = open_clip.get_tokenizer("ViT-B-32")
        tokens = tokenizer([text]).numpy().astype(np.int64)
-        input_name = self._textual.get_inputs()[0].name
-        out = self._textual.run(None, {input_name: tokens})[0][0]
+        # Compute EOT indices outside ONNX (avoids ArgMax(13) op)
+        eot_indices = tokens.argmax(axis=-1).astype(np.int64)
+        inputs = self._textual.get_inputs()
+        out = self._textual.run(None, {
+            inputs[0].name: tokens,
+            inputs[1].name: eot_indices,
+        })[0][0]
        out = out / np.linalg.norm(out)
        return out.astype(np.float32)

--- a/backend/app/services/vision/export_models.py
+++ b/backend/app/services/vision/export_models.py
@@ -0,0 +1,177 @@
+"""
+Export / download all vision model weights to ONNX format.
+
+Run ONCE on any machine with Python + pip (doesn't need GPU):
+
+    pip install open-clip-torch ultralytics onnx
+    python -m app.services.vision.export_models [--models-dir /data/models]
+
+This produces:
+    embed/visual.onnx        (~350 MB)
+    embed/textual.onnx       (~250 MB)
+    detect/yolov8n.onnx      (~12 MB)
+
+YuNet and SFace are downloaded by bootstrap_models.py at worker boot
+(Apache 2.0, lightweight, no export step needed).
+
+After export, copy the /data/models directory into your Docker volume:
+    docker cp /data/models mulita-worker:/data/models
+Or mount a host path in docker-compose.yml.
+"""
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def export_openclip(models_dir: Path):
+    """Export OpenCLIP ViT-B/32 to two ONNX files (visual + textual)."""
+    import torch
+    import open_clip
+
+    out_dir = models_dir / "embed"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    visual_path = out_dir / "visual.onnx"
+    textual_path = out_dir / "textual.onnx"
+
+    if visual_path.exists() and textual_path.exists():
+        logger.info("OpenCLIP ONNX files already exist, skipping export")
+        return
+
+    logger.info("Loading OpenCLIP ViT-B-32 laion2b_s34b_b79k...")
+    model, _, preprocess = open_clip.create_model_and_transforms(
+        "ViT-B-32", pretrained="laion2b_s34b_b79k"
+    )
+    model.eval()
+
+    # Use dynamo=False to get the legacy TorchScript exporter which
+    # produces IR version 9 (compatible with onnxruntime 1.17.x).
+    # The new torch.onnx.export default (dynamo=True) emits IR 10.
+    export_kwargs = dict(opset_version=14, dynamo=False)
+
+    # ── Visual encoder ────────────────────────────────────────────────
+    if not visual_path.exists():
+        logger.info("Exporting visual encoder → %s", visual_path)
+        dummy_image = torch.randn(1, 3, 224, 224)
+        torch.onnx.export(
+            model.visual,
+            dummy_image,
+            str(visual_path),
+            input_names=["image"],
+            output_names=["embedding"],
+            dynamic_axes={"image": {0: "batch"}},
+            **export_kwargs,
+        )
+        size_mb = visual_path.stat().st_size / 1e6
+        logger.info("Visual encoder exported (%.1f MB)", size_mb)
+
+    # ── Textual encoder ───────────────────────────────────────────────
+    if not textual_path.exists():
+        logger.info("Exporting textual encoder → %s", textual_path)
+        tokenizer = open_clip.get_tokenizer("ViT-B-32")
+        dummy_text = tokenizer(["a photo"]).to(torch.int64)
+
+        class TextEncoder(torch.nn.Module):
+            """Wrap the CLIP text encoder to avoid argmax in the ONNX graph.
+            OpenCLIP uses argmax to find the EOT token position, but ORT
+            ARM64 doesn't support ArgMax(13). We pre-compute the EOT index
+            from the token sequence and pass it directly."""
+            def __init__(self, clip_model):
+                super().__init__()
+                self.transformer = clip_model.transformer
+                self.token_embedding = clip_model.token_embedding
+                self.positional_embedding = clip_model.positional_embedding
+                self.ln_final = clip_model.ln_final
+                self.text_projection = clip_model.text_projection
+
+            def forward(self, text, eot_indices):
+                x = self.token_embedding(text)
+                x = x + self.positional_embedding
+                x = x.permute(1, 0, 2)  # NLD -> LND
+                x = self.transformer(x)
+                x = x.permute(1, 0, 2)  # LND -> NLD
+                x = self.ln_final(x)
+                # Take the feature at the EOT token. The EOT index is
+                # passed in as a separate input (computed outside ONNX)
+                # to avoid ArgMax(13) which ORT ARM64 doesn't support.
+                x = x[torch.arange(x.shape[0]), eot_indices]
+                x = x @ self.text_projection
+                return x
+
+        text_enc = TextEncoder(model)
+        text_enc.eval()
+
+        # Compute EOT indices from dummy tokens (argmax of token ids)
+        dummy_eot = dummy_text.argmax(dim=-1)
+
+        torch.onnx.export(
+            text_enc,
+            (dummy_text, dummy_eot),
+            str(textual_path),
+            input_names=["text", "eot_indices"],
+            output_names=["embedding"],
+            dynamic_axes={"text": {0: "batch"}, "eot_indices": {0: "batch"}},
+            **export_kwargs,
+        )
+        size_mb = textual_path.stat().st_size / 1e6
+        logger.info("Textual encoder exported (%.1f MB)", size_mb)
+
+
+def export_yolov8n(models_dir: Path):
+    """Export YOLOv8n to ONNX."""
+    out_dir = models_dir / "detect"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    onnx_path = out_dir / "yolov8n.onnx"
+
+    if onnx_path.exists():
+        logger.info("YOLOv8n ONNX already exists, skipping export")
+        return
+
+    logger.info("Exporting YOLOv8n → %s", onnx_path)
+
+    from ultralytics import YOLO
+
+    model = YOLO("yolov8n.pt")
+    model.export(format="onnx", imgsz=640, simplify=True)
+
+    # ultralytics exports to cwd as yolov8n.onnx — move to target
+    exported = Path("yolov8n.onnx")
+    if exported.exists():
+        exported.rename(onnx_path)
+
+    size_mb = onnx_path.stat().st_size / 1e6
+    logger.info("YOLOv8n exported (%.1f MB)", size_mb)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export vision model weights to ONNX")
+    parser.add_argument(
+        "--models-dir",
+        type=Path,
+        default=Path("/data/models"),
+        help="Directory to write model files (default: /data/models)",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s %(message)s",
+    )
+
+    models_dir = args.models_dir
+    models_dir.mkdir(parents=True, exist_ok=True)
+
+    logger.info("Exporting models to %s", models_dir)
+
+    export_openclip(models_dir)
+    export_yolov8n(models_dir)
+
+    logger.info("Done. Run bootstrap_models.py next to download YuNet + SFace.")
+
+
+if __name__ == "__main__":
+    main()
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -35,7 +35,7 @@ pyexiftool==0.5.6
 watchfiles==0.21.0

 # Vision pipeline (ONNX Runtime CPU inference)
-onnxruntime==1.17.1
+onnxruntime==1.18.1
 open-clip-torch==2.24.0     # tokenizer + export helper; inference via ONNX
 rapidocr-onnxruntime==1.3.22
 scikit-learn==1.4.0          # DBSCAN for face clustering
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -62,7 +62,7 @@ services:
      context: ./backend
      dockerfile: Dockerfile
    container_name: mulita-worker
-    command: celery -A app.tasks.celery worker --loglevel=${LOG_LEVEL:-info} --concurrency=${CELERYD_CONCURRENCY:-4}
+    command: sh -c "python -m app.services.vision.bootstrap_models && celery -A app.tasks.celery worker --loglevel=${LOG_LEVEL:-info} --concurrency=${CELERYD_CONCURRENCY:-4}"
    volumes:
      - ./mulita.yml:/app/config/mulita.yml:ro
      - ${PHOTO_DIRS:-./photos}:/photos:rw