deepcw_model.py
Rama main · réplica de lectura del repositorio original.
"""
deepcw_model.py — server-side management of the DeepCW model
- Downloads the ONNX model from GitHub (e04/web-deep-cw-decoder)
- Serves it locally via /api/deepcw/model
- Checks for updates (commit SHA from the GitHub API)
- Caches it in a local file, deepcw/model_balanced.onnx
"""
import asyncio
import hashlib
import json
import os
import pathlib
import time
import aiohttp
# Data directory (APPDATA) — the model survives product updates.
try:
from config import DATA as _DATA
MODEL_DIR = pathlib.Path(_DATA) / "deepcw"
except Exception:
MODEL_DIR = pathlib.Path("deepcw")
MODEL_FILE = MODEL_DIR / "model_balanced.onnx"
META_FILE = MODEL_DIR / "meta.json"
GITHUB_API = "https://api.github.com/repos/e04/web-deep-cw-decoder/commits?path=public/model_balanced.onnx&per_page=1"
MODEL_URL = "https://cdn.jsdelivr.net/gh/e04/web-deep-cw-decoder@main/public/model_balanced.onnx"
CHECK_INTERVAL = 6 * 3600 # check every 6h
class DeepCWModelManager:
def __init__(self):
self.model_bytes: bytes | None = None
self.meta: dict = {}
self._lock = asyncio.Lock()
def _load_meta(self) -> dict:
try:
return json.loads(META_FILE.read_text())
except Exception:
return {}
def _save_meta(self, meta: dict):
MODEL_DIR.mkdir(exist_ok=True)
META_FILE.write_text(json.dumps(meta, indent=2))
def get_status(self) -> dict:
m = self._load_meta()
has_model = MODEL_FILE.exists()
return {
"hasModel": has_model,
"sizeBytes": MODEL_FILE.stat().st_size if has_model else 0,
"sha": m.get("sha", ""),
"downloadedAt": m.get("downloaded_at", ""),
"checkedAt": m.get("checked_at", ""),
"latestSha": m.get("latest_sha", ""),
"updateAvailable": (
m.get("latest_sha") and
m.get("sha") and
m.get("latest_sha") != m.get("sha")
),
}
async def check_update(self) -> dict:
"""Check the latest commit SHA from the GitHub API."""
try:
async with aiohttp.ClientSession() as session:
async with session.get(
GITHUB_API,
headers={"Accept": "application/vnd.github.v3+json"},
timeout=aiohttp.ClientTimeout(total=10),
) as r:
if r.status != 200:
return {"ok": False, "error": f"GitHub API HTTP {r.status}"}
data = await r.json()
if not data:
return {"ok": False, "error": "No commits found"}
latest_sha = data[0]["sha"][:12]
meta = self._load_meta()
meta["latest_sha"] = latest_sha
meta["checked_at"] = time.strftime("%Y-%m-%d %H:%M UTC", time.gmtime())
self._save_meta(meta)
current = meta.get("sha", "")
update_available = bool(current and current != latest_sha)
return {
"ok": True,
"latestSha": latest_sha,
"currentSha": current,
"updateAvailable": update_available,
}
except Exception as e:
return {"ok": False, "error": str(e)}
async def download_model(self, broadcast_fn=None) -> dict:
"""Download the model from the CDN and save it locally."""
async with self._lock:
try:
save_path = str(MODEL_FILE.resolve())
MODEL_DIR.mkdir(exist_ok=True)
async def _bcast(msg, pct, received=0, total=0):
if not broadcast_fn: return
detail = f"{received/1e6:.1f} / {total/1e6:.1f} MB" if total else f"{received/1e6:.1f} MB"
await broadcast_fn({
"type": "deepcw_progress",
"msg": msg,
"pct": pct,
"detail": detail,
"savePath": save_path,
})
await _bcast("Łączenie z CDN...", 0)
async with aiohttp.ClientSession() as session:
async with session.get(
MODEL_URL,
timeout=aiohttp.ClientTimeout(total=120),
) as r:
if r.status != 200:
return {"ok": False, "error": f"HTTP {r.status}"}
total = int(r.headers.get("Content-Length", 0))
data = bytearray()
async for chunk in r.content.iter_chunked(65536):
data.extend(chunk)
pct = int(len(data) / total * 100) if total else 0
if len(data) % (256*1024) < 65536: # roughly every 256KB
await _bcast(f"Pobieranie modelu DeepCW...", pct, len(data), total)
model_bytes = bytes(data)
MODEL_FILE.write_bytes(model_bytes)
# We do NOT keep a copy in RAM — the engine (deepcw_engine)
# reads the file from disk via ONNX Runtime. An in-memory
# copy would just be 15 MB wasted.
sha = hashlib.sha256(model_bytes).hexdigest()[:12]
meta = self._load_meta()
meta["sha"] = meta.get("latest_sha", sha)
meta["sha256"] = sha
meta["downloaded_at"] = time.strftime("%Y-%m-%d %H:%M UTC", time.gmtime())
meta["size_bytes"] = len(model_bytes)
self._save_meta(meta)
await _bcast(f"✓ Model gotowy ({len(model_bytes)/1e6:.1f} MB) → {save_path}", 100, len(model_bytes), len(model_bytes))
print(f"[deepcw] Model downloaded: {len(model_bytes)/1e6:.1f} MB -> {save_path}", flush=True)
return {"ok": True, "sizeBytes": len(model_bytes), "sha": sha, "savePath": save_path}
except Exception as e:
print(f"[deepcw] Download error: {e}", flush=True)
if broadcast_fn:
await broadcast_fn({"type": "deepcw_progress", "msg": f"✗ Błąd: {e}", "pct": -1})
return {"ok": False, "error": str(e)}
def load_from_disk(self):
"""Check whether the model is present (without loading it into RAM).
This used to load the whole file into memory — 15 MB held for no
reason, since inference is done by deepcw_engine reading the file directly."""
if MODEL_FILE.exists():
_mb = MODEL_FILE.stat().st_size / 1e6
print(f"[deepcw] Model on disk: {_mb:.1f} MB ({MODEL_FILE})", flush=True)
async def auto_check_loop(self, broadcast_fn=None):
"""Update-check loop, every 6h."""
await asyncio.sleep(30) # startup delay
while True:
try:
result = await self.check_update()
if result.get("updateAvailable") and broadcast_fn:
await broadcast_fn({
"type": "deepcw_update",
"msg": f"Dostępna aktualizacja modelu DeepCW (SHA: {result['latestSha']})",
"latestSha": result["latestSha"],
})
print(f"[deepcw] Update: {result['latestSha']}", flush=True)
except Exception as e:
print(f"[deepcw] check_update error: {e}", flush=True)
await asyncio.sleep(CHECK_INTERVAL)
# Singleton
deepcw_manager = DeepCWModelManager()