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.")