| |
| """StarVector-1B HTTP server for HF Inference Endpoints (and local docker). |
| |
| One contract, wherever it runs (PROPOSAL.md §5.1 L2 / §11): |
| POST /generate {"image_b64": "<png b64>", "max_tokens": 4096, "num_beams": 2} |
| -> {"svg": "<svg ...>", "elapsed_s": 12.3} |
| GET /health -> {"ok": true, "model_loaded": bool, "device": "..."} |
| |
| License notes: see LICENSE/NOTICE in this directory. Model weights are |
| self-contained in the starvector safetensors — we never touch the gated |
| bigcode/starcoderbase-1b repo (same trick as the estate bridge). |
| """ |
| import base64 |
| import importlib.util |
| import io |
| import os |
| import sys |
| import textwrap |
| import time |
|
|
| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel |
|
|
| MODEL_NAME = os.environ.get("STARVECTOR_MODEL", "starvector/starvector-1b-im2svg") |
| MAX_TOKENS_DEFAULT = int(os.environ.get("STARVECTOR_MAX_TOKENS", "4096")) |
| PORT = int(os.environ.get("PORT", "8080")) |
|
|
| app = FastAPI(title="IconEdge StarVector bridge", version="0.1.0") |
| _state = {"model": None, "device": "cpu", "loaded": False} |
|
|
|
|
| def _patch_starvector() -> None: |
| """Same patch as IconTreasury's bridge: load tokenizer+LLM weights from the |
| starvector repo itself instead of the gated bigcode/starcoderbase-1b repo.""" |
| spec = importlib.util.find_spec("starvector") |
| if spec is None or not spec.submodule_search_locations: |
| raise RuntimeError("starvector package not installed") |
| pkg_dir = list(spec.submodule_search_locations)[0] |
| patch_path = os.path.join(pkg_dir, "model", "llm", "starcoder_patched.py") |
| if os.path.exists(patch_path): |
| return |
| os.makedirs(os.path.dirname(patch_path), exist_ok=True) |
| with open(patch_path, "w") as f: |
| f.write(textwrap.dedent('''\ |
| """Patched StarCoderModel — loads tokenizer+weights from starvector repo |
| instead of the gated bigcode/starcoderbase-1b repo.""" |
| import torch.nn as nn |
| from transformers import AutoTokenizer |
| from starvector.model.gpt_bigcode.configuration_gpt_bigcode import GPTBigCodeConfig |
| from starvector.model.gpt_bigcode.modeling_gpt_bigcode import GPTBigCodeForCausalLM |
| |
| class StarCoderModelPatched(nn.Module): |
| def __init__(self, config, **kwargs): |
| super().__init__() |
| repo = config._name_or_path if hasattr(config, "_name_or_path") else "starvector/starvector-1b-im2svg" |
| self.init_tokenizer(repo) |
| self.max_length = config.max_length |
| dropout = config.dropout if hasattr(config, "dropout") else 0.1 |
| model_config = GPTBigCodeConfig( |
| vocab_size=config.vocab_size, |
| n_positions=config.max_position_embeddings, |
| n_embd=config.hidden_size, |
| n_layer=config.num_hidden_layers, |
| n_head=config.num_attention_heads, |
| n_inner=config.hidden_size * 4, |
| activation_function="gelu_pytorch_tanh", |
| resid_pdrop=dropout, embd_pdrop=dropout, attn_pdrop=dropout, |
| layer_norm_epsilon=1e-5, scale_attn_weights=True, |
| use_cache=config.use_cache, |
| bos_token_id=self.tokenizer.bos_token_id, |
| eos_token_id=self.tokenizer.eos_token_id, |
| pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id else self.tokenizer.eos_token_id, |
| multi_query=config.multi_query if hasattr(config, "multi_query") else True, |
| ) |
| model = GPTBigCodeForCausalLM(model_config) |
| model.resize_token_embeddings(len(self.tokenizer)) |
| self.transformer = model |
| self.prompt = "<svg" |
| |
| def init_tokenizer(self, model_name): |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False) |
| if self.tokenizer.eos_token_id is None: |
| self.tokenizer.add_special_tokens({"eos_token": "[EOS]"}) |
| if self.tokenizer.pad_token_id is None: |
| self.tokenizer.add_special_tokens({"pad_token": "[PAD]"}) |
| self.svg_start_token = "<svg-start>" |
| self.image_start_token = "<image-start>" |
| self.text_start_token = "<caption-start>" |
| self.tokenizer.add_tokens([self.svg_start_token, self.image_start_token, self.text_start_token]) |
| self.svg_start_token_id = self.tokenizer.encode(self.svg_start_token)[0] |
| ''')) |
|
|
|
|
| def _load_model() -> None: |
| import torch |
|
|
| _patch_starvector() |
| import starvector.model.models.starvector_v1 as v1 |
| from starvector.model.llm.starcoder_patched import StarCoderModelPatched |
| v1.StarVectorStarCoder._get_svg_transformer = ( |
| lambda self, config, **kwargs: StarCoderModelPatched(config, **kwargs) |
| ) |
| from starvector.model.starvector_arch import StarVectorForCausalLM |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| dtype = torch.float16 |
| print(f"Loading {MODEL_NAME} on {device}...", flush=True) |
| t0 = time.time() |
| model = StarVectorForCausalLM.from_pretrained( |
| MODEL_NAME, torch_dtype=dtype, trust_remote_code=True, |
| ) |
| model.config.use_flash_attn = False |
| model.eval() |
| model = model.to(device) |
| model.model_precision = dtype |
| _state.update(model=model, device=device, loaded=True) |
| print(f"Model loaded in {time.time() - t0:.1f}s", flush=True) |
|
|
|
|
| class GenerateRequest(BaseModel): |
| image_b64: str |
| max_tokens: int = MAX_TOKENS_DEFAULT |
| num_beams: int = 2 |
|
|
|
|
| @app.get("/health") |
| def health(): |
| return {"ok": True, "model_loaded": _state["loaded"], "device": _state["device"]} |
|
|
|
|
| @app.post("/generate") |
| def generate(req: GenerateRequest): |
| if not _state["loaded"]: |
| raise HTTPException(503, "model still loading — retry shortly") |
| import torch |
| from PIL import Image |
|
|
| try: |
| png_bytes = base64.b64decode(req.image_b64) |
| except Exception as e: |
| raise HTTPException(400, f"image_b64 is not valid base64: {e}") |
|
|
| image = Image.open(io.BytesIO(png_bytes)).convert("RGB") |
| model = _state["model"] |
| image_tensor = model.process_images([image])[0].to(_state["device"], dtype=torch.float16) |
|
|
| t0 = time.time() |
| with torch.no_grad(): |
| raw_svgs = model.generate_im2svg( |
| {"image": image_tensor}, max_length=req.max_tokens, |
| num_beams=req.num_beams, temperature=1.0, |
| ) |
| return {"svg": raw_svgs[0], "elapsed_s": round(time.time() - t0, 2)} |
|
|
|
|
| @app.on_event("startup") |
| def startup(): |
| if os.environ.get("LOAD_ON_STARTUP", "1") == "1": |
| _load_model() |
|
|
|
|
| if __name__ == "__main__": |
| import uvicorn |
| uvicorn.run(app, host="0.0.0.0", port=PORT) |
|
|
