fix: LAION V2 import — correct module (aesthetics_predictor), class (V2Linear), and model ID (sac-logos-ava1-l14-linearMSE)
This commit is contained in:
+11
-7
@@ -17,20 +17,24 @@ class LaionScorer:
|
|||||||
"""Load model weights (downloads ~2GB on first run)."""
|
"""Load model weights (downloads ~2GB on first run)."""
|
||||||
try:
|
try:
|
||||||
import torch # noqa: F401
|
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" Loading LAION V2 model... (first run downloads ~2GB)")
|
||||||
print(f" Cache directory: {self._cache_dir}")
|
print(f" Cache directory: {self._cache_dir}")
|
||||||
from simple_aesthetics_predictor import AestheticsPredictorV1
|
|
||||||
from transformers import CLIPProcessor
|
|
||||||
|
|
||||||
self.model = AestheticsPredictorV1.from_pretrained(
|
# V2 trained on SAC + LAION-Logos + AVA (current best per research)
|
||||||
"shunk031/aesthetics-predictor-v2-vit-large-patch14",
|
model_id = "shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE"
|
||||||
|
|
||||||
|
self.model = AestheticsPredictorV2Linear.from_pretrained(
|
||||||
|
model_id,
|
||||||
cache_dir=self._cache_dir,
|
cache_dir=self._cache_dir,
|
||||||
)
|
)
|
||||||
self.processor = CLIPProcessor.from_pretrained(
|
self.processor = CLIPProcessor.from_pretrained(
|
||||||
"shunk031/aesthetics-predictor-v2-vit-large-patch14",
|
model_id,
|
||||||
cache_dir=self._cache_dir,
|
cache_dir=self._cache_dir,
|
||||||
)
|
)
|
||||||
self.model.eval()
|
self.model.eval()
|
||||||
|
|||||||
Reference in New Issue
Block a user