fix: self-contained LAION V2 (CLIP via transformers + official .pth head), no remote code, no stale wrapper

requirements.txt

@@ -1,13 +1,13 @@
 # Photo Judgers — Requirements
 # Install in conda env: pip install -r requirements.txt
 #
-# We load the LAION V2 model directly via transformers (trust_remote_code),
+# LAION V2 scorer uses OpenAI CLIP via transformers (stable API, no remote
-# so we do NOT depend on the stale `simple-aesthetics-predictor` wrapper
+# code) plus a small linear head loaded from the official .pth. This avoids
-# (breaks with modern transformers). torch is required by transformers.
+# the stale `simple-aesthetics-predictor` wrapper and transformers remote
+# code, both of which break on modern transformers.
 
-transformers>=4.40
+transformers>=4.30
 torch>=2.1
 torchvision
 Pillow
-tqdm
+tqdm
-numpy

src/scorers/laion.py

@@ -1,13 +1,50 @@
 """LAION Aesthetic Predictor V2 scorer.
 
-Loads the V2 model directly from HuggingFace via transformers'
+Self-contained implementation following the canonical
-`trust_remote_code` mechanism. This avoids the `simple-aesthetics-predictor`
+christophschuhmann/improved-aesthetic-predictor repo:
-wrapper package, which is stale (last release Dec 2024) and breaks with
+
-modern transformers (circular import in `dependency_versions_check`).
+1. Encode the image with OpenAI CLIP ViT-L/14 (via transformers, stable API).
+2. L2-normalize the 768-d embedding.
+3. Run it through the V2 linear MLP head (loaded from the official .pth).
+
+This deliberately avoids HuggingFace *remote code* for the model (which
+breaks against modern transformers: the hosted config code expects
+transformers ~4.30 and errors on 4.57) and the stale
+`simple-aesthetics-predictor` wrapper package.
 """
 
+import os
+import urllib.request
 from typing import Optional
 
+import torch
+import torch.nn as nn
+from PIL import Image
+
+
+class _MLPHead(nn.Module):
+    """V2 linear head: 768 -> 1024 -> 128 -> 64 -> 16 -> 1 (dropout between).
+
+    Mirrors the `MLP` LightningModule from the improved-aesthetic-predictor
+    repo so the state_dict keys (`layers.*`) line up with the official .pth.
+    """
+
+    def __init__(self, input_size: int = 768) -> None:
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(input_size, 1024),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            nn.Dropout(0.1),
+            nn.Linear(64, 16),
+            nn.Linear(16, 1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
 
 class LaionScorer:
     """LAION V2 aesthetic scorer — CLIP-based, 0-10 scale."""
@@ -15,58 +52,88 @@ class LaionScorer:
     name = "LAION Aesthetic Predictor V2"
     description = "CLIP-based aesthetic scoring (0-10 scale)"
 
-    # V2 trained on SAC + LAION-Logos + AVA (current best per research)
+    # Official V2 weights (trained on SAC + LAION-Logos + AVA) from the
-    MODEL_ID = "shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE"
+    # canonical repo; same URL used by the shunk031 wrapper internally.
+    HEAD_URL = (
+        "https://github.com/christophschuhmann/improved-aesthetic-predictor"
+        "/raw/main/sac%2Blogos%2Bava1-l14-linearMSE.pth"
+    )
+    HEAD_FILENAME = "laion_v2_sac_logos_ava1_l14_linearMSE.pth"
+
+    CLIP_ID = "openai/clip-vit-large-patch14"
 
     def __init__(self, cache_dir: str = "models") -> None:
-        self.model = None  # type: ignore[assignment]
+        self.clip_model = None  # type: ignore[assignment]
         self.processor = None  # type: ignore[assignment]
+        self.head = None  # type: ignore[assignment]
         self._version = "v2"
         self._cache_dir = cache_dir
+        self._device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def _head_path(self) -> str:
+        return os.path.join(self._cache_dir, self.HEAD_FILENAME)
+
+    def _download_head(self) -> None:
+        """Download the official V2 head weights if not already cached."""
+        dest = self._head_path()
+        if os.path.exists(dest):
+            print(f"  Head weights present: {dest}")
+            return
+        os.makedirs(self._cache_dir, exist_ok=True)
+        print(f"  Downloading V2 head weights... (small, ~7MB)")
+        try:
+            urllib.request.urlretrieve(self.HEAD_URL, dest)
+        except Exception as e:  # pragma: no cover - network edge case
+            print(f"  ERROR: Failed to download head weights: {e}")
+            raise
+        print(f"  Saved to: {dest}")
 
     def load(self) -> None:
-        """Load model weights (downloads ~2GB on first run)."""
+        """Load CLIP encoder + V2 head (downloads ~1.7GB CLIP on first run)."""
         try:
-            import torch  # noqa: F401
+            from transformers import CLIPModel, CLIPProcessor
-            from transformers import AutoProcessor, AutoModel
 
-            print("  Loading LAION V2 model... (first run downloads ~2GB)")
+            print("  Loading CLIP ViT-L/14 encoder... (first run downloads ~1.7GB)")
             print(f"  Cache directory: {self._cache_dir}")
 
-            self.model = AutoModel.from_pretrained(
+            self.clip_model = CLIPModel.from_pretrained(
-                self.MODEL_ID,
+                self.CLIP_ID,
-                cache_dir=self._cache_dir,
-                trust_remote_code=True,
-            )
-            self.processor = AutoProcessor.from_pretrained(
-                self.MODEL_ID,
                 cache_dir=self._cache_dir,
             )
-            self.model.eval()
+            self.processor = CLIPProcessor.from_pretrained(
+                self.CLIP_ID,
+                cache_dir=self._cache_dir,
+            )
+            self.clip_model.to(self._device)
+            self.clip_model.eval()
+
+            self._download_head()
+            self.head = _MLPHead(input_size=768)
+            state = torch.load(self._head_path(), map_location=self._device)
+            self.head.load_state_dict(state)
+            self.head.to(self._device)
+            self.head.eval()
+
             print("  Model loaded.")
         except ImportError:
             print("  ERROR: Required packages not installed.")
-            print(
+            print("  Run: pip install -r requirements.txt")
-                "  Run: pip install simple-aesthetics-predictor transformers "
-                "torch torchvision Pillow tqdm"
-            )
             raise
 
     def score(self, image_path: str) -> Optional[dict]:
         """Score a single image. Returns dict with score or None on failure."""
         try:
-            import torch  # noqa: F401
+            if self.processor is None or self.clip_model is None or self.head is None:
-            from PIL import Image
-
-            if self.processor is None or self.model is None:
                 raise RuntimeError("Model not loaded. Call load() first.")
 
             image = Image.open(image_path).convert("RGB")
             inputs = self.processor(images=image, return_tensors="pt")
+            inputs = {k: v.to(self._device) for k, v in inputs.items()}
 
             with torch.no_grad():
-                outputs = self.model(**inputs)
+                embeds = self.clip_model.get_image_features(**inputs)
-                score = outputs.logits.squeeze().item()
+                embeds = nn.functional.normalize(embeds, dim=-1)
+                score = self.head(embeds).squeeze().item()
 
             return {
                 "laion_score": round(score, 2),
@@ -78,5 +145,6 @@ class LaionScorer:
 
     def unload(self) -> None:
         """Clean up model references."""
-        self.model = None  # type: ignore[assignment]
+        self.clip_model = None  # type: ignore[assignment]
-        self.processor = None  # type: ignore[assignment]
+        self.processor = None  # type: ignore[assignment]
+        self.head = None  # type: ignore[assignment]