import gradio as gr
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Set default device to CUDA for GPU acceleration
device = 'cuda' if torch.cuda.is_available() else "cpu"
# torch.set_default_device("cuda")

# Initialize the model and tokenizer
model = AutoModelForCausalLM.from_pretrained("ManishThota/Sparrow").to(device)
tokenizer = AutoTokenizer.from_pretrained("ManishThota/Sparrow", trust_remote_code=True)

def predict_answer(image, question):
    # Convert PIL image to RGB if not already
    image = image.convert("RGB")
    
    # # Format the text input for the model
    # text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{question} ASSISTANT:"
    
    # Tokenize the text input
    encoding = tokenizer(image, question, return_tensors='pt').to(device)

    out = model.generate(**encoding)
    # Preprocess the image for the model
    generated_text = tokenizer.decode(out[0], skip_special_tokens=True)
    
    # # Generate the answer
    # output_ids = model.generate(
    #     input_ids,
    #     max_new_tokens=100,
    #     images=image_tensor,
    #     use_cache=True)[0]
    
    # # Decode the generated tokens to get the answer
    # answer = tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
    
    return generated_text

def gradio_predict(image, question):
    answer = predict_answer(image, question)
    return answer

# Define the Gradio interface
iface = gr.Interface(
    fn=gradio_predict,
    inputs=[gr.Image(type="pil", label="Upload or Drag an Image"), gr.Textbox(label="Question", placeholder="e.g. What are the colors of the bus in the image?", scale=4)],
    outputs=gr.TextArea(label="Answer"),
    title="Sparrow-based Visual Question Answering",
    description="An interactive chat model that can answer questions about images.",
)

# Launch the app
iface.queue().launch(debug=True)