fix: extract pooler_output from BaseModelOutputWithPooling in get_image_features (newer transformers)

2026-08-23 18:15:50 +00:00
parent 016a68e0a2
commit af26f3c369
1 changed files with 4 additions and 1 deletions
@@ -131,7 +131,10 @@ class LaionScorer:
            inputs = {k: v.to(self._device) for k, v in inputs.items()}

            with torch.no_grad():
-                embeds = self.clip_model.get_image_features(**inputs)
+                out = self.clip_model.get_image_features(**inputs)
+                # Newer transformers return a `BaseModelOutputWithPooling`
+                # object; older ones return the tensor directly. Handle both.
+                embeds = getattr(out, "pooler_output", out)
                embeds = nn.functional.normalize(embeds, dim=-1)
                score = self.head(embeds).squeeze().item()