import streamlit as st
import torch
from PIL import Image
import requests
from transformers import BlipForConditionalGeneration, BlipForQuestionAnswering, BlipProcessor


## Loading BLIP image caption model
processor = BlipProcessor.from_pretrained("adit94/blip_humour")
model = BlipForConditionalGeneration.from_pretrained("adit94/blip_humour")

device = "cuda" if torch.cuda.is_available() else "cpu"


def blip_generate(img):
    #raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
    img = img.convert('RGB')
    inputs = processor(img, return_tensors="pt").to(device)
    
    outputs = model.generate(
        **inputs,
        do_sample=False,
        num_beams=5,
        max_length=256,
        min_length=1,
        top_p=0.9,
        repetition_penalty=1.5,
        length_penalty=1.0,
        temperature=1,
    )

    generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()

    return generated_text

st.set_page_config(
    page_title="Humour Detection",
    initial_sidebar_state = 'auto'
)
st.title(":new_moon_with_face: Humour Detection")
st.divider()

st.text("Please upload an image file")
file = st.file_uploader("", type=["jpg", "png"])

if file is not None:
    image = Image.open(file)
    st.image(image)

    caption = blip_generate(image)

    st.header('Predicted output')
    st.text(caption)