Update app.py

app.py

@@ -1,26 +1,14 @@
-import io, base64
+# app.py
+import os, io, base64
 import gradio as gr
-import spaces
 from PIL import Image
 import torch
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    AutoImageProcessor,
-)
+from transformers import AutoTokenizer, AutoImageProcessor, AutoModelForCausalLM
 
 MODEL_ID = "starvector/starvector-8b-im2svg"
 
-# Load separate components explicitly
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    use_fast=False,              # ensure GPT2Tokenizer (python) not Siglip fast
-    trust_remote_code=True
-)
-image_processor = AutoImageProcessor.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True
-)
+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,
@@ -29,37 +17,28 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 ).eval()
 
-# Safety: some GPT2 checkpoints have no pad token
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-def _prep_inputs(image: Image.Image | None, text: str):
+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
-    return batch
-
-@spaces.GPU(duration=180)
-def run_starvector(image: Image.Image | None, text: str) -> str:
-    inputs = _prep_inputs(image, text)
-    # move tensors to the model device(s)
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    batch = {k: v.to(model.device) for k, v in batch.items()}
     with torch.no_grad():
         out = model.generate(
-            **inputs,
+            **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
+            eos_token_id=tokenizer.eos_token_id,
         )
-    svg = tokenizer.batch_decode(out, skip_special_tokens=True)[0]
-    return svg
+    return tokenizer.batch_decode(out, skip_special_tokens=True)[0]
 
-# --- Your UI wiring (example) ---
 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)")
@@ -68,9 +47,6 @@ with gr.Blocks(title="StarVector: Image/Text → SVG") as demo:
     code = gr.Code(label="SVG Output", language="xml")
     btn.click(run_starvector, [img, txt], code)
 
-import os
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    # Gradio will listen on 0.0.0.0 for Spaces
     demo.launch(server_name="0.0.0.0", server_port=port)
-