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

# Load the smaller StarVector model (lighter, runs fine on free Hugging Face hardware)
model_id = "starvector/starvector-1b-im2svg"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16)

def convert_to_svg(image):
    # Prepare the image and run inference
    inputs = tokenizer(image, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=1024)
    svg_code = tokenizer.decode(outputs[0])
    return svg_code

demo = gr.Interface(
    fn=convert_to_svg,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Code(language="svg", label="Generated SVG Code"),
    title="StarVector Image → SVG Converter"
)

demo.launch()