File size: 4,442 Bytes
15b8951 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """FastAPI server cho VLM Web Console."""
import os
import base64
import asyncio
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from ..common.vlm_engine import parse_boxes, draw_boxes, encode_frame_jpeg
WEBUI = Path(__file__).resolve().parent.parent / "webui" / "index.html"
def create_app(source, engine):
app = FastAPI()
app.state.busy = False
@app.get("/", response_class=HTMLResponse)
async def index():
return WEBUI.read_text(encoding="utf-8")
@app.get("/status")
async def status():
return JSONResponse({
"connected": bool(source.is_connected),
"mock": bool(source.is_mock),
})
@app.get("/video_feed")
async def video_feed():
async def gen():
while True:
frame = source.get_latest_frame()
if frame is not None:
jpg = encode_frame_jpeg(frame)
yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n"
+ jpg + b"\r\n")
await asyncio.sleep(0.05) # ~20 fps
return StreamingResponse(
gen(), media_type="multipart/x-mixed-replace; boundary=frame"
)
@app.websocket("/ws")
async def ws(websocket: WebSocket):
await websocket.accept()
try:
while True:
data = await websocket.receive_json()
prompt = (data or {}).get("prompt", "").strip()
if not prompt:
continue
if app.state.busy:
await websocket.send_json({"type": "busy"})
continue
app.state.busy = True
try:
await _handle_prompt(websocket, source, engine, prompt)
except Exception as e:
await websocket.send_json({"type": "error", "message": str(e)})
finally:
app.state.busy = False
except WebSocketDisconnect:
return
return app
async def _handle_prompt(websocket, source, engine, prompt):
frame = source.get_latest_frame()
if frame is None:
await websocket.send_json(
{"type": "error", "message": "Chưa có ảnh từ robot."}
)
return
snapshot = frame.copy()
loop = asyncio.get_event_loop()
full = []
# Bơm generator (blocking) qua executor, đẩy token ra queue của event loop.
def produce(q):
try:
for chunk in engine.stream_infer(snapshot, prompt):
asyncio.run_coroutine_threadsafe(q.put(("token", chunk)), loop)
asyncio.run_coroutine_threadsafe(q.put(("end", None)), loop)
except Exception as e:
asyncio.run_coroutine_threadsafe(q.put(("err", str(e))), loop)
q: asyncio.Queue = asyncio.Queue()
loop.run_in_executor(None, produce, q)
while True:
kind, payload = await q.get()
if kind == "token":
full.append(payload)
await websocket.send_json({"type": "token", "text": payload})
elif kind == "err":
await websocket.send_json({"type": "error", "message": payload})
return
else: # end
break
text = "".join(full)
h, w = snapshot.shape[:2]
boxes = parse_boxes(text, w, h)
if boxes:
drawn = draw_boxes(snapshot, boxes, prompt)
b64 = base64.b64encode(encode_frame_jpeg(drawn)).decode("ascii")
await websocket.send_json({"type": "image", "data": b64})
await websocket.send_json({"type": "done"})
def main():
import uvicorn
from ..common.frame_source import FrameSource
from ..common.vlm_engine import VLMEngine
source = FrameSource(os.getenv("ROBOT_IP"))
engine = VLMEngine()
if os.getenv("VLM_SKIP_MODEL") == "1":
print("[server] VLM_SKIP_MODEL=1 -> bỏ qua load model (chỉ xem GUI). "
"Gửi lệnh sẽ báo lỗi cho tới khi bỏ cờ này.")
else:
print("[server] Loading VLM model... (lần đầu sẽ tải ~vài GB)")
engine.load()
app = create_app(source, engine)
@app.on_event("startup")
async def _startup():
await source.start()
uvicorn.run(app, host="127.0.0.1", port=8000)
if __name__ == "__main__":
main()
|