import gradio as gr
import torch.nn.functional as F
# from PIL import Image
import uform

model = uform.get_model('unum-cloud/uform-vl-english')

def find_score(img, txt, if_fine_grained):
    txt = model.preprocess_text(txt)
    img = model.preprocess_image(img)
    txt_features, txt_emb = model.encode_text(txt, return_features = True)
    img_features, img_emb = model.encode_image(img, return_features = True)
    if if_fine_grained:
        joint_embedding = model.encode_multimodal(
        image_features=img_features,
        text_features=txt_features,
        attention_mask=txt['attention_mask'])
        score = model.get_matching_scores(joint_embedding)
    else:
        score = F.cosine_similarity(txt_emb, img_emb)
    return score

def find_score_img(img1, img2, if_fine_grained):
    img1 = model.preprocess_image(img1)
    img2 = model.preprocess_image(img2)
    img_features1, img_emb1 = model.encode_image(img1, return_features = True)
    img_features2, img_emb2 = model.encode_image(img2, return_features = True)
    if if_fine_grained:
        joint_embedding = model.encode_multimodal(
        image_features=img_features1,
        text_features=img_features2,
        attention_mask=img1['attention_mask'])
        score = model.get_matching_scores(joint_embedding)
    else:
        score = F.cosine_similarity(img_emb1, img_emb2)
    return score

def find_score_txt(txt1, txt2, if_fine_grained):
    txt1 = model.preprocess_text(txt1)
    txt2 = model.preprocess_text(txt2)
    txt_features1, txt_emb1 = model.encode_text(txt1, return_features = True)
    txt_features2, txt_emb2 = model.encode_text(txt2, return_features = True)
    if if_fine_grained:
        joint_embedding = model.encode_multimodal(
        image_features=txt_features1,
        text_features=txt_features2,
        attention_mask=txt1['attention_mask'])
        score = model.get_matching_scores(joint_embedding)
    else:
        score = F.cosine_similarity(txt_emb1, txt_emb2)
    return score
    
with gr.Blocks(theme = gr.themes.Glass()) as demo_mix:
    gr.Markdown('# Find similarity between images and text.')
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(source = 'upload', type = 'pil', label = "Drop your image here", shape = [256, 256])
        with gr.Column():
            txt_input = gr.Textbox(label = 'Enter your text here:' , lines = 1,)
            if_fine_grained = gr.Checkbox(label = "Check for a more fine-grained comparison")
            btn = gr.Button("Find similarity")
    score = gr.Number(label='Similarity score')
    btn.click(find_score, inputs=[img_input, txt_input, if_fine_grained], outputs=[score])
    gr.Markdown('If the box for a more fine-grained comparison is checked, the similarity score will be in (0,1), otherwise, it will be in (-1,1)')
    gr.Markdown('### Image examples')
    gr.Examples(['imgs/red_panda.jpg', 'imgs/trash_raccoon.jpg', 'imgs/baby_panda.jpg', 'imgs/rocket.jpg'], inputs=[img_input])
    gr.Markdown('### Text examples')
    gr.Examples(['baby red panda staring into the camera', \
        'trash raccoon peeking from a trash bin', 
        'a cartoonish raccoon wearing a blue and red suit', 
        "a person holding a baby panda"], inputs=[txt_input])  


with gr.Blocks() as demo_img:
    gr.Markdown('# Find similarity between images.')
    with gr.Row():
        with gr.Column():
            img_input1 = gr.Image(source = 'upload', type = 'pil', label = "Drop your image here", shape = [256, 256])
            if_fine_grained = gr.Checkbox(label = "Check for a more fine-grained comparison")
        with gr.Column():
            img_input2 = gr.Image(source = 'upload', type = 'pil', label = "Drop your image here", shape = [256, 256])
            btn = gr.Button("Find similarity")
            score = gr.Number(label='Similarity score')
    btn.click(find_score_img, inputs=[img_input1, img_input2, if_fine_grained], outputs=[score])
    gr.Markdown('If the box for a more fine-grained comparison is checked, the similarity score will be in (0,1), otherwise, it will be in (-1,1)')
    gr.Markdown('### Image examples')
    gr.Examples(['imgs/red_panda.jpg', 'imgs/trash_raccoon.jpg', 'imgs/baby_panda.jpg', 'imgs/rocket.jpg'], inputs=[img_input1])


with gr.Blocks() as demo_txt:
    gr.Markdown('# Find similarity between short descriptions.')
    with gr.Row():
        with gr.Column():
            txt_input1 = gr.Textbox(label = 'Enter your text here:' , lines = 1)
            txt_input2 = gr.Textbox(label = 'Enter your text here:' , lines = 1)
        with gr.Column():
            if_fine_grained = gr.Checkbox(label = "Check for a more fine-grained comparison")
            btn = gr.Button("Find similarity")
            score = gr.Number(label='Similarity score')
    btn.click(find_score_txt, inputs=[img_input, txt_input, if_fine_grained], outputs=[score])
    gr.Markdown('If the box for a more fine-grained comparison is checked, the similarity score will be in (0,1), otherwise, it will be in (-1,1)')
    gr.Markdown('### Text examples')
    gr.Examples(['baby red panda staring into the camera', \
        'trash raccoon peeking from a trash bin', 
        'a cartoonish raccoon wearing a blue and red suit', 
        "a person holding a baby panda"], inputs=[txt_input1]) 

demo = gr.TabbedInterface([demo_mix, demo_img, demo_txt], ["img2txt", "img2img", "txt2txt"])

if __name__ == "__main__":
    demo.launch(share=True)