| from __future__ import annotations
|
|
|
| import json
|
| import uuid
|
| from contextlib import asynccontextmanager
|
| from pathlib import Path
|
| from typing import Annotated, Optional
|
|
|
| from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
| from fastapi.responses import JSONResponse, StreamingResponse
|
|
|
| from audio.asr import create_engine
|
| from audio.config import Settings, get_settings
|
| from audio.metrics import gpu_memory
|
| from audio.pipeline import AudioPipeline
|
| from audio.queue import GpuSerialQueue, QueueFullError
|
| from audio.translation import TranslationClient
|
|
|
| _queue: GpuSerialQueue | None = None
|
|
|
|
|
| @asynccontextmanager
|
| async def lifespan(app: FastAPI):
|
| global _queue
|
| settings = get_settings()
|
| asr = create_engine(settings)
|
| pipeline = AudioPipeline(settings, asr, TranslationClient(settings))
|
| _queue = GpuSerialQueue(
|
| pipeline,
|
| max_queue=settings.max_queue,
|
| gpu_concurrency=asr.replicas,
|
| )
|
| app.state.settings = settings
|
| yield
|
|
|
|
|
| app = FastAPI(title="Translation Audio Service", lifespan=lifespan)
|
|
|
|
|
| def queue() -> GpuSerialQueue:
|
| if _queue is None:
|
| raise RuntimeError("service not started")
|
| return _queue
|
|
|
|
|
| def _save_upload(upload_dir: str, filename: str | None, data: bytes) -> Path:
|
| suffix = Path(filename or "audio.bin").suffix or ".bin"
|
| dest = Path(upload_dir) / f"{uuid.uuid4().hex}{suffix}"
|
| dest.parent.mkdir(parents=True, exist_ok=True)
|
| dest.write_bytes(data)
|
| return dest
|
|
|
|
|
| @app.get("/health")
|
| def health() -> dict:
|
| snap = gpu_memory()
|
| settings: Settings = get_settings()
|
| engine = queue().pipeline.asr
|
| return {
|
| "ok": True,
|
| "layer": "audio",
|
| "asr_backend": getattr(engine, "backend", settings.asr_backend),
|
| "asr_model": getattr(engine, "model_name", settings.asr_model),
|
| "device": getattr(engine, "device", settings.asr_device),
|
| "compute_type": getattr(engine, "compute_type", settings.asr_compute_type),
|
| "batch_size": getattr(engine, "batch_size", settings.asr_batch_size),
|
| "replicas": getattr(engine, "replicas", settings.asr_replicas),
|
| "gpu": {
|
| "name": snap.name,
|
| "used_mb": snap.used_mb,
|
| "total_mb": snap.total_mb,
|
| "free_mb": snap.free_mb,
|
| "compute_capability": (
|
| f"{snap.compute_major}.{snap.compute_minor}"
|
| if snap.compute_major is not None
|
| else None
|
| ),
|
| },
|
| }
|
|
|
|
|
| @app.post("/translate")
|
| async def translate_media(
|
| file: Annotated[UploadFile, File()],
|
| target_language: Annotated[str, Form()] = "en",
|
| source_language: Annotated[Optional[str], Form()] = None,
|
| translate: Annotated[bool, Form()] = True,
|
| ) -> JSONResponse:
|
| settings: Settings = get_settings()
|
| dest = _save_upload(settings.upload_dir, file.filename, await file.read())
|
| try:
|
| result = await queue().run(
|
| source=dest,
|
| target_language=target_language,
|
| source_language=source_language or None,
|
| translate=translate,
|
| )
|
| except QueueFullError as exc:
|
| raise HTTPException(status_code=429, detail=str(exc)) from exc
|
| except Exception as exc:
|
| raise HTTPException(status_code=500, detail=str(exc)) from exc
|
| return JSONResponse(result)
|
|
|
|
|
| @app.post("/translate/stream")
|
| async def translate_stream(
|
| file: Annotated[UploadFile, File()],
|
| source_language: Annotated[Optional[str], Form()] = None,
|
| ) -> StreamingResponse:
|
| settings: Settings = get_settings()
|
| dest = _save_upload(settings.upload_dir, file.filename, await file.read())
|
|
|
| async def events():
|
| try:
|
| async for item in queue().stream(
|
| source=dest,
|
| source_language=source_language or None,
|
| ):
|
| yield f"data: {json.dumps(item, ensure_ascii=False)}\n\n"
|
| except QueueFullError as exc:
|
| yield f"data: {json.dumps({'type': 'error', 'detail': str(exc)})}\n\n"
|
| except Exception as exc:
|
| yield f"data: {json.dumps({'type': 'error', 'detail': str(exc)})}\n\n"
|
|
|
| return StreamingResponse(events(), media_type="text/event-stream")
|
|
|
