Spaces:
Runtime error
Runtime error
File size: 2,124 Bytes
d0e644a cd6e083 d0e644a fb38179 d0e644a 1859a6d d0e644a cd6e083 d0e644a cd6e083 d0e644a cd6e083 d0e644a adee17c cd6e083 d0e644a adee17c cd6e083 907d98c | 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 | import asyncio
import numpy as np
import logging
import warnings
import os
import threading
logging.getLogger("nemo_logger").setLevel(logging.ERROR)
logging.getLogger("nemo").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from asr import ASRPipeline
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
MODEL_DIR = "/app/models"
pipeline = ASRPipeline(MODEL_DIR)
@app.api_route("/", methods=["GET", "HEAD"])
def root():
return {"status": "CaptionAI server running"}
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
print("=" * 50)
print("Client connected")
print("=" * 50)
loop = asyncio.get_event_loop()
caption_queue = asyncio.Queue()
async def send_results():
while True:
text = await caption_queue.get()
if text is None:
break
try:
await websocket.send_text(text)
except:
break
sender = asyncio.create_task(send_results())
try:
while True:
data = await websocket.receive_bytes()
audio_chunk = np.frombuffer(data, dtype=np.float32)
def process(chunk=audio_chunk):
result = pipeline.transcribe_chunk(chunk)
if result:
asyncio.run_coroutine_threadsafe(
caption_queue.put(result),
loop
)
threading.Thread(target=process, daemon=True).start()
except WebSocketDisconnect:
pipeline.stop()
await caption_queue.put(None)
await sender
print("=" * 50)
print("Client disconnected")
print("=" * 50)
except Exception as e:
print(f"Error: {e}")
pipeline.stop()
await caption_queue.put(None)
await websocket.close() |