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)

def compute_similarity(image, text):
    image = Image.fromarray(image)  # Convert NumPy array to PIL Image

    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

        # Compute cosine similarity
        similarity_score = (text_embeds @ image_embeds.T).item()

    return similarity_score

# 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."
)

demo.launch()