diff --git a/src/scorers/laion.py b/src/scorers/laion.py index 11c3fb7..d7171c1 100644 --- a/src/scorers/laion.py +++ b/src/scorers/laion.py @@ -17,20 +17,24 @@ class LaionScorer: """Load model weights (downloads ~2GB on first run).""" try: import torch # noqa: F401 - from simple_aesthetics_predictor import AestheticsPredictorV1 # noqa: F401 - from transformers import CLIPProcessor # noqa: F401 + + # Package on PyPI is `simple-aesthetics-predictor`, but the + # importable module is `aesthetics_predictor` (no `simple_`). + from aesthetics_predictor import AestheticsPredictorV2Linear + from transformers import CLIPProcessor 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", + # V2 trained on SAC + LAION-Logos + AVA (current best per research) + model_id = "shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE" + + self.model = AestheticsPredictorV2Linear.from_pretrained( + model_id, cache_dir=self._cache_dir, ) self.processor = CLIPProcessor.from_pretrained( - "shunk031/aesthetics-predictor-v2-vit-large-patch14", + model_id, cache_dir=self._cache_dir, ) self.model.eval()