import gradio as gr
from PIL import Image
from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration

# SMART RUNTIME LINK: Change 'your-username/my-chart-weights' to your exact HF paths
model_repo_path = "your-username/my-chart-weights"

print("Downloading and caching model weights from model repo...")
processor = Pix2StructProcessor.from_pretrained(model_repo_path)
model = Pix2StructForConditionalGeneration.from_pretrained(model_repo_path)
print("? Model loaded into 16GB RAM successfully!")

def extract_chart_data(img):
    image = Image.fromarray(img).convert("RGB")
    inputs = processor(images=image, text="Generate underlying data table of the figure below:", return_tensors="pt")
    generated_ids = model.generate(**inputs, max_new_tokens=512)
    extracted_table = processor.decode(generated_ids[0], skip_special_tokens=True)
    return extracted_table

demo = gr.Interface(fn=extract_chart_data, inputs="image", outputs="text")
demo.launch()