From 016a68e0a200df476703a1835a7b3743674f6b4e Mon Sep 17 00:00:00 2001 From: Kareem Horstink Date: Sun, 23 Aug 2026 17:31:13 +0000 Subject: [PATCH] fix: self-contained LAION V2 (CLIP via transformers + official .pth head), no remote code, no stale wrapper --- requirements.txt | 12 ++-- src/scorers/laion.py | 130 ++++++++++++++++++++++++++++++++----------- 2 files changed, 105 insertions(+), 37 deletions(-) diff --git a/requirements.txt b/requirements.txt index e936b0d..384a59f 100644 --- a/requirements.txt +++ b/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), -# so we do NOT depend on the stale `simple-aesthetics-predictor` wrapper -# (breaks with modern transformers). torch is required by transformers. +# LAION V2 scorer uses OpenAI CLIP via transformers (stable API, no remote +# code) plus a small linear head loaded from the official .pth. This avoids +# 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 -numpy \ No newline at end of file +tqdm \ No newline at end of file diff --git a/src/scorers/laion.py b/src/scorers/laion.py index 188c201..38b9bce 100644 --- a/src/scorers/laion.py +++ b/src/scorers/laion.py @@ -1,13 +1,50 @@ """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`). +Self-contained implementation following the canonical +christophschuhmann/improved-aesthetic-predictor repo: + +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) - MODEL_ID = "shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE" + # Official V2 weights (trained on SAC + LAION-Logos + AVA) from the + # 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 AutoProcessor, AutoModel + from transformers import CLIPModel, CLIPProcessor - 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.MODEL_ID, - cache_dir=self._cache_dir, - trust_remote_code=True, - ) - self.processor = AutoProcessor.from_pretrained( - self.MODEL_ID, + self.clip_model = CLIPModel.from_pretrained( + self.CLIP_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( - " Run: pip install simple-aesthetics-predictor transformers " - "torch torchvision Pillow tqdm" - ) + print(" Run: pip install -r requirements.txt") 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 - from PIL import Image - - if self.processor is None or self.model is None: + if self.processor is None or self.clip_model is None or self.head 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) - score = outputs.logits.squeeze().item() + embeds = self.clip_model.get_image_features(**inputs) + 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.processor = None # type: ignore[assignment] \ No newline at end of file + self.clip_model = None # type: ignore[assignment] + self.processor = None # type: ignore[assignment] + self.head = None # type: ignore[assignment] \ No newline at end of file