Spaces:
Running on Zero
Running on Zero
| """ | |
| The Hugging Face Space entry point (Gradio SDK on ZeroGPU hardware, free for a personal account): the REST API of | |
| serving/app/main.py unchanged, plus a demo page at /ui. scripts/deploy_hf.py uploads this file as the Space's app.py. | |
| ZeroGPU lends a GPU per call. A request needs at most one (serving/app/vision.py): the pill model and every reading | |
| (Florence-2-large: imprints, labels, the input check) run together in one @spaces.GPU call; requests that only embed | |
| pills stay on the CPU, so the free allowance (a daily number of runs) lasts. When no GPU is granted (allowance used | |
| up, queue too long) the request runs on the CPU with Florence-2-base within a time budget, and stays on the CPU for | |
| COOLDOWN seconds, so answers keep coming. | |
| Off ZeroGPU it is the plain CPU API with the demo page: python serving/space_app.py (http://localhost:7860/ui) | |
| """ | |
| import copy | |
| import os | |
| import sys | |
| from pathlib import Path | |
| HERE = Path(__file__).resolve().parent | |
| BASE = HERE if (HERE / "serving").is_dir() else HERE.parent # the Space root (app.py), or the epills folder | |
| sys.path[:0] = [str(BASE / "src"), str(BASE)] | |
| try: | |
| import spaces # before torch: on ZeroGPU it patches torch | |
| except ImportError: | |
| spaces = None | |
| import gradio as gr | |
| import torch | |
| import uvicorn | |
| from fastapi import HTTPException | |
| from fastapi.responses import RedirectResponse | |
| import threading | |
| from epills.core.imprint.ocr import GPU_MODEL | |
| from epills.core.inference.pill_finder import find_pills | |
| from epills.core.inference.runtime import cpu_cores | |
| from serving.app.fallback import GpuOrCpu | |
| from serving.app.main import IDENTIFY_NOTE, LABEL_NOTE, app, asset, identify_pill, imprint_adapter_path, vision | |
| ZERO_GPU = spaces is not None and os.environ.get("SPACES_ZERO_GPU", "").lower() in ("1", "t", "true") | |
| COOLDOWN = 600 # seconds on the CPU after ZeroGPU refused a call | |
| if os.environ.get("EPILLS_MODEL_REPO"): | |
| token = os.environ.get("HF_TOKEN") | |
| try: | |
| from huggingface_hub import whoami | |
| who = whoami(token)["name"] if token else "no token" | |
| except Exception as error: | |
| who = f"token rejected ({type(error).__name__})" | |
| print(f"assets from {os.environ['EPILLS_MODEL_REPO']}: HF_TOKEN {'set' if token else 'NOT set'}, {who}", flush=True) | |
| for name in ("safety_db", "embedder", "catalogue", "reader"): # now: ZeroGPU needs the models before start-up | |
| try: | |
| asset(name) | |
| except HTTPException as error: | |
| print(f"{name}: {error.detail}", flush=True) | |
| embedder, reader = app.state.embedder, app.state.reader | |
| print(f"CPU: {cpu_cores()} usable cores (os.cpu_count {os.cpu_count()}); torch threads {torch.get_num_threads()}", | |
| flush=True) | |
| pipeline = vision() if embedder is not None else None | |
| if pipeline is not None and reader is not None: | |
| pipeline.warm_up() # times a CPU reading, for the CPU time budget | |
| print(f"CPU reading: {pipeline.seconds_per_image:.2f} s per image; budget {pipeline.budget:.0f} s", flush=True) | |
| threading.Thread(target=lambda: asset("label_reader").vocab, daemon=True).start() # RxNorm names, ~3 s | |
| if ZERO_GPU and pipeline is not None: | |
| from transformers import Florence2ForConditionalGeneration | |
| gpu_model = copy.deepcopy(embedder.model).to("cuda") # placed now; ZeroGPU moves them to a real GPU per call | |
| gpu_ocr = Florence2ForConditionalGeneration.from_pretrained(GPU_MODEL, torch_dtype=torch.float16).eval() | |
| adapter = imprint_adapter_path() # trained on pill imprints; switched on for imprint reads only | |
| if adapter is not None: # added on the CPU, then moved with the model | |
| from epills.core.imprint.lora import load_lora, set_lora | |
| app.state.gpu_imprint_adapter = load_lora(gpu_ocr, adapter).get("name", "imprint adapter") | |
| print(f"imprint adapter: {app.state.gpu_imprint_adapter}", flush=True) | |
| gpu_ocr = gpu_ocr.to("cuda") | |
| def gpu_compute(payload): | |
| """One request's work in one GPU call: the pill model, then every reading (large model).""" | |
| out = {"embed": None, "gen": [], "device": "zerogpu"} | |
| with torch.no_grad(): | |
| if payload["pixels"] is not None: | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| out["embed"] = gpu_model(payload["pixels"].cuda())["embedding"].float().cpu().numpy() | |
| for job in payload["jobs"]: | |
| if adapter is not None: | |
| set_lora(gpu_ocr, job["task"] == "read") | |
| out["gen"].append(gpu_ocr.generate(input_ids=job["ids"].cuda(), pixel_values=job["px"].cuda().half(), | |
| max_new_tokens=job["tokens"], num_beams=3).cpu()) | |
| return out | |
| pipeline.compute = GpuOrCpu(gpu_compute, pipeline.cpu_compute, COOLDOWN) | |
| # ---- demo page ------------------------------------------------------------------------------------------------------- | |
| def candidates(vectors, imprint=None, top_k=3, reads=None): | |
| try: | |
| return identify_pill(vectors, imprint=imprint or None, top_k=top_k, reads=reads)["candidates"] | |
| except HTTPException as error: | |
| raise gr.Error(str(error.detail)) | |
| def label(c): | |
| read = ", imprint read" if c.get("imprint_read") else "" | |
| return f"{c['name']} (imprint {c['imprint'] or 'none'}{read}), similarity {c['similarity']:.2f}" | |
| def ui_identify(photo, imprint): | |
| if photo is None or embedder is None: | |
| raise gr.Error("Add a photo of one pill." if embedder else "The model isn't loaded; see /health.") | |
| cat, out = app.state.catalogue, [] | |
| crops = embedder.frame(photo) | |
| r = vision().run(crops, None if imprint else crops) | |
| for c in candidates(r["embeddings"], imprint, top_k=6, reads=(r["reads"] or [None])[0]): | |
| images = cat.product(c["pill_id"])["images"] | |
| if images: | |
| out.append((str(cat.image_file(c["pill_id"], images[0]["n"])), label(c))) | |
| return out | |
| def ui_scan(photo): | |
| if photo is None or embedder is None: | |
| raise gr.Error("Add a photo of the pills." if embedder else "The model isn't loaded; see /health.") | |
| boxes = find_pills(photo) | |
| if not boxes: | |
| raise gr.Error("No pill found. Spread the pills apart on a plain background and photograph them again.") | |
| rows, marks = [], [] | |
| crops = embedder.frame(photo, boxes) | |
| r = vision().run(crops, crops) | |
| for k, (box, vector, read) in enumerate(zip(boxes, r["embeddings"], r["reads"] or [None] * len(crops)), 1): | |
| rows.append([k] + [label(c) for c in candidates(vector[None], reads=read)]) | |
| marks.append(([round(v) for v in box], f"pill {k}")) | |
| return (photo, marks), rows | |
| def ui_label(photo): | |
| if photo is None or reader is None: | |
| raise gr.Error("Add a photo of a box or label." if reader else "Text reading isn't loaded; see /health.") | |
| try: | |
| found = asset("label_reader").read(vision().run(text=photo)["lines"] or []) | |
| except Exception as error: | |
| raise gr.Error(f"Couldn't look the names up: {error}") | |
| if not found: | |
| raise gr.Error("No medicine name found. Photograph the label closer, straight on and in good light.") | |
| return [[c["name"], ", ".join(i["name"] for i in c["ingredients"]), c["rxcui"], c["rxnorm_name"] or ""] | |
| for c in found] | |
| with gr.Blocks(title="PharmaScan") as demo: | |
| gr.Markdown("# PharmaScan\nA demo of the PharmaScan API ([REST API docs](/docs)). " + IDENTIFY_NOTE) | |
| with gr.Tab("Identify a pill"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| pill = gr.Image(type="pil", label="One pill in the middle of the photo, filling about 70% of it") | |
| imprint = gr.Textbox(label="Imprint (optional)", placeholder="e.g. M 547") | |
| identify_button = gr.Button("Identify", variant="primary") | |
| gallery = gr.Gallery(label="Most alike pills", columns=3) | |
| identify_button.click(ui_identify, [pill, imprint], gallery) | |
| with gr.Tab("Scan a photo"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| full = gr.Image(type="pil", label="Several pills, apart, on a plain background") | |
| scan_button = gr.Button("Scan", variant="primary") | |
| found = gr.AnnotatedImage(label="Pills found") | |
| table = gr.Dataframe(headers=["pill", "most alike", "2nd", "3rd"], label="Catalogue matches") | |
| scan_button.click(ui_scan, full, [found, table]) | |
| with gr.Tab("Read a box or label"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| box_photo = gr.Image(type="pil", label="The medicine's box or pharmacy label, close up") | |
| label_button = gr.Button("Read", variant="primary") | |
| medicines = gr.Dataframe(headers=["medicine", "ingredients", "RxNorm code", "RxNorm name"], | |
| label=LABEL_NOTE) | |
| label_button.click(ui_label, box_photo, medicines) | |
| app = gr.mount_gradio_app(app, demo, path="/ui", ssr_mode=False) # SSR's Node server would take port 7860 | |
| def root(): | |
| return RedirectResponse("/ui") | |
| if __name__ == "__main__": | |
| if ZERO_GPU and pipeline is not None: | |
| from spaces.zero import startup # what demo.launch() would run: hands the placed weights to ZeroGPU | |
| startup() | |
| uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860))) | |