diff --git a/backend/app/scripts/__init__.py b/backend/app/scripts/__init__.py new file mode 100644 index 0000000..0d3c54b --- /dev/null +++ b/backend/app/scripts/__init__.py @@ -0,0 +1 @@ +"""Operational scripts run from the container entrypoint or CLI.""" diff --git a/backend/app/scripts/download_models.py b/backend/app/scripts/download_models.py new file mode 100644 index 0000000..39131a1 --- /dev/null +++ b/backend/app/scripts/download_models.py @@ -0,0 +1,53 @@ +"""Self-heal model weights into /models. Idempotent — only missing files +are fetched. Called by the ml-worker entrypoint before Celery starts. +""" + +import os +import sys +from pathlib import Path + +MODEL_ROOT = Path(os.environ.get("ML_MODEL_DIR", "/models")) +CAMIE_REPO = os.environ.get("CAMIE_HF_REPO", "Camais03/camie-tagger-v2") +SIGLIP_REPO = os.environ.get( + "SIGLIP_HF_REPO", "google/siglip-so400m-patch14-384" +) + + +def _snapshot(repo_id: str, dest: Path, allow_patterns: list[str] | None) -> None: + from huggingface_hub import snapshot_download + + dest.mkdir(parents=True, exist_ok=True) + snapshot_download( + repo_id=repo_id, + local_dir=str(dest), + allow_patterns=allow_patterns, + ) + + +def ensure_camie() -> None: + dest = MODEL_ROOT / "camie" + if (dest / "model.onnx").is_file() and (dest / "selected_tags.csv").is_file(): + print(f"[download_models] Camie present at {dest}") + return + print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}") + _snapshot(CAMIE_REPO, dest, ["model.onnx", "selected_tags.csv", "*.json"]) + + +def ensure_siglip() -> None: + dest = MODEL_ROOT / "siglip" + if (dest / "config.json").is_file() and any(dest.glob("*.safetensors")): + print(f"[download_models] SigLIP present at {dest}") + return + print(f"[download_models] Fetching {SIGLIP_REPO} -> {dest}") + _snapshot(SIGLIP_REPO, dest, None) + + +def main() -> int: + ensure_camie() + ensure_siglip() + print("[download_models] Done.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/backend/app/tasks/import_file.py b/backend/app/tasks/import_file.py index e350909..bcc83db 100644 --- a/backend/app/tasks/import_file.py +++ b/backend/app/tasks/import_file.py @@ -95,11 +95,13 @@ def import_media_file(self, import_task_id: int) -> dict: session.add(task) session.commit() - # Enqueue the thumbnail task for newly imported images. + # Enqueue the thumbnail + ML tasks for newly imported images. if result.status == "imported" and result.image_id is not None: + from .ml import tag_and_embed from .thumbnail import generate_thumbnail generate_thumbnail.delay(result.image_id) + tag_and_embed.delay(result.image_id) # If this was the last task in the batch, mark the batch complete. remaining = session.execute( diff --git a/entrypoint.sh b/entrypoint.sh index 53a6f6a..3d9cdf5 100755 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -37,6 +37,8 @@ case "$ROLE" in ;; ml-worker) + echo "[entrypoint] Ensuring ML models present in /models..." + python -m backend.app.scripts.download_models echo "[entrypoint] Starting ML Celery worker (ml queue)" exec celery -A backend.app.celery_app:celery worker \ --loglevel=info \ diff --git a/tests/test_download_models.py b/tests/test_download_models.py new file mode 100644 index 0000000..edd13be --- /dev/null +++ b/tests/test_download_models.py @@ -0,0 +1,45 @@ +"""download_models tests. No network in CI; we test the 'already present → +skip' short-circuit by faking the expected files, and that main() wires +both ensure_* calls. +""" + +from unittest.mock import patch + +from backend.app.scripts import download_models as dm + + +def test_ensure_camie_skips_when_present(tmp_path, monkeypatch): + monkeypatch.setattr(dm, "MODEL_ROOT", tmp_path) + camie = tmp_path / "camie" + camie.mkdir(parents=True) + (camie / "model.onnx").write_bytes(b"x") + (camie / "selected_tags.csv").write_text("tag_id,name,category,count\n") + with patch.object(dm, "_snapshot") as snap: + dm.ensure_camie() + snap.assert_not_called() + + +def test_ensure_camie_downloads_when_missing(tmp_path, monkeypatch): + monkeypatch.setattr(dm, "MODEL_ROOT", tmp_path) + with patch.object(dm, "_snapshot") as snap: + dm.ensure_camie() + snap.assert_called_once() + + +def test_ensure_siglip_skips_when_present(tmp_path, monkeypatch): + monkeypatch.setattr(dm, "MODEL_ROOT", tmp_path) + sig = tmp_path / "siglip" + sig.mkdir(parents=True) + (sig / "config.json").write_text("{}") + (sig / "model.safetensors").write_bytes(b"x") + with patch.object(dm, "_snapshot") as snap: + dm.ensure_siglip() + snap.assert_not_called() + + +def test_main_calls_both(monkeypatch): + calls = [] + monkeypatch.setattr(dm, "ensure_camie", lambda: calls.append("camie")) + monkeypatch.setattr(dm, "ensure_siglip", lambda: calls.append("siglip")) + assert dm.main() == 0 + assert calls == ["camie", "siglip"]