diff --git a/.env.example b/.env.example index ea55c5cc..4137682c 100644 --- a/.env.example +++ b/.env.example @@ -81,6 +81,16 @@ ML_MODEL_IDLE_TTL_SECONDS=300 # Increase this if you have plenty of VRAM/RAM (e.g., 5-8 for full functionality). # Decrease this if you experience OOM (Out of Memory) errors. ML_MAX_LOADED_MODELS=5 +# Directory where downloaded model packs (caption/OCR/object/embedding/face +# models) are cached for reuse. Only relevant for installer/desktop mode — +# Docker/full/mock dev modes are unaffected. +# Leave empty to use a platform-appropriate default app-data/cache directory. +MODEL_CACHE_DIR= +# When true, model loaders must use only what is already verified in +# MODEL_CACHE_DIR and must not attempt any network download. Intended for +# offline desktop use after first-run setup is complete. Defaults to false +# so existing Docker/full/mock dev workflows keep working unchanged. +ML_OFFLINE_ONLY=false # RQ worker class. SimpleWorker keeps ML models warm in the worker process so # ML_MODEL_IDLE_TTL_SECONDS can control when they are unloaded. RQ_WORKER_CLASS=rq.worker.worker_classes.SimpleWorker diff --git a/backend/src/find_api/core/config.py b/backend/src/find_api/core/config.py index 2301799d..a46d32f8 100644 --- a/backend/src/find_api/core/config.py +++ b/backend/src/find_api/core/config.py @@ -47,6 +47,8 @@ class Settings(BaseSettings): REMOTE_ML_FEATURES: str = "embed,caption,detect,ocr,cluster" ML_MODEL_IDLE_TTL_SECONDS: int = 300 ML_MAX_LOADED_MODELS: int = 5 + MODEL_CACHE_DIR: str = "" + ML_OFFLINE_ONLY: bool = False CLIP_MODEL: str = "ViT-B-16-SigLIP" CLIP_PRETRAINED: str = "webli" BLIP_MODEL: str = "microsoft/Florence-2-base" diff --git a/backend/src/find_api/core/model_pack.py b/backend/src/find_api/core/model_pack.py new file mode 100644 index 00000000..4a5ec1f7 --- /dev/null +++ b/backend/src/find_api/core/model_pack.py @@ -0,0 +1,127 @@ +""" +Thin interfaces for versioned model pack metadata and caching. + +This module defines the CONTRACT for installer-mode model download/cache +behavior. It intentionally does not implement real downloading yet — see +docs/plans/partial/model-cache-design.md for the full design. Existing ML +loaders in find_api/ml/ are NOT modified by this module. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Protocol + + +class PackCategory(str, Enum): + """Which ML capability a pack provides.""" + + CAPTION = "caption" + OCR = "ocr" + OBJECTS = "objects" + EMBEDDINGS = "embeddings" + FACES = "faces" + + +@dataclass(frozen=True) +class ModelPack: + """Versioned metadata for a single downloadable model pack. + + Maps onto the existing loader/config surface: + - EMBEDDINGS -> settings.CLIP_MODEL / settings.CLIP_PRETRAINED (clip_embedder.py) + - CAPTION -> settings.BLIP_MODEL (captioner.py) + - OBJECTS -> settings.YOLO_MODEL (object_detector.py) + - OCR -> PaddleOCR "en" pipeline (ocr.py) + - FACES -> InsightFace "antelopev2" (face_detector.py) + """ + + pack_id: str + category: PackCategory + version: str + source_url: str + license: str + size_bytes: int + checksum_sha256: str + compatible_app_versions: str # e.g. ">=1.0.0,<2.0.0" + config_key: str # matches the loader's config_key format, e.g. "model=..." + description: str = "" + + +class PackStatus(str, Enum): + NOT_INSTALLED = "not_installed" + DOWNLOADING = "downloading" + VERIFYING = "verifying" + INSTALLED = "installed" + CORRUPTED = "corrupted" + FAILED = "failed" + + +@dataclass +class PackProgress: + """Progress snapshot for a single pack download.""" + + pack_id: str + status: PackStatus + bytes_downloaded: int = 0 + bytes_total: int = 0 + error: str | None = None + + +class PackCache(Protocol): + """Contract for a model-pack cache implementation. + + This is intentionally a thin interface: no real download/verify logic + ships in this PR. A concrete implementation (e.g. FilesystemPackCache) + will be added in a follow-up once this contract is agreed. + """ + + def is_installed(self, pack: ModelPack) -> bool: + """Return True if the pack is present and passes checksum verification.""" + ... + + def status(self, pack: ModelPack) -> PackProgress: + """Return current status/progress for a pack.""" + ... + + def install( + self, + pack: ModelPack, + on_progress: Callable[[PackProgress], None] | None = None, + ) -> None: + """Download, verify, and atomically install a pack. Must support + resume/retry and must never leave a partially-installed pack marked + as installed.""" + ... + + def verify(self, pack: ModelPack) -> bool: + """Re-check an already-installed pack's checksum (corruption recovery).""" + ... + + def remove(self, pack: ModelPack) -> None: + """Delete a cached pack from disk.""" + ... + + +class NotImplementedPackCache: + """Placeholder PackCache used until a real implementation lands. + + Every method raises NotImplementedError on purpose — this class exists + only so other code can type-check against PackCache today without a + working cache backend yet. + """ + + def is_installed(self, pack: ModelPack) -> bool: + raise NotImplementedError("PackCache implementation not yet added") + + def status(self, pack: ModelPack) -> PackProgress: + raise NotImplementedError("PackCache implementation not yet added") + + def install(self, pack: ModelPack, on_progress=None) -> None: + raise NotImplementedError("PackCache implementation not yet added") + + def verify(self, pack: ModelPack) -> bool: + raise NotImplementedError("PackCache implementation not yet added") + + def remove(self, pack: ModelPack) -> None: + raise NotImplementedError("PackCache implementation not yet added") diff --git a/backend/tests/test_model_pack.py b/backend/tests/test_model_pack.py new file mode 100644 index 00000000..2acbe85b --- /dev/null +++ b/backend/tests/test_model_pack.py @@ -0,0 +1,69 @@ +"""Tests for the thin ModelPack/PackCache contract (no real downloads).""" + +import pytest + +from find_api.core.model_pack import ( + ModelPack, + PackCategory, + PackProgress, + PackStatus, + NotImplementedPackCache, +) + + +def _sample_pack(**overrides) -> ModelPack: + defaults = dict( + pack_id="siglip-vit-b-16", + category=PackCategory.EMBEDDINGS, + version="1.0.0", + source_url="https://example.com/siglip-vit-b-16.tar", + license="Apache-2.0", + size_bytes=350_000_000, + checksum_sha256="a" * 64, + compatible_app_versions=">=1.0.0,<2.0.0", + config_key="model=ViT-B-16-SigLIP|pretrained=webli", + ) + defaults.update(overrides) + return ModelPack(**defaults) + + +def test_model_pack_holds_expected_fields(): + pack = _sample_pack() + assert pack.category == PackCategory.EMBEDDINGS + assert pack.checksum_sha256 == "a" * 64 + assert pack.size_bytes > 0 + + +def test_model_pack_is_immutable(): + pack = _sample_pack() + with pytest.raises(AttributeError): + pack.version = "2.0.0" # frozen dataclass must reject mutation + + +def test_pack_progress_defaults(): + progress = PackProgress(pack_id="siglip-vit-b-16", status=PackStatus.NOT_INSTALLED) + assert progress.bytes_downloaded == 0 + assert progress.bytes_total == 0 + assert progress.error is None + + +def test_not_implemented_pack_cache_raises_for_every_method(): + cache = NotImplementedPackCache() + pack = _sample_pack() + + with pytest.raises(NotImplementedError): + cache.is_installed(pack) + with pytest.raises(NotImplementedError): + cache.status(pack) + with pytest.raises(NotImplementedError): + cache.install(pack) + with pytest.raises(NotImplementedError): + cache.verify(pack) + with pytest.raises(NotImplementedError): + cache.remove(pack) + + +@pytest.mark.parametrize("category", [c for c in PackCategory]) +def test_all_pack_categories_constructible(category): + pack = _sample_pack(pack_id=f"pack-{category.value}", category=category) + assert pack.category == category diff --git a/docs/plans/partial/local-first-roadmap.md b/docs/plans/partial/local-first-roadmap.md index 3fe24567..fd31d30f 100644 --- a/docs/plans/partial/local-first-roadmap.md +++ b/docs/plans/partial/local-first-roadmap.md @@ -204,7 +204,7 @@ The roadmap should be broken into implementation issues after this architecture - Desktop shell bootstrap and sidecar lifecycle management. - Local data-store replacement plan. -- Model cache and first-run download UX. +- [Model cache and first-run download UX](./model-cache-design.md) - Mobile PWA installability and pairing flow. - Remote-acceleration trust model and settings UI. diff --git a/docs/plans/partial/model-cache-design.md b/docs/plans/partial/model-cache-design.md new file mode 100644 index 00000000..a86d7a7e --- /dev/null +++ b/docs/plans/partial/model-cache-design.md @@ -0,0 +1,129 @@ +# Model Pack Cache & Download Design (Installer Mode) + +**Status:** Draft — design + thin interfaces only +**Related:** [Issue #45](https://github.com/Abhash-Chakraborty/Find/issues/45), [local-first-roadmap.md](./local-first-roadmap.md) + +## Goal + +Let Find ship a small installer while still running ML fully locally after +first use. Models are not bundled in the installer and are not downloaded +silently — the user explicitly chooses what to download, sees size/progress, +and can cancel, retry, and work fully offline afterward. + +## Current state (verified against source) + +Find has 5 ML capabilities, each with its own model source and no shared +cache/version concept today: + +| Category | Loader file | Model identifier(s) | +|------------|------------------------------------|--------------------------------------------------------| +| Embeddings | `ml/clip_embedder.py` | `settings.CLIP_MODEL` (`ViT-B-16-SigLIP`) + `settings.CLIP_PRETRAINED` (`webli`), via open_clip | +| Caption | `ml/captioner.py` | `settings.BLIP_MODEL` (`microsoft/Florence-2-base`), via HF `from_pretrained` | +| Objects | `ml/object_detector.py` | `settings.YOLO_MODEL` (`yolo26n.pt`), Ultralytics auto-download | +| OCR | `ml/ocr.py` | PaddleOCR, `lang="en"`, own internal cache dir | +| Faces | `ml/face_detector.py` | InsightFace `antelopev2` (has a known nested-folder extraction quirk already handled in code) | + +`core/model_manager.py` only manages in-memory lifecycle (lazy load, LRU +eviction, idle unload). It has no concept of "on disk," "downloaded," or +"version" — this design adds that layer above it, without touching it. + +## Versioned pack manifest + +Each pack (see `core/model_pack.py::ModelPack`) records: + +- `pack_id`, `category`, `version` +- `source_url`, `license` +- `size_bytes` (shown to user before download) +- `checksum_sha256` (verified after download and on every load) +- `compatible_app_versions` +- `config_key` — mirrors the existing loader `config_key` format already + used by `ModelManager.get_model()`, so a pack swap can invalidate an + in-memory model the same way a config change does today. + +Manifests will initially live as static JSON (or Python data) shipped with +the app; a remote-updatable manifest is out of scope for this PR. + +## First-run selection & download UX + +1. On first run, show each category (Embeddings/Caption/Objects/OCR/Faces) + with its pack size and license. +2. Preflight: check free disk space against sum of selected pack sizes + before starting any download; block with a clear message if insufficient. +3. Per-pack progress (bytes downloaded / total), an overall progress + summary, and a cancel button per pack. +4. Cancel leaves the pack `NOT_INSTALLED` (no partial pack is ever marked + installed). +5. Resume: partial downloads are resumable across app restarts by keeping + the partial file plus a small sidecar state file recording bytes-so-far; + if the sidecar is missing/corrupt, restart the download instead of + guessing. +6. Retry: failed downloads move to `PackStatus.FAILED` with a stored error + and a manual retry action; automatic retry uses capped exponential + backoff (not implemented in this PR). + +## Atomic install & corrupted-cache recovery + +- Download to a temp path inside `MODEL_CACHE_DIR` (e.g. `.part`). +- Verify SHA-256 against the manifest. +- Atomically rename into the final cache path only after verification + succeeds. +- Mark `PackStatus.INSTALLED` only after the rename succeeds. +- On every subsequent load, `PackCache.verify()` re-checks the checksum + cheaply (or on a cadence, to avoid hashing large files every launch); + a failed verify moves the pack to `PackStatus.CORRUPTED`, quarantines + the bad file, and re-offers download instead of crashing the loader. + +## Cache location & offline reuse + +- New setting `MODEL_CACHE_DIR` (`config.py`) — empty string means "resolve + a platform-appropriate app-data/cache directory at the call site." +- New setting `ML_OFFLINE_ONLY` (`config.py`) — when true, a pack cache + implementation must refuse any network call and rely solely on what + `PackCache.is_installed()` confirms is already verified on disk. +- Underlying libraries have their own cache conventions (HF `HF_HOME` for + Florence-2, Ultralytics' own weights directory, PaddleOCR's internal + cache dir, InsightFace's `~/.insightface`). The install-mode cache should + point these at subdirectories of `MODEL_CACHE_DIR` via their respective + env vars, rather than reimplementing per-library caching. This mapping is + a follow-up implementation task, not covered by the thin interfaces in + this PR. + +## Preserving existing dev modes + +- `ML_MODE=mock` / `ML_MODE=full` under Docker are unaffected: the pack + cache is additive infrastructure primarily consumed by future installer + (Tauri) flows, not a requirement for contributor/dev workflows. +- `ML_MODE=remote` is untouched — remote mode does not need local packs. + +## Thin interfaces added in this PR + +See `backend/src/find_api/core/model_pack.py`: + +- `ModelPack` (frozen dataclass) — the manifest shape. +- `PackCategory`, `PackStatus` — enums. +- `PackProgress` — progress/status snapshot. +- `PackCache` (Protocol) — the cache contract: `is_installed`, `status`, + `install`, `verify`, `remove`. +- `NotImplementedPackCache` — placeholder so other code can type against + `PackCache` before a real backend exists. + +**No existing loader in `find_api/ml/` is modified by this PR.** A real +`PackCache` implementation, and wiring loaders to consult it, is explicitly +out of scope here and should be a follow-up issue once this contract is +agreed. + +## Out of scope for this PR + +- Bundling any model weights into the installer. +- Any silent/automatic download. +- Replacing or modifying any of the 5 existing ML loaders. +- A real filesystem-backed `PackCache` implementation. +- Tauri-side UI/IPC for the download screen (tracked separately in the + local-first roadmap's Phase 2). + +## Open questions for maintainer review + +- Should pack manifests be bundled at build time or fetched from a remote + index? (This design assumes bundled/static for the first version.) +- Should `MODEL_CACHE_DIR` also apply to Docker/full mode, or stay + installer-only? \ No newline at end of file