diff --git a/requirements.txt b/requirements.txt index e6a30c5..e936b0d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,13 +1,13 @@ # Photo Judgers — Requirements # Install in conda env: pip install -r requirements.txt # -# simple-aesthetics-predictor automatically installs torch + transformers. -# We avoid pinning torchvision/torchaudio here to prevent version mismatch -# with the torch that simple-aesthetics-predictor resolves. Pillow is needed -# for image loading; tqdm for progress bars. +# We load the LAION V2 model directly via transformers (trust_remote_code), +# so we do NOT depend on the stale `simple-aesthetics-predictor` wrapper +# (breaks with modern transformers). torch is required by transformers. -simple-aesthetics-predictor -transformers +transformers>=4.40 +torch>=2.1 +torchvision Pillow tqdm numpy \ No newline at end of file diff --git a/src/scorers/laion.py b/src/scorers/laion.py index d7171c1..188c201 100644 --- a/src/scorers/laion.py +++ b/src/scorers/laion.py @@ -1,12 +1,23 @@ -"""LAION Aesthetic Predictor V2 scorer.""" +"""LAION Aesthetic Predictor V2 scorer. + +Loads the V2 model directly from HuggingFace via transformers' +`trust_remote_code` mechanism. This avoids the `simple-aesthetics-predictor` +wrapper package, which is stale (last release Dec 2024) and breaks with +modern transformers (circular import in `dependency_versions_check`). +""" from typing import Optional + + class LaionScorer: """LAION V2 aesthetic scorer — CLIP-based, 0-10 scale.""" name = "LAION Aesthetic Predictor V2" description = "CLIP-based aesthetic scoring (0-10 scale)" + # V2 trained on SAC + LAION-Logos + AVA (current best per research) + MODEL_ID = "shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE" + def __init__(self, cache_dir: str = "models") -> None: self.model = None # type: ignore[assignment] self.processor = None # type: ignore[assignment] @@ -17,24 +28,18 @@ class LaionScorer: """Load model weights (downloads ~2GB on first run).""" try: import torch # noqa: F401 + from transformers import AutoProcessor, AutoModel - # 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(" Loading LAION V2 model... (first run downloads ~2GB)") print(f" Cache directory: {self._cache_dir}") - # 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, + self.model = AutoModel.from_pretrained( + self.MODEL_ID, cache_dir=self._cache_dir, + trust_remote_code=True, ) - self.processor = CLIPProcessor.from_pretrained( - model_id, + self.processor = AutoProcessor.from_pretrained( + self.MODEL_ID, cache_dir=self._cache_dir, ) self.model.eval() @@ -74,4 +79,4 @@ class LaionScorer: def unload(self) -> None: """Clean up model references.""" self.model = None # type: ignore[assignment] - self.processor = None # type: ignore[assignment] + self.processor = None # type: ignore[assignment] \ No newline at end of file