Imrayya/photo_judgers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add local model cache directory

Browse Source
This commit is contained in:
Kareem Horstink
2026-08-23 15:36:10 +00:00
parent 4b28ac6de0
commit 2d6d9085f5
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/.gitkeep
View File
src/scorers/laion.py
+8 -4
View File
@@ -9,10 +9,11 @@ class LaionScorer:
name = "LAION Aesthetic Predictor V2" name = "LAION Aesthetic Predictor V2"
description = "CLIP-based aesthetic scoring (0-10 scale)" 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.model = None # type: ignore[assignment]
self.processor = None # type: ignore[assignment] self.processor = None # type: ignore[assignment]
self._version = "v2" self._version = "v2"
self._cache_dir = cache_dir
def load(self) -> None: def load(self) -> None:
"""Load model weights (downloads ~2GB on first run).""" """Load model weights (downloads ~2GB on first run)."""
@@ -21,15 +22,18 @@ class LaionScorer:
from simple_aesthetics_predictor import AestheticsPredictorV1 # noqa: F401 from simple_aesthetics_predictor import AestheticsPredictorV1 # noqa: F401
from transformers import CLIPProcessor # 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 simple_aesthetics_predictor import AestheticsPredictorV1
from transformers import CLIPProcessor from transformers import CLIPProcessor
self.model = AestheticsPredictorV1.from_pretrained( 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( 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() self.model.eval()
print(" Model loaded.") print(" Model loaded.")