Update app.py

app.py

@@ -1,80 +1,90 @@
 import gradio as gr
 from transformers import AutoModel
 from PIL import Image
-import numpy as np
 import torch
+import numpy as np
 
 # Load JinaAI CLIP model
-model = AutoModel.from_pretrained("jinaai/jina-clip-v1", trust_remote_code=True)
+model = AutoModel.from_pretrained('jinaai/jina-clip-v1', trust_remote_code=True)
 
-# Function to process input
-def process_input(input_data, input_type):
-    if input_type == "Text":
-        return model.encode_text([input_data]) if input_data.strip() else None
-    elif input_type == "Image":
-        if isinstance(input_data, np.ndarray):  # Gradio provides NumPy array for images
-            image = Image.fromarray(input_data)  # Convert NumPy to PIL Image
-            return model.encode_image(image)  # Directly pass image (no list)
-        return None  # If input is not valid
-    return None
+def compute_similarity(input1, input2):
+    """
+    Computes similarity between:
+    - Image and Text
+    - Image and Image
+    - Text and Text
+    """
+    
+    # Detect input types
+    input1_is_text = isinstance(input1, str) and input1.strip() != ""
+    input2_is_text = isinstance(input2, str) and input2.strip() != ""
+    input1_is_image = isinstance(input1, np.ndarray)
+    input2_is_image = isinstance(input2, np.ndarray)
 
-# Function to compute similarity
-def compute_similarity(input1, input2, input1_type, input2_type):
-    # Validate inputs
-    if input1_type == "Text" and not input1.strip():
-        return "Error: Input 1 is empty!"
-    if input2_type == "Text" and not input2.strip():
-        return "Error: Input 2 is empty!"
-    if input1_type == "Image" and input1 is None:
-        return "Error: Image 1 is missing!"
-    if input2_type == "Image" and input2 is None:
-        return "Error: Image 2 is missing!"
+    # Ensure valid input
+    if not (input1_is_text or input1_is_image) or not (input2_is_text or input2_is_image):
+        return "Error: Both inputs must be valid (image or text)!"
 
-    # Process inputs
-    embedding1 = process_input(input1, input1_type)
-    embedding2 = process_input(input2, input2_type)
+    try:
+        with torch.no_grad():
+            if input1_is_text and input2_is_text:
+                # Text-Text Similarity
+                emb1 = model.encode_text([input1])
+                emb2 = model.encode_text([input2])
+            elif input1_is_image and input2_is_image:
+                # Image-Image Similarity
+                image1 = Image.fromarray(input1)
+                image2 = Image.fromarray(input2)
+                emb1 = model.encode_image([image1])
+                emb2 = model.encode_image([image2])
+            else:
+                # Image-Text Similarity
+                if input1_is_image:
+                    image = Image.fromarray(input1)
+                    text = input2
+                    emb1 = model.encode_image([image])
+                    emb2 = model.encode_text([text])
+                else:
+                    image = Image.fromarray(input2)
+                    text = input1
+                    emb1 = model.encode_text([text])
+                    emb2 = model.encode_image([image])
 
-    if embedding1 is None or embedding2 is None:
-        return "Error: Failed to process input!"
+            # Compute cosine similarity
+            similarity_score = (emb1 @ emb2.T).item()
 
-    # Compute cosine similarity
-    similarity_score = (embedding1 @ embedding2.T).item()
-    return f"Similarity Score: {similarity_score:.4f}"
+        return similarity_score
 
-# Function to update UI dynamically
-def update_visibility(input1_type, input2_type):
-    return (
-        gr.update(visible=(input1_type == "Text"), value="" if input1_type == "Image" else None),
-        gr.update(visible=(input1_type == "Image"), value=None),
-        gr.update(visible=(input2_type == "Text"), value="" if input2_type == "Image" else None),
-        gr.update(visible=(input2_type == "Image"), value=None)
-    )
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 # Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## JinaAI CLIP Multimodal Similarity")
-
-    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")
+demo = gr.Interface(
+    fn=compute_similarity,
+    inputs=[
+        gr.Radio(["Text", "Image"], label="Input 1 Type", value="Text"),
+        gr.Textbox(label="Text Input 1", visible=True),
+        gr.Image(type="numpy", label="Image Input 1", visible=False),
 
-    with gr.Row():
-        input1_text = gr.Textbox(label="Text Input 1", visible=True)
-        input1_image = gr.Image(type="numpy", interactive=True, label="Image Input 1", visible=False)
+        gr.Radio(["Text", "Image"], label="Input 2 Type", value="Text"),
+        gr.Textbox(label="Text Input 2", visible=True),
+        gr.Image(type="numpy", label="Image Input 2", visible=False),
+    ],
+    outputs=gr.Textbox(label="Similarity Score / Error", interactive=False),
+    title="JinaAI CLIP Multimodal Similarity",
+    description="Compare similarity between two inputs (Text, Image, or both)."
+)
 
-    with gr.Row():
-        input2_text = gr.Textbox(label="Text Input 2", visible=False)
-        input2_image = gr.Image(type="numpy", interactive=True, label="Image Input 2", visible=True)
-
-    output = gr.Textbox(label="Similarity Score / Error", interactive=False)
-
-    # Toggle visibility of inputs dynamically
-    input1_type.change(update_visibility, inputs=[input1_type, input2_type], 
-                       outputs=[input1_text, input1_image, input2_text, input2_image])
-    input2_type.change(update_visibility, inputs=[input1_type, input2_type], 
-                       outputs=[input1_text, input1_image, input2_text, input2_image])
+# Update visibility dynamically
+def update_visibility(input1_type, input2_type):
+    return (
+        input1_type == "Text",  # Text input 1 visibility
+        input1_type == "Image", # Image input 1 visibility
+        input2_type == "Text",  # Text input 2 visibility
+        input2_type == "Image"  # Image input 2 visibility
+    )
 
-    btn = gr.Button("Compute Similarity")
-    btn.click(compute_similarity, inputs=[input1_text, input2_text, input1_type, input2_type], outputs=output)
+# Add event handlers for input type change
+demo.load(update_visibility, inputs=["Input 1 Type", "Input 2 Type"], outputs=["Text Input 1", "Image Input 1", "Text Input 2", "Image Input 2"])
 
 demo.launch()