import gradio as gr
from PIL import Image
from transformers import CLIPModel, AutoTokenizer, AutoProcessor
import torch

# Ensure required dependencies are installed
try:
    import timm
except ImportError:
    import subprocess
    subprocess.run(["pip", "install", "timm"], check=True)

# Load Jina CLIP model with trust_remote_code=True
model_name = "jinaai/jina-clip-v1"
model = CLIPModel.from_pretrained(model_name, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)

def compute_similarity(input1, input2, text1, text2, type1, type2):
    # Process input1
    if type1 == "Image":
        if not input1:
            return "Error: No image provided for Input 1"
        image1 = Image.open(input1).convert("RGB")
        input1_tensor = processor(images=image1, return_tensors="pt")["pixel_values"]
    elif type1 == "Text":
        if not text1.strip():
            return "Error: No text provided for Input 1"
        input1_tensor = tokenizer(text1, return_tensors="pt")["input_ids"]
    else:
        return "Error: Invalid input type for Input 1"
    
    # Process input2
    if type2 == "Image":
        if not input2:
            return "Error: No image provided for Input 2"
        image2 = Image.open(input2).convert("RGB")
        input2_tensor = processor(images=image2, return_tensors="pt")["pixel_values"]
    elif type2 == "Text":
        if not text2.strip():
            return "Error: No text provided for Input 2"
        input2_tensor = tokenizer(text2, return_tensors="pt")["input_ids"]
    else:
        return "Error: Invalid input type for Input 2"
    
    # Compute embeddings
    with torch.no_grad():
        if type1 == "Image":
            embedding1 = model.get_image_features(pixel_values=input1_tensor)
        else:
            embedding1 = model.get_text_features(input_ids=input1_tensor)
        
        if type2 == "Image":
            embedding2 = model.get_image_features(pixel_values=input2_tensor)
        else:
            embedding2 = model.get_text_features(input_ids=input2_tensor)
    
    # Normalize embeddings
    embedding1 = embedding1 / embedding1.norm(dim=-1, keepdim=True)
    embedding2 = embedding2 / embedding2.norm(dim=-1, keepdim=True)
    
    # Compute cosine similarity
    similarity = torch.nn.functional.cosine_similarity(embedding1, embedding2).item()
    return f"Similarity Score: {similarity:.4f}"

with gr.Blocks() as demo:
    gr.Markdown("# CLIP-based Similarity Comparison")
    
    with gr.Row():
        type1 = gr.Radio(["Image", "Text"], label="Input 1 Type", value="Image")
        type2 = gr.Radio(["Image", "Text"], label="Input 2 Type", value="Text")
    
    with gr.Row():
        input1 = gr.Image(type="filepath", label="Upload Image 1")
        input2 = gr.Image(type="filepath", label="Upload Image 2")
        text1 = gr.Textbox(label="Enter Text 1")
        text2 = gr.Textbox(label="Enter Text 2")
    
    compare_btn = gr.Button("Compare")
    output = gr.Textbox(label="Similarity Score")
    
    compare_btn.click(
        compute_similarity, 
        inputs=[
            input1,
            input2,
            text1,
            text2,
            type1,
            type2
        ], 
        outputs=output
    )

demo.launch()