fix: self-contained LAION V2 (CLIP via transformers + official .pth head), no remote code, no stale wrapper
This commit is contained in:
Kareem Horstink
+5
-5
@@ -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
|
||||
+98
-30
@@ -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.clip_model = None # type: ignore[assignment]
|
||||
self.processor = None # type: ignore[assignment]
|
||||
self.head = None # type: ignore[assignment]
|
||||
Reference in New Issue
Block a user