Update app.py

app.py

@@ -1,40 +1,24 @@
 import gradio as gr
-from transformers import CLIPModel, CLIPFeatureExtractor, BertTokenizer
+from transformers import AutoModel
 from PIL import Image
 import torch
 import torch.nn.functional as F
+import requests
+from io import BytesIO
 
-# Load model and processors separately
-model_name = "jinaai/jina-clip-v1"
-model = CLIPModel.from_pretrained(model_name)
-feature_extractor = CLIPFeatureExtractor.from_pretrained(model_name)
-tokenizer = BertTokenizer.from_pretrained(model_name)
+# Load model with remote code support
+model = AutoModel.from_pretrained('jinaai/jina-clip-v1', trust_remote_code=True)
 
 def compute_similarity(image, text):
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
 
-    # Process image
-    image_inputs = feature_extractor(images=image, return_tensors="pt")
-
-    # Process text (Remove `token_type_ids`)
-    text_inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    text_inputs.pop("token_type_ids", None)  
-
     with torch.no_grad():
-        # Extract embeddings
-        image_embeds = model.get_image_features(**image_inputs)
-        text_embeds = model.get_text_features(**text_inputs)
-
-        # Print to debug
-        print("Image Embedding:", image_embeds)
-        print("Text Embedding:", text_embeds)
-
-        # Normalize embeddings
-        image_embeds = F.normalize(image_embeds, p=2, dim=-1)
-        text_embeds = F.normalize(text_embeds, p=2, dim=-1)
+        # Encode text and image using JinaAI CLIP model
+        text_embeds = model.encode_text([text])  # Expecting list input
+        image_embeds = model.encode_image([image])  # Expecting list input
 
         # Compute cosine similarity
-        similarity_score = (image_embeds @ text_embeds.T).item()
+        similarity_score = (text_embeds @ image_embeds.T).item()
 
     return similarity_score
 
@@ -47,4 +31,4 @@ demo = gr.Interface(
     description="Upload an image and enter a text prompt to get the similarity score."
 )
 
-demo.launch()
+demo.launch()