Add Docker Space app for StarVector

Co-Authored-By: Oz <oz-agent@warp.dev>

Dockerfile

@@ -0,0 +1,32 @@
+FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel
+
+# System deps
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    build-essential \
+    libcairo2 \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender1 \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Install StarVector from GitHub
+# Pin to main branch by default; you can pin a commit for reproducibility.
+RUN git clone --depth 1 https://github.com/joanrod/star-vector.git /app/star-vector
+
+# Install python dependencies
+RUN pip install --upgrade pip \
+    && pip install -e /app/star-vector
+
+# Copy Space app
+COPY app.py /app/app.py
+
+# Hugging Face Spaces sets PORT; default to 7860 locally.
+ENV PORT=7860
+
+EXPOSE 7860
+
+CMD ["python", "/app/app.py"]

README.md

@@ -1,11 +1,19 @@
 ---
-title: Starvector Gpu
+title: Starvector GPU
 emoji: 👀
 colorFrom: indigo
 colorTo: indigo
 sdk: docker
 pinned: false
-short_description: starvector-gpu
+short_description: StarVector Image→SVG on a CUDA GPU
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This is a Docker Space that runs **StarVector** (image→SVG generation) with **CUDA**.
+
+Setup
+- In Space settings, pick a GPU hardware tier.
+- Add a Space secret `HUGGING_FACE_HUB_TOKEN` with access to the gated model `bigcode/starcoderbase-1b`.
+
+Env vars (optional)
+- `STARVECTOR_MODEL` (default: `starvector/starvector-1b-im2svg`)
+- `STARVECTOR_MAX_LENGTH` (default: `4000`)

app.py

@@ -0,0 +1,85 @@
+import os
+import traceback
+
+import gradio as gr
+import torch
+from PIL import Image
+
+from starvector.model.starvector_arch import StarVectorForCausalLM
+from starvector.data.util import process_and_rasterize_svg
+
+
+def _ensure_hf_token_env() -> None:
+    # The repo docs mention HUGGING_FACE_HUB_TOKEN; huggingface_hub also looks for HF_TOKEN.
+    tok = os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
+    if tok and not os.environ.get("HF_TOKEN"):
+        os.environ["HF_TOKEN"] = tok
+
+
+_ensure_hf_token_env()
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+
+MODEL_NAME = os.environ.get("STARVECTOR_MODEL", "starvector/starvector-1b-im2svg")
+MAX_LENGTH = int(os.environ.get("STARVECTOR_MAX_LENGTH", "4000"))
+
+print(f"Starting StarVector Space on device={DEVICE} dtype={DTYPE} model={MODEL_NAME}", flush=True)
+
+starvector = StarVectorForCausalLM.from_pretrained(MODEL_NAME)
+starvector = starvector.to(device=DEVICE)
+starvector.eval()
+
+
+def im2svg(image: Image.Image, max_length: int = MAX_LENGTH):
+    if image is None:
+        return "", None, "No image provided."
+
+    try:
+        # Preprocess image
+        image_tensor = starvector.process_images([image])[0]
+        image_tensor = image_tensor.to(device=DEVICE)
+        if DEVICE == "cuda":
+            image_tensor = image_tensor.to(dtype=DTYPE)
+
+        batch = {"image": image_tensor}
+
+        # Generate raw svg
+        raw_svg = starvector.generate_im2svg(batch, max_length=max_length)[0]
+
+        # Clean + rasterize preview
+        svg, raster = process_and_rasterize_svg(raw_svg)
+
+        # raster may be a PIL image
+        return svg, raster, ""
+
+    except Exception:
+        return "", None, traceback.format_exc()
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "# StarVector (GPU)\n"
+        "Upload an icon/logo/diagram-like image and generate SVG code.\n\n"
+        "Notes:\n"
+        "- This Space requires a GPU and a Hugging Face token with access to the gated `bigcode/starcoderbase-1b` model.\n"
+        "- Set Space secret `HUGGING_FACE_HUB_TOKEN`."
+    )
+
+    with gr.Row():
+        inp = gr.Image(type="pil", label="Input image")
+        preview = gr.Image(type="pil", label="Rasterized preview")
+
+    max_len = gr.Slider(256, 8000, value=MAX_LENGTH, step=128, label="max_length")
+
+    out_svg = gr.Code(language="xml", label="SVG")
+    err = gr.Textbox(label="Error", visible=True)
+
+    btn = gr.Button("Generate SVG")
+    btn.click(im2svg, inputs=[inp, max_len], outputs=[out_svg, preview, err])
+
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", "7860"))
+    demo.queue(concurrency_count=1)
+    demo.launch(server_name="0.0.0.0", server_port=port)