Browse Source

add img embedding.py

luojunhui 4 weeks ago
parent
commit
5de7d95178
1 changed files with 4 additions and 1 deletions
  1. 4 1
      applications/clip_embedding/embedding.py

+ 4 - 1
applications/clip_embedding/embedding.py

@@ -32,7 +32,10 @@ async def embed_image_url(img_url_list: List[str]):
         inputs = image_processor(images=chunk, return_tensors="pt")
         inputs = {k: v.to(DEVICE, dtype=DTYPE) if hasattr(v, "to") else v for k, v in inputs.items()}
 
-        feats = model.get_image_features(**inputs)
+        outputs = model(**inputs)
+        # 某些实现是 outputs.last_hidden_state,某些是 outputs.image_embeds
+        feats = outputs.image_embeds if hasattr(outputs, "image_embeds") else outputs.last_hidden_state
+
         feats = _normalize(feats)
         outputs.extend(_to_list(feats))