Update app.py

app.py

@@ -2,33 +2,72 @@ import gradio as gr
 from transformers import AutoModel
 from PIL import Image
 import torch
-import torch.nn.functional as F
-import requests
-from io import BytesIO
 
-# Load model with remote code support
-model = AutoModel.from_pretrained('jinaai/jina-clip-v1', trust_remote_code=True)
+# Load JinaAI CLIP model
+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
+# Function to compute similarity
+def compute_similarity(input1, input2, input1_type, input2_type):
+    # Check if inputs are empty
+    if (input1_type == "Text" and not input1.strip()) or (input1_type == "Image" and input1 is None):
+        return "Error: Input 1 is empty!"
+    if (input2_type == "Text" and not input2.strip()) or (input2_type == "Image" and input2 is None):
+        return "Error: Input 2 is empty!"
+    
+    inputs = []
+    
+    # Process first input
+    if input1_type == "Text":
+        text1_embedding = model.encode_text([input1])
+        inputs.append(text1_embedding)
+    elif input1_type == "Image":
+        image1_embedding = model.encode_image([Image.fromarray(input1)])
+        inputs.append(image1_embedding)
+    
+    # Process second input
+    if input2_type == "Text":
+        text2_embedding = model.encode_text([input2])
+        inputs.append(text2_embedding)
+    elif input2_type == "Image":
+        image2_embedding = model.encode_image([Image.fromarray(input2)])
+        inputs.append(image2_embedding)
+    
+    # Compute cosine similarity
+    similarity_score = (inputs[0] @ inputs[1].T).item()
+    
+    return similarity_score
 
-    with torch.no_grad():
-        # 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
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Multimodal Similarity: Text-Text, Text-Image, Image-Image")
 
-        # Compute cosine similarity
-        similarity_score = (text_embeds @ image_embeds.T).item()
+    with gr.Row():
+        input1_type = gr.Radio(["Text", "Image"], label="Input 1 Type", value="Text")
+        input2_type = gr.Radio(["Text", "Image"], label="Input 2 Type", value="Image")
 
-    return similarity_score
+    with gr.Row():
+        input1 = gr.Textbox(label="Text Input 1", visible=True)
+        image1 = gr.Image(type="numpy", label="Image Input 1", visible=False)
+    
+    with gr.Row():
+        input2 = gr.Textbox(label="Text Input 2", visible=False)
+        image2 = gr.Image(type="numpy", label="Image Input 2", visible=True)
 
-# Gradio UI
-demo = gr.Interface(
-    fn=compute_similarity,
-    inputs=[gr.Image(type="numpy"), gr.Textbox(label="Enter text")],
-    outputs=gr.Number(label="Similarity Score"),
-    title="JinaAI CLIP Image-Text Similarity",
-    description="Upload an image and enter a text prompt to get the similarity score."
-)
+    output = gr.Textbox(label="Similarity Score / Error", interactive=False)
+
+    # Function to toggle visibility based on selected types
+    def update_visibility(input1_type, input2_type):
+        return (
+            input1_type == "Text",
+            input1_type == "Image",
+            input2_type == "Text",
+            input2_type == "Image"
+        )
+
+    input1_type.change(update_visibility, inputs=[input1_type, input2_type], outputs=[input1, image1, input2, image2])
+    input2_type.change(update_visibility, inputs=[input1_type, input2_type], outputs=[input1, image1, input2, image2])
+
+    btn = gr.Button("Compute Similarity")
+    btn.click(compute_similarity, inputs=[input1, input2, input1_type, input2_type], outputs=output)
 
 demo.launch()