muaalem-proxy / app.py
hatimqman's picture
add GET /model (serve 126MB streaming model for browser on-device; CF Pages 25MiB limit) + 5/h/IP limit
b2afe94 verified
Raw
History Blame Contribute Delete
12.6 kB
# -*- coding: utf-8 -*-
"""ูˆุณูŠุท ยซุงู„ู…ุฌูˆู‘ุฏ/ุจูŠุงู†ยป (HF Space) โ€” ุซู„ุงุซ ู…ู‡ุงู…ู‘:
1) ุจุฑูˆูƒุณูŠ ุขู…ู† ู„ุฎุงุฏู…ูŽูŠ HF ุงู„ู…ุญู…ูŠู‘ูŠู† (/stt ู†ุตู‘ูŠู‘ุŒ /analyze ููˆู†ูŠู…ูŠู‘) โ€” ุงู„ุชูˆูƒู† ุณุฑู‘ Space.
2) ุญู…ุงูŠุฉ ุฅุณุงุกุฉ ุงู„ุงุณุชุฎุฏุงู…: ุญุฏู‘ ู…ุนุฏู‘ู„ ู„ูƒู„ู‘ IP (ู†ุงูุฐุฉ ู…ู†ุฒู„ู‚ุฉ) + ุณู‚ู ุฌู„ุณุงุช ุจุซู‘ ู…ุชุฒุงู…ู†ุฉ.
3) /stream: ุจุซู‘ ู„ุญุธูŠู‘ (WebSocket) ุจู†ู…ูˆุฐุฌ Muno459 streaming (cache-aware CTC) ูŠุนู…ู„ ู‡ู†ุง
ุนู„ู‰ CPU โ€” ุฅุทุงุฑุงุช int16@16k โ†’ ู†ุตู‘ ุชุฑุงูƒู…ูŠู‘ + ู†ุงูุฐุฉ recent (~4.5ุซ) ู„ูƒู„ู‘ ู‚ุทุนุฉ 112-mel.
ุงู„ุซูˆุงุจุช ุงู„ุตูˆุชูŠู‘ุฉ (ู†ุงูุฐุฉ/mel-fb/CMVN) ู…ู† mel_const.json โ€” ู†ูุณ ู…ู„ูู‘ ุงู„ุนู…ูŠู„ุŒ ุชุทุงุจู‚ูŒ ุจุชู‘ูŠู‘.
"""
import os, io, json, time, base64, asyncio, threading
from collections import deque, defaultdict
import numpy as np
import requests
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
EP = os.environ.get("MUAALEM_ENDPOINT_URL", "").rstrip("/")
TOK = os.environ.get("MUAALEM_TOKEN", "")
STT_EP = os.environ.get("STT_ENDPOINT_URL", "").rstrip("/")
STT_TOK = os.environ.get("STT_TOKEN", "") or TOK
app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ ุญู…ุงูŠุฉ ุงู„ู…ุนุฏู‘ู„ (ู„ูƒู„ู‘ IP) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
_BUCKETS: dict = defaultdict(deque)
_WS_LIVE: dict = defaultdict(int)
_LOCK = threading.Lock()
def _client_ip(scope_headers, client):
for k, v in scope_headers:
if k == b"x-forwarded-for":
return v.decode().split(",")[0].strip()
return client.host if client else "?"
def _allow(ip: str, route: str, limit: int, per: int = 60) -> bool:
with _LOCK:
q = _BUCKETS[(ip, route)]
now = time.time()
while q and now - q[0] > per:
q.popleft()
if len(q) >= limit:
return False
q.append(now)
return True
def _ip_http(req: Request) -> str:
return (req.headers.get("x-forwarded-for", "").split(",")[0].strip()
or (req.client.host if req.client else "?"))
# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ ู…ุญุฑู‘ูƒ ุงู„ุจุซู‘ (ุนู„ู‰ ุงู„ูˆุณูŠุท) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
_HERE = os.path.dirname(os.path.abspath(__file__))
_SESS = None # onnxruntime session (ูŠูุญู…ูŽู‘ู„ ุจุฎูŠุท ุฎู„ููŠู‘ ุนู†ุฏ ุงู„ุฅู‚ู„ุงุน)
_MC = None # ุซูˆุงุจุช mel
_VOCAB = None
_WINDOW = _MELFB = None
_CMVN = {}
BLANK, NL, DM, LC, TC, CH = 1024, 17, 512, 70, 8, 112
WIN = HOP = NFFT = NR = NMEL = 0
CS = ADV = 0
def _b64f32(s):
return np.frombuffer(base64.b64decode(s), dtype="<f4").copy()
def _load_model():
global _SESS, _MC, _VOCAB, _WINDOW, _MELFB, _CMVN, WIN, HOP, NFFT, NR, NMEL, CS, ADV
try:
import onnxruntime as ort
_MC = json.load(io.open(os.path.join(_HERE, "mel_const.json"), encoding="utf-8"))
_VOCAB = json.load(io.open(os.path.join(_HERE, "vocab.json"), encoding="utf-8"))
WIN, HOP, NFFT, NR, NMEL = _MC["win_n"], _MC["hop_n"], _MC["n_fft"], _MC["n_rfft"], _MC["n_mels"]
CS = WIN + (CH - 1) * HOP
ADV = CH * HOP
_WINDOW = _b64f32(_MC["window_b64"])
_MELFB = _b64f32(_MC["mel_fb_b64"]).reshape(NMEL, NR)
for kind in ("clean", "tlog"):
_CMVN[kind] = (_b64f32(_MC[f"{kind}_mean_b64"])[:, None], _b64f32(_MC[f"{kind}_std_b64"])[:, None])
so = ort.SessionOptions()
so.intra_op_num_threads = max(1, (os.cpu_count() or 2) - 1)
_SESS = ort.InferenceSession(os.path.join(_HERE, "model.q8.onnx"), so, providers=["CPUExecutionProvider"])
print("stream model ready", flush=True)
except Exception as e:
print("stream model load FAILED:", e, flush=True)
threading.Thread(target=_load_model, daemon=True).start()
def _decode(ids):
return "".join((_VOCAB[i] if 0 <= i < len(_VOCAB) else "") for i in ids).replace("โ–", " ").strip()
class _Stream:
def __init__(self, kind):
self.kind = kind if kind in ("clean", "tlog") else "tlog"
self.buf = np.zeros(0, np.float32)
self.prev_raw = 0.0
self.clc = np.zeros((1, NL, LC, DM), np.float32)
self.clt = np.zeros((1, NL, DM, TC), np.float32)
self.cll = np.zeros((1,), np.int64)
self.ids, self.born, self.prev, self.n = [], [], -1, 0
def feed(self, raw: bytes):
self.buf = np.concatenate([self.buf, np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0])
def ready(self):
return len(self.buf) >= CS
def _mel(self, seg):
x = np.empty_like(seg)
x[0] = seg[0] - 0.97 * self.prev_raw
x[1:] = seg[1:] - 0.97 * seg[:-1]
nf = 1 + (len(x) - WIN) // HOP
from numpy.lib.stride_tricks import as_strided
fr = as_strided(x, shape=(nf, WIN), strides=(x.strides[0] * HOP, x.strides[0])).copy() * _WINDOW
sp = np.fft.rfft(fr, n=NFFT, axis=1)
pw = np.real(sp * np.conj(sp)).astype(np.float32)
mel = np.log(pw @ _MELFB.T + 2 ** -24).T
gm, gs = _CMVN[self.kind]
return ((mel - gm) / (gs + 1e-5)).astype(np.float32)
def step(self):
seg = self.buf[:CS]
mel = self._mel(seg)[None]
self.prev_raw = float(self.buf[ADV - 1])
self.buf = self.buf[ADV:]
o = _SESS.run(None, {"audio_signal": mel, "length": np.array([mel.shape[2]], np.int64),
"cache_last_channel": self.clc, "cache_last_time": self.clt,
"cache_last_channel_len": self.cll})
logp, self.clc, self.clt, self.cll = o[0], o[2], o[3], o[4]
for t in logp[0].argmax(-1):
t = int(t)
if t != self.prev and t != BLANK:
self.ids.append(t)
self.born.append(self.n)
self.prev = t
self.n += 1
cut = self.n - 4
recent = [i for i, b in zip(self.ids, self.born) if b >= cut]
return {"t": self.n, "text": _decode(self.ids), "recent": _decode(recent)}
def flush(self):
if len(self.buf) > HOP:
self.buf = np.concatenate([self.buf, np.zeros(CS - len(self.buf), np.float32)])
return self.step()
return None
@app.websocket("/stream")
async def stream(ws: WebSocket):
ip = _client_ip(ws.scope.get("headers", []), ws.client)
if not _allow(ip, "ws-dial", 20):
await ws.close(code=1013)
return
with _LOCK:
if _WS_LIVE[ip] >= 3:
await ws.close(code=1013)
return
_WS_LIVE[ip] += 1
await ws.accept()
try:
# ุงู„ู†ู…ูˆุฐุฌ ู‚ุฏ ูŠูƒูˆู† ู…ุง ุฒุงู„ ูŠูุญู…ูŽู‘ู„ (ุฅู‚ู„ุงุน ุจุงุฑุฏ) โ€” ุงู†ุชุธุฑ ุญุชู‰ 25ุซ
waited = 0.0
while _SESS is None and waited < 25:
await asyncio.sleep(0.5)
waited += 0.5
if _SESS is None:
await ws.close(code=1011)
return
st = _Stream(ws.query_params.get("kind", "tlog"))
while True:
msg = await ws.receive()
if msg.get("type") == "websocket.disconnect":
break
if msg.get("bytes") is not None:
st.feed(msg["bytes"])
# ุณู‚ู ู…ุฎุฒูˆู†: 120ุซ ุตูˆุช ุบูŠุฑ ู…ุนุงู„ูŽุฌ = ุฅุณุงุกุฉ/ุชุณุฑูŠุจ โ†’ ุฃุบู„ู‚
if len(st.buf) > 16000 * 120:
await ws.close(code=1009)
break
while st.ready():
r = await asyncio.to_thread(st.step)
await ws.send_text(json.dumps(r, ensure_ascii=False))
elif msg.get("text") == "end":
r = await asyncio.to_thread(st.flush)
if r:
await ws.send_text(json.dumps(r, ensure_ascii=False))
await ws.close()
break
except WebSocketDisconnect:
pass
finally:
with _LOCK:
_WS_LIVE[ip] = max(0, _WS_LIVE[ip] - 1)
# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ ุงู„ุตุญู‘ุฉ + ุงู„ุจุฑูˆูƒุณูŠ (ูƒู…ุง ูƒุงู† + ุญุฏูˆุฏ ู…ุนุฏู‘ู„) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
@app.get("/")
def health():
return {"ok": True, "muaalem": bool(EP and TOK), "stt": bool(STT_EP and STT_TOK),
"stream": _SESS is not None}
@app.get("/model")
def model_file(req: Request):
"""ู†ู…ูˆุฐุฌ ุงู„ุจุซู‘ (126MB) ู„ู…ูŠุฒุฉ ยซุนู„ู‰ ุงู„ุฌู‡ุงุฒยป ููŠ ุงู„ู…ุชุตูู‘ุญ โ€” Cloudflare Pages ูŠุฑูุถ ู…ู„ูู‘ุงุช >25MiB
ูู†ู‚ุฏู‘ู…ู‡ ู…ู† ู‡ู†ุง (ุงู„ู†ุทุงู‚ ู…ุณู…ูˆุญ ููŠ CSP ุงู„ุชุทุจูŠู‚ูŠู† ุฃุตู„ู‹ุง). ูƒุงุด ุณู†ุฉ + ุญุฏู‘ 5 ุชู†ุฒูŠู„ุงุช/ุณุงุนุฉ/IP."""
if not _allow(_ip_http(req), "model", 5, per=3600):
return JSONResponse({"error": "rate_limited"}, status_code=429)
p = os.path.join(_HERE, "model.q8.onnx")
if not os.path.exists(p):
return JSONResponse({"error": "no_model"}, status_code=404)
return FileResponse(p, media_type="application/octet-stream",
headers={"Cache-Control": "public, max-age=31536000, immutable"})
@app.post("/stt")
async def stt(req: Request):
if not _allow(_ip_http(req), "stt", 120):
return JSONResponse({"error": "rate_limited"}, status_code=429)
try:
body = await req.json()
except Exception:
return JSONResponse({"error": "invalid_json"}, status_code=400)
pcm = body.get("pcm")
if not pcm or not isinstance(pcm, str):
return JSONResponse({"error": "missing_pcm"}, status_code=400)
if not (STT_EP and STT_TOK):
return JSONResponse({"error": "stt_not_configured"}, status_code=503)
ref = body.get("ref")
payload = {"inputs": pcm, "ref": ref} if ref else {"inputs": pcm}
headers = {"Authorization": f"Bearer {STT_TOK}", "Content-Type": "application/json"}
for _ in range(30):
try:
r = requests.post(STT_EP, json=payload, headers=headers, timeout=180)
except Exception as e:
return JSONResponse({"error": "stt_error", "detail": str(e)[:200]}, status_code=502)
if r.status_code in (502, 503):
time.sleep(6)
continue
try:
data = r.json()
except Exception:
return JSONResponse({"text": "", "words": None}, status_code=200)
obj = data[0] if isinstance(data, list) and data else (data if isinstance(data, dict) else {})
return JSONResponse({"text": obj.get("text", ""), "words": obj.get("words"),
"acoustic": obj.get("acoustic")}, status_code=200)
return JSONResponse({"text": "", "words": None}, status_code=200)
@app.post("/analyze")
async def analyze(req: Request):
if not _allow(_ip_http(req), "analyze", 40):
return JSONResponse({"error": "rate_limited"}, status_code=429)
try:
body = await req.json()
except Exception:
return JSONResponse({"error": "invalid_json"}, status_code=400)
pcm = body.get("pcm")
if not pcm or not isinstance(pcm, str):
return JSONResponse({"error": "missing_pcm"}, status_code=400)
if len(pcm) > 30 * 1024 * 1024:
return JSONResponse({"error": "too_large"}, status_code=413)
inputs = {"pcm": pcm, "surah": body.get("surah")}
if body.get("ayahs"):
inputs["ayahs"] = body["ayahs"]
else:
inputs["ayah"] = body.get("ayah")
if not (EP and TOK):
return JSONResponse({"error": "not_configured"}, status_code=503)
headers = {"Authorization": f"Bearer {TOK}", "Content-Type": "application/json"}
for _ in range(30): # ูŠุนุงู„ุฌ ุงู„ุจุฏุก ุงู„ุจุงุฑุฏ ู„ู„ุฎุงุฏู… (scale-to-zero โ†’ 502/503)
try:
r = requests.post(EP, json={"inputs": inputs}, headers=headers, timeout=180)
except Exception as e:
return JSONResponse({"error": "proxy_error", "detail": str(e)[:200]}, status_code=502)
if r.status_code in (502, 503):
time.sleep(6)
continue
try:
data = r.json()
except Exception:
data = {"error": "bad_upstream", "status": r.status_code}
return JSONResponse(data, status_code=200)
return JSONResponse({"error": "timeout"}, status_code=200)