import streamlit as st
from PIL import Image
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM


st.set_page_config(page_title="Multimodal VQA Chatbot", page_icon="👁️", layout="wide")
st.title("👁️ Multimodal Visual Question Answering Chatbot")
st.markdown("Upload an image and ask questions! This app cleanly connects Vision with Text.")


@st.cache_resource
def load_vqa_models():
    # Load BLIP Vision Components
    blip_id = "Salesforce/blip-image-captioning-base"
    blip_processor = BlipProcessor.from_pretrained(blip_id)
    blip_model = BlipForConditionalGeneration.from_pretrained(blip_id)
    
    # Load LaMini Text Components
    text_model_id = "MBZUAI/LaMini-Flan-T5-248M"
    text_tokenizer = AutoTokenizer.from_pretrained(text_model_id)
    text_model = AutoModelForSeq2SeqLM.from_pretrained(text_model_id)
    
    return blip_processor, blip_model, text_tokenizer, text_model

blip_processor, blip_model, text_tokenizer, text_model = load_vqa_models()


# 3. CREATE TWO COLUMNS ON THE WEB SCREEN
col1, col2 = st.columns([1, 1])

# Column 1: For uploading and viewing the picture
with col1:
    st.subheader("Step 1: Upload Source Image")
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image Source', use_column_width=True)

# Column 2: For typing questions and seeing the answer
with col2:
    st.subheader("Step 2: Ask the AI Chatbot")
    user_question = st.text_input("What would you like to know about this image?")
    
    # When the user clicks the blue button, run this logic:
    if st.button("Analyze & Generate Answer", type="primary"):
        if uploaded_file is not None and user_question.strip():
            with st.spinner("Processing visual components..."):
                try:
                    
                    inputs = blip_processor(images=image, text=user_question,return_tensors="pt")
                    with torch.no_grad():
                        vision_outputs = blip_model.generate(
                            **inputs,
                            no_repeat_ngram_size=3,
                            repetition_penalty=2.2,
                            max_new_tokens=80
                        )
                    clean_context_text = blip_processor.decode(vision_outputs[0], skip_special_tokens=True)
                    st.info(f"**BLIP Extracted Context:** {clean_context_text}")
                    
               
                    prompt = f"Context: {clean_context_text}\nQuestion: {user_question}\nAnswer:"
                    
                  
                    text_inputs = text_tokenizer(prompt, return_tensors="pt")
                    
                  
                    with torch.no_grad():
                        text_outputs = text_model.generate(
                            **text_inputs,
                            max_new_tokens=150,
                            repetition_penalty=2.0,
                            no_repeat_ngram_size=3
                        )
                    
                 
                    final_answer = text_tokenizer.decode(text_outputs[0], skip_special_tokens=True)
                    
            
                    st.success("### AI Chatbot Response:")
                    st.write(final_answer.strip())
                    
                except Exception as e:
                    st.error(f"An unexpected error occurred during execution: {str(e)}")
        else:
            st.warning("⚠️ Please upload an image first and type a question.")