app.py

import gradio as gr
from transformers import AutoProcessor, AutoModelForVision2Seq
from PIL import Image
import torch

# Load LLaVA (Apache-2.0 license, safe to use)
model_id = "llava-hf/llava-1.5-7b-hf"

processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForVision2Seq.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
).to("cuda")

def analyze_chart(image, question="Describe this chart"):
    # Preprocess
    inputs = processor(images=image, text=question, return_tensors="pt").to("cuda", torch.float16)
    
    # Generate
    output_ids = model.generate(**inputs, max_new_tokens=200)
    response = processor.decode(output_ids[0], skip_special_tokens=True)
    return response

demo = gr.Interface(
    fn=analyze_chart,
    inputs=[gr.Image(type="pil"), gr.Textbox(value="Describe this chart")],
    outputs="text",
    title="Chart Analyzer (LLaVA)",
    description="Upload a chart and ask questions like 'What does this chart show?' or 'Which month has highest sales?'."
)

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