from __future__ import annotations import argparse import gc import json import os from pathlib import Path import threading from typing import Any from fastapi import FastAPI, File, Form, HTTPException, UploadFile from fastapi.responses import FileResponse from fastapi.staticfiles import StaticFiles import psutil from asr_onnx_runtime import OnnxAsrEngine, OnnxCacheAsrEngine, collect_metrics APP_DIR = Path(__file__).resolve().parent STATIC_DIR = APP_DIR / "static" ENGINE: OnnxAsrEngine | None = None ENGINE_INFO: dict[str, Any] = {} ENGINE_ARGS: argparse.Namespace | None = None ENGINE_LOCK = threading.RLock() PROCESS = psutil.Process() SERVICE_PEAK_RSS = PROCESS.memory_info().rss LAST_REQUEST_PEAK_RSS = SERVICE_PEAK_RSS app = FastAPI(title="Audio8 ASR Local") app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") def release_engine() -> None: global ENGINE ENGINE = None gc.collect() class RequestMemorySampler: def __init__(self, sample_interval_s: float = 0.05) -> None: self.sample_interval_s = float(sample_interval_s) self.peak_rss_bytes = PROCESS.memory_info().rss self._stop = threading.Event() self._thread: threading.Thread | None = None def _run(self) -> None: while not self._stop.wait(self.sample_interval_s): self.peak_rss_bytes = max(self.peak_rss_bytes, PROCESS.memory_info().rss) def start(self) -> None: self.peak_rss_bytes = PROCESS.memory_info().rss self._thread = threading.Thread(target=self._run, daemon=True) self._thread.start() def stop(self) -> int: self.peak_rss_bytes = max(self.peak_rss_bytes, PROCESS.memory_info().rss) self._stop.set() if self._thread is not None: self._thread.join(timeout=1.0) return self.peak_rss_bytes def update_service_peak(candidate: int | None = None) -> None: global SERVICE_PEAK_RSS current = PROCESS.memory_info().rss values = [SERVICE_PEAK_RSS, current] if candidate is not None: values.append(int(candidate)) SERVICE_PEAK_RSS = max(values) def available_cache_precisions(bundle_dir: str | Path) -> list[str]: bundle_path = Path(bundle_dir).expanduser().resolve() metadata_path = bundle_path / "metadata.json" if not metadata_path.exists(): return [] metadata = json.loads(metadata_path.read_text(encoding="utf-8")) prefill_graph = metadata.get("graphs", {}).get("lm_cache_prefill") graph = metadata.get("graphs", {}).get("lm_cache_decode") if not graph or not prefill_graph: return [] prefill_graph_path = bundle_path / prefill_graph["path"] graph_path = bundle_path / graph["path"] values = [] if prefill_graph_path.exists() and graph_path.exists(): values.append("fp32") prefill_int8_path = prefill_graph_path.with_name(f"{prefill_graph_path.stem}_int8{prefill_graph_path.suffix}") int8_path = graph_path.with_name(f"{graph_path.stem}_int8{graph_path.suffix}") if prefill_int8_path.exists() and int8_path.exists(): values.append("int8") prefill_int4_path = prefill_graph_path.with_name(f"{prefill_graph_path.stem}_int4{prefill_graph_path.suffix}") int4_path = graph_path.with_name(f"{graph_path.stem}_int4{graph_path.suffix}") if prefill_int4_path.exists() and int4_path.exists(): values.append("int4") return values def available_audio_precisions(bundle_dir: str | Path) -> list[str]: bundle_path = Path(bundle_dir).expanduser().resolve() metadata_path = bundle_path / "metadata.json" if not metadata_path.exists(): return [] metadata = json.loads(metadata_path.read_text(encoding="utf-8")) graphs = metadata.get("graphs", {}) values = [] fp32_graph = graphs.get("audio_hidden") if fp32_graph and (bundle_path / fp32_graph["path"]).exists(): values.append("fp32") int8_graph = graphs.get("audio_hidden_int8") if int8_graph and (bundle_path / int8_graph["path"]).exists(): values.append("int8") return values def runtime_snapshot() -> dict[str, Any]: update_service_peak() bundle_dir = ENGINE_INFO.get("bundle_dir") return { **ENGINE_INFO, "runtime_loaded": ENGINE is not None, "available_cache_precisions": available_cache_precisions(bundle_dir) if bundle_dir else [], "available_audio_precisions": available_audio_precisions(bundle_dir) if bundle_dir else [], "service_peak_rss_bytes": int(SERVICE_PEAK_RSS), "last_request_peak_rss_bytes": int(LAST_REQUEST_PEAK_RSS), } def load_engine( args: argparse.Namespace, *, backend: str | None = None, cache_precision: str | None = None, audio_precision: str | None = None, release_existing: bool = False, ) -> None: global ENGINE, ENGINE_INFO if release_existing: release_engine() selected_backend = backend or args.backend selected_precision = cache_precision or args.cache_precision selected_audio_precision = audio_precision if audio_precision is not None else getattr(args, "audio_precision", None) threads = args.threads if args.threads > 0 else None if selected_backend in {"auto", "onnx_cache"}: try: engine = OnnxCacheAsrEngine( args.bundle_dir, provider=args.provider, intra_op_num_threads=threads, cache_precision=selected_precision, audio_precision=selected_audio_precision, ) ENGINE = engine ENGINE_INFO = { "bundle_dir": str(engine.bundle_dir), "backend": "onnx_cache", "cache_precision": engine.cache_precision, "cache_prefill_graph": str(engine.prefill_graph_path), "cache_graph": str(engine.cache_graph_path), "audio_precision": engine.audio_precision, "audio_graph": str(engine.audio_graph_path), "providers": engine.providers, } update_service_peak() return except Exception: if selected_backend == "onnx_cache": raise engine = OnnxAsrEngine( args.bundle_dir, provider=args.provider, intra_op_num_threads=threads, audio_precision=selected_audio_precision or "fp32", ) ENGINE = engine ENGINE_INFO = { "bundle_dir": str(engine.bundle_dir), "backend": "onnx", "cache_precision": None, "cache_graph": None, "audio_precision": engine.audio_precision, "audio_graph": str(engine.audio_graph_path), "providers": engine.providers, } update_service_peak() @app.get("/") async def index() -> FileResponse: return FileResponse(STATIC_DIR / "index.html") @app.get("/health") async def health() -> dict[str, Any]: if ENGINE is None: return {"ok": False, "error": "engine not loaded"} return {"ok": True, **runtime_snapshot()} @app.get("/api/runtime") async def runtime() -> dict[str, Any]: return {"ok": ENGINE is not None, "runtime": runtime_snapshot()} @app.post("/api/reload") async def reload_runtime( backend: str = Form("onnx_cache"), cache_precision: str | None = Form(None), audio_precision: str | None = Form(None), ) -> dict[str, Any]: if ENGINE_ARGS is None: raise HTTPException(status_code=503, detail="Server arguments are not initialized") if backend not in {"onnx_cache", "onnx", "auto"}: raise HTTPException(status_code=400, detail=f"Unsupported backend: {backend}") if cache_precision: cache_precision = cache_precision.lower().strip() if cache_precision not in {"fp32", "int8", "int4", "auto"}: raise HTTPException(status_code=400, detail=f"Unsupported cache precision: {cache_precision}") else: cache_precision = str(ENGINE_INFO.get("cache_precision") or "").lower().strip() or ENGINE_ARGS.cache_precision if audio_precision: audio_precision = audio_precision.lower().strip() if audio_precision not in {"fp32", "int8", "auto"}: raise HTTPException(status_code=400, detail=f"Unsupported audio precision: {audio_precision}") else: audio_precision = str(ENGINE_INFO.get("audio_precision") or "").lower().strip() or None with ENGINE_LOCK: try: load_engine( ENGINE_ARGS, backend=backend, cache_precision=cache_precision, audio_precision=audio_precision, release_existing=True, ) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc return {"ok": True, "runtime": runtime_snapshot()} @app.get("/metrics") async def metrics() -> dict[str, Any]: data = collect_metrics() data["runtime"] = runtime_snapshot() return data @app.post("/asr") async def asr( audio: UploadFile = File(...), language: str | None = Form(None), max_new_tokens: int = Form(128), cache_precision: str | None = Form(None), audio_precision: str | None = Form(None), hotwords: str | None = Form(None), hotword_topk: int = Form(50), hotword_start_boost: float = Form(6.0), hotword_continuation_boost: float = Form(8.0), ) -> dict[str, Any]: global LAST_REQUEST_PEAK_RSS if ENGINE is None: raise HTTPException(status_code=503, detail="ASR engine is not loaded") audio_bytes = await audio.read() if not audio_bytes: raise HTTPException(status_code=400, detail="Empty audio upload") if ENGINE_ARGS is not None: requested_cache = cache_precision.lower().strip() if cache_precision else None requested_audio = audio_precision.lower().strip() if audio_precision else None if requested_cache and requested_cache not in {"fp32", "int8", "int4", "auto"}: raise HTTPException(status_code=400, detail=f"Unsupported cache precision: {requested_cache}") if requested_audio and requested_audio not in {"fp32", "int8", "auto"}: raise HTTPException(status_code=400, detail=f"Unsupported audio precision: {requested_audio}") current_cache = str(ENGINE_INFO.get("cache_precision") or "").lower().strip() current_audio = str(ENGINE_INFO.get("audio_precision") or "").lower().strip() target_cache = requested_cache or current_cache or ENGINE_ARGS.cache_precision target_audio = requested_audio or current_audio or getattr(ENGINE_ARGS, "audio_precision", None) should_reload = ( ENGINE_INFO.get("backend") == "onnx_cache" and ( (requested_cache is not None and target_cache != current_cache) or (requested_audio is not None and target_audio != current_audio) ) ) if should_reload: with ENGINE_LOCK: try: load_engine( ENGINE_ARGS, backend="onnx_cache", cache_precision=target_cache, audio_precision=target_audio, release_existing=True, ) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) from exc sampler = RequestMemorySampler() sampler.start() try: with ENGINE_LOCK: if ENGINE is None: raise HTTPException(status_code=503, detail="ASR engine is not loaded") result = ENGINE.transcribe( audio_bytes, language=language or None, max_new_tokens=max_new_tokens, hotwords=hotwords or None, hotword_topk=int(hotword_topk), hotword_start_boost=float(hotword_start_boost), hotword_continuation_boost=float(hotword_continuation_boost), ) LAST_REQUEST_PEAK_RSS = max(int(sampler.stop()), PROCESS.memory_info().rss) update_service_peak(LAST_REQUEST_PEAK_RSS) result["request_peak_rss_bytes"] = int(LAST_REQUEST_PEAK_RSS) result["runtime"] = runtime_snapshot() return result except Exception as exc: LAST_REQUEST_PEAK_RSS = max(int(sampler.stop()), PROCESS.memory_info().rss) update_service_peak(LAST_REQUEST_PEAK_RSS) raise HTTPException(status_code=500, detail=str(exc)) from exc def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run the local Audio8 ASR ONNX web UI.") parser.add_argument( "--bundle_dir", default=os.environ.get("AUDIO8_ASR_BUNDLE", "model_bundle"), ) parser.add_argument( "--backend", choices=["auto", "onnx", "onnx_cache"], default=os.environ.get("ASR_BACKEND", "auto"), ) parser.add_argument("--host", default=os.environ.get("HOST", "127.0.0.1")) parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "7860"))) parser.add_argument("--provider", default=os.environ.get("ORT_PROVIDER", "CPUExecutionProvider")) parser.add_argument("--cache_precision", default=os.environ.get("ASR_CACHE_PRECISION", "auto")) parser.add_argument("--audio_precision", choices=["fp32", "int8", "auto"], default=os.environ.get("ASR_AUDIO_PRECISION", "auto")) parser.add_argument("--threads", type=int, default=int(os.environ.get("ORT_THREADS", "0"))) return parser.parse_args() def main() -> None: global ENGINE_ARGS args = parse_args() ENGINE_ARGS = args load_engine(args) _run_uvicorn(args) def _run_uvicorn(args: argparse.Namespace) -> None: import uvicorn uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": main()