Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import CLIPModel, CLIPFeatureExtractor, BertTokenizer
 from PIL import Image
 import torch
+import torch.nn.functional as F
 
 # Load model and processors separately
 model_name = "jinaai/jina-clip-v1"
@@ -17,13 +18,20 @@ def compute_similarity(image, text):
 
     # Process text (Remove `token_type_ids`)
     text_inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    text_inputs.pop("token_type_ids", None)  # Remove token_type_ids to avoid TypeError
+    text_inputs.pop("token_type_ids", None)  
 
     with torch.no_grad():
-        outputs = model(**image_inputs, **text_inputs)
-        logits_per_image = outputs.logits_per_image  # Image-to-text similarity
-        similarity_score = logits_per_image.item()
-    
+        # Extract embeddings
+        image_embeds = model.get_image_features(**image_inputs)
+        text_embeds = model.get_text_features(**text_inputs)
+
+        # Normalize embeddings
+        image_embeds = F.normalize(image_embeds, p=2, dim=-1)
+        text_embeds = F.normalize(text_embeds, p=2, dim=-1)
+
+        # Compute cosine similarity
+        similarity_score = (image_embeds @ text_embeds.T).item()
+
     return similarity_score
 
 # Gradio UI