|
|
|
|
|
import os, io, base64 |
|
|
import gradio as gr |
|
|
from PIL import Image |
|
|
import torch |
|
|
from transformers import AutoTokenizer, AutoImageProcessor, AutoModelForCausalLM |
|
|
|
|
|
MODEL_ID = "starvector/starvector-8b-im2svg" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False, trust_remote_code=True) |
|
|
image_processor = AutoImageProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
MODEL_ID, |
|
|
torch_dtype=torch.bfloat16, |
|
|
low_cpu_mem_usage=True, |
|
|
device_map="auto", |
|
|
trust_remote_code=True |
|
|
).eval() |
|
|
|
|
|
if tokenizer.pad_token_id is None: |
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
def run_starvector(image: Image.Image | None, text: str) -> str: |
|
|
text = text or "" |
|
|
toks = tokenizer(text, return_tensors="pt", add_special_tokens=True) |
|
|
batch = {"input_ids": toks.input_ids} |
|
|
if image is not None: |
|
|
pix = image_processor(images=image, return_tensors="pt").pixel_values |
|
|
batch["pixel_values"] = pix |
|
|
batch = {k: v.to(model.device) for k, v in batch.items()} |
|
|
with torch.no_grad(): |
|
|
out = model.generate( |
|
|
**batch, |
|
|
max_new_tokens=2048, |
|
|
temperature=0.2, |
|
|
do_sample=False, |
|
|
pad_token_id=tokenizer.pad_token_id, |
|
|
eos_token_id=tokenizer.eos_token_id, |
|
|
) |
|
|
return tokenizer.batch_decode(out, skip_special_tokens=True)[0] |
|
|
|
|
|
with gr.Blocks(title="StarVector: Image/Text → SVG") as demo: |
|
|
gr.Markdown("# StarVector: Image/Text → SVG") |
|
|
img = gr.Image(type="pil", label="Upload image (optional)") |
|
|
txt = gr.Textbox(label="Text prompt (optional)") |
|
|
btn = gr.Button("Generate SVG") |
|
|
code = gr.Code(label="SVG Output", language="xml") |
|
|
btn.click(run_starvector, [img, txt], code) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
port = int(os.environ.get("PORT", 7860)) |
|
|
demo.launch(server_name="0.0.0.0", server_port=port) |
|
|
|
