fix: extract pooler_output from BaseModelOutputWithPooling in get_image_features (newer transformers)
This commit is contained in:
@@ -131,7 +131,10 @@ class LaionScorer:
|
|||||||
inputs = {k: v.to(self._device) for k, v in inputs.items()}
|
inputs = {k: v.to(self._device) for k, v in inputs.items()}
|
||||||
|
|
||||||
with torch.no_grad():
|
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)
|
embeds = nn.functional.normalize(embeds, dim=-1)
|
||||||
score = self.head(embeds).squeeze().item()
|
score = self.head(embeds).squeeze().item()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user