Spaces:
Running on Zero
Running on Zero
File size: 2,390 Bytes
5ecb125 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """Turn the demo form into a /v1/systemone-style request for vev-4b and the answers into label dicts."""
import base64
import io
import os
from PIL import Image
from vev.model import CheckpointEngine
MODEL = os.environ.get("VEV_MODEL", "CountingSheep/vev-4b")
N_QUESTIONS = 4
def load() -> CheckpointEngine:
return CheckpointEngine(MODEL, dtype="bf16", device="cuda")
def data_url(img: Image.Image) -> str:
buf = io.BytesIO()
img.convert("RGB").save(buf, format="PNG")
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")
def build(rows: list[tuple[str, str]]) -> dict[str, dict]:
"""rows: (question, options). Two or more options separated by ';' make a choice; otherwise yes/no."""
questions = {}
for i, (q, opts) in enumerate(rows):
q = (q or "").strip()
if not q:
continue
labels = [o.strip() for o in (opts or "").split(";") if o.strip()]
if len(labels) >= 2:
questions[f"q{i}"] = {"type": "choice", "instructions": q, "criteria": {k: None for k in labels}}
else:
questions[f"q{i}"] = {"type": "noul", "instructions": q, "criteria": None}
return questions
def ask(engine: CheckpointEngine, image: Image.Image, context: str, rows: list[tuple[str, str]]):
"""Returns one {label: probability} dict (or None for an empty row) per row, and the engine time in ms."""
questions = build(rows)
if image is None or not questions:
raise ValueError("Upload an image and write at least one question.")
state = {"screen": {"image": {"url": data_url(image)}}}
if (context or "").strip():
state["context"] = context.strip()
res = engine.run(state, questions)
out = []
for i in range(len(rows)):
a = res.answers.get(f"q{i}")
if a is None:
out.append(None)
elif a["type"] == "noul":
out.append({"yes": a["noul"], "no": 1 - a["noul"]})
else:
out.append(a["probabilities"])
return out, res.extensions["timing"]["total_ms"]
EXAMPLE = ["checkout.png", "",
"Does the screen show an error message?", "",
"Which checkout step is the user on?", "shipping; payment; review",
"What should the user do next?", "try another card; wait for the order to ship; nothing, the order went through",
"", ""]
|