app.py

import gradio as gr
from PIL import Image
import os
import sys
from llava_inference import LLaVAHelper

# Add error handling for module imports
try:
    model = LLaVAHelper()
except Exception as e:
    print(f"Failed to initialize LLaVA model: {e}")
    # Continue execution to show error in the UI
    model = None

def answer_question(image, question):
    if model is None:
        return "Model initialization failed. Please check server logs."
    
    if image is None or question.strip() == "":
        return "Please upload an image and enter a question."
    
    try:
        return model.generate_answer(image, question)
    except Exception as e:
        return f"Error processing request: {str(e)}"

# Create examples directory if it doesn't exist
os.makedirs("assets", exist_ok=True)

demo = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Image(type="pil", label="Upload Public Transport Signage"),
        gr.Textbox(label="Ask a question (e.g., 'When is the next train to London?')")
    ],
    outputs=gr.Textbox(label="Answer"),
    title="UK Public Transport Assistant",
    description="Upload an image of UK public transport signage (like train timetables or metro maps), and ask a question related to it. Powered by LLaVA-1.5.",
    examples=[
        # Only use examples if the example file exists
        ["assets/example.jpg", "Where is platform 3?"] if os.path.exists("assets/example.jpg") else None
    ]
)

if __name__ == "__main__":
    demo.launch(share=True)  # Added share=True to make it accessible on a public URL