Update app.py

app.py

@@ -1,69 +1,57 @@
+import io, base64
 import gradio as gr
+import spaces
 from PIL import Image
-from transformers import AutoModelForCausalLM
-from starvector.data.util import process_and_rasterize_svg
 import torch
+from transformers import AutoProcessor, AutoModelForCausalLM
 
-print("Loading model...")
-model_name = "starvector/starvector-8b-im2svg"
-starvector = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    torch_dtype=torch.float16, 
+MODEL_ID = "starvector/starvector-8b-im2svg"
+
+# Load once at startup; ZeroGPU allocates GPU when the decorated function runs.
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
     trust_remote_code=True
-)
-processor = starvector.model.processor
-tokenizer = starvector.model.svg_transformer.tokenizer
+).eval()
+
+def _prep_inputs(image: Image.Image | None, text: str):
+    text = text or ""
+    if image is not None:
+        return processor(images=image, text=text, return_tensors="pt")
+    return processor(text=text, return_tensors="pt")
 
-# Move to GPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
-starvector.to(device)
-starvector.eval()
-print(f"Model loaded on {device}!")
+@spaces.GPU(duration=180)  # request ZeroGPU for up to 180s per call
+def run_starvector(image: Image.Image | None, text: str) -> str:
+    inputs = _prep_inputs(image, text)
+    # Move only tensors to the GPU at call time
+    inputs = {k: v.to("cuda") if hasattr(v, "to") else v for k, v in inputs.items()}
+    with torch.no_grad():
+        out = model.generate(
+            **inputs,
+            max_new_tokens=2048,
+            temperature=0.2,
+            do_sample=False
+        )
+    svg = processor.batch_decode(out, skip_special_tokens=True)[0]
+    return svg
 
-def convert_image_to_svg(image_pil):
-    """Convert uploaded image to SVG"""
-    try:
-        # Process image
-        image = processor(image_pil, return_tensors="pt")['pixel_values'].to(device)
-        if image.shape[0] != 1:
-            image = image.unsqueeze(0)
-        
-        batch = {"image": image}
-        
-        # Generate SVG
-        raw_svg = starvector.generate_im2svg(batch, max_length=4000)[0]
-        svg, raster_image = process_and_rasterize_svg(raw_svg)
-        
-        # Save SVG to file
-        svg_path = "output.svg"
-        with open(svg_path, 'w') as f:
-            f.write(svg)
-        
-        return svg_path, raster_image
-    
-    except Exception as e:
-        return None, f"Error: {str(e)}"
+def preview_svg(svg_code: str) -> str:
+    # Safe inline preview wrapper
+    return f"<div style='height:480px;overflow:auto;border:1px solid #ccc'><pre>{gr.utils.sanitize_html(svg_code)}</pre></div>"
 
-# Create Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 Image to SVG Converter")
-    gr.Markdown("Convert your images to SVG format using StarVector AI")
-    
+with gr.Blocks(title="StarVector: Image/Text → SVG") as demo:
+    gr.Markdown("# StarVector: Image/Text → SVG")
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Upload Image")
-            convert_btn = gr.Button("Convert to SVG", variant="primary")
-        
-        with gr.Column():
-            output_file = gr.File(label="Download SVG")
-            output_preview = gr.Image(label="Preview")
-    
-    convert_btn.click(
-        fn=convert_image_to_svg,
-        inputs=input_image,
-        outputs=[output_file, output_preview]
+        img = gr.Image(type="pil", label="Upload image (optional)")
+        txt = gr.Textbox(label="Text prompt (optional)")
+    btn = gr.Button("Generate SVG")
+    svg_code = gr.Code(label="SVG Output", language="xml")
+    svg_render = gr.HTML(label="Preview")
+    btn.click(fn=run_starvector, inputs=[img, txt], outputs=svg_code).then(
+        fn=preview_svg, inputs=svg_code, outputs=svg_render
     )
-    
-    gr.Markdown("### Example: Upload a PNG/JPG image and get an SVG file!")
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()