diff --git a/backend/app/services/vision/bootstrap_models.py b/backend/app/services/vision/bootstrap_models.py index 0238dbf..5ac7fc2 100644 --- 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) diff --git a/backend/app/services/vision/embed.py b/backend/app/services/vision/embed.py index 36e5a80..b5357fa 100644 --- 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) diff --git a/backend/app/services/vision/export_models.py b/backend/app/services/vision/export_models.py new file mode 100644 index 0000000..eaa6cbf --- /dev/null +++ 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() diff --git a/backend/requirements.txt b/backend/requirements.txt index d41d66f..5e59f06 100644 --- 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 diff --git a/docker-compose.yml b/docker-compose.yml index 16d5f0c..d1e1e57 100644 --- 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