From 2d6d9085f550bb9ed1375bf96eb83bcca0f6a0e3 Mon Sep 17 00:00:00 2001 From: Kareem Horstink Date: Sun, 23 Aug 2026 15:36:10 +0000 Subject: [PATCH] feat: add local model cache directory --- models/.gitkeep | 0 src/scorers/laion.py | 12 ++++++++---- 2 files changed, 8 insertions(+), 4 deletions(-) create mode 100644 models/.gitkeep diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/src/scorers/laion.py b/src/scorers/laion.py index 2b2186b..c5c8bd3 100644 --- a/src/scorers/laion.py +++ b/src/scorers/laion.py @@ -9,10 +9,11 @@ class LaionScorer: name = "LAION Aesthetic Predictor V2" description = "CLIP-based aesthetic scoring (0-10 scale)" - def __init__(self) -> None: + def __init__(self, cache_dir: str = "models") -> None: self.model = None # type: ignore[assignment] self.processor = None # type: ignore[assignment] self._version = "v2" + self._cache_dir = cache_dir def load(self) -> None: """Load model weights (downloads ~2GB on first run).""" @@ -21,15 +22,18 @@ class LaionScorer: from simple_aesthetics_predictor import AestheticsPredictorV1 # noqa: F401 from transformers import CLIPProcessor # noqa: F401 - print(" Loading LAION V2 model... (first run downloads ~2GB)") + print(f" Loading LAION V2 model... (first run downloads ~2GB)") + print(f" Cache directory: {self._cache_dir}") from simple_aesthetics_predictor import AestheticsPredictorV1 from transformers import CLIPProcessor self.model = AestheticsPredictorV1.from_pretrained( - "shunk031/aesthetics-predictor-v2-vit-large-patch14" + "shunk031/aesthetics-predictor-v2-vit-large-patch14", + cache_dir=self._cache_dir, ) self.processor = CLIPProcessor.from_pretrained( - "shunk031/aesthetics-predictor-v2-vit-large-patch14" + "shunk031/aesthetics-predictor-v2-vit-large-patch14", + cache_dir=self._cache_dir, ) self.model.eval() print(" Model loaded.")