""" NeMo STT Server - Arabic FastConformer LiveKit-compatible HTTP API for speech-to-text. Model: nvidia/stt_ar_fastconformer_hybrid_large_pcd_v1.0 Input: 16kHz mono PCM or WAV """ import logging import os import tempfile import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.responses import JSONResponse MODEL_NAME = "nvidia/stt_ar_fastconformer_hybrid_large_pcd_v1.0" _MODEL_FILENAME = "stt_ar_fastconformer_hybrid_large_pcd_v1.0.nemo" # Prefer env; then local nemo_stt/models/ (no HF download); else Docker /app/ _server_dir = os.path.dirname(os.path.abspath(__file__)) _local_model = os.path.join(_server_dir, "models", _MODEL_FILENAME) MODEL_PATH = os.getenv("NEMO_MODEL_PATH") or ( _local_model if os.path.isfile(_local_model) else f"/app/{_MODEL_FILENAME}" ) SAMPLE_RATE = 16000 CATT_CKPT = os.getenv("CATT_CKPT") or os.path.join(_server_dir, "models", "catt", "best_ed_mlm_ns_epoch_178.pt") logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI(title="NeMo STT Server", version="0.1.0") asr_model = None diacritizer = None def load_model(): global asr_model if asr_model is not None: return try: import nemo.collections.asr as nemo_asr if os.path.isfile(MODEL_PATH): logger.info("Loading model from %s", MODEL_PATH) asr_model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(MODEL_PATH) else: logger.info("Model file not found, loading from_pretrained %s", MODEL_NAME) asr_model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.from_pretrained(model_name=MODEL_NAME) asr_model.eval() # Disable CUDA graphs — two separate flags both need to be off. # use_cuda_graphs controls the greedy path; use_cuda_graph_decoder # controls the loop_labels path. Both hit the same broken cu_call() # that returns 5 values instead of 6 on this CUDA/PyTorch combo. try: from omegaconf import open_dict with open_dict(asr_model.cfg): asr_model.cfg.decoding.greedy.use_cuda_graphs = False asr_model.cfg.decoding.greedy.use_cuda_graph_decoder = False asr_model.change_decoding_strategy(asr_model.cfg.decoding) logger.info("CUDA graphs disabled for RNNT decoding") except Exception as _e: logger.warning("Could not disable CUDA graphs: %s", _e) logger.info("Model loaded successfully") except Exception as e: logger.exception("Failed to load model: %s", e) raise def load_diacritizer(): """Best-effort: forces transcripts to come out with tashkeel via vendored CATT. Never raises — /transcribe falls back to plain (undiacritized) text if this fails, exactly like Fasih-TTS's own diacritizer loading does.""" global diacritizer if diacritizer is not None: return try: from diacritize import Diacritizer device = "cuda" if asr_model is not None and next(asr_model.parameters()).is_cuda else None diacritizer = Diacritizer(ckpt=CATT_CKPT, device=device) logger.info("CATT diacritizer loaded (device=%s)", diacritizer.device) except Exception as e: logger.warning("CATT diacritizer unavailable, transcripts will be plain text: %s", e) def _diacritize(text: str) -> str: """Best-effort: '' on empty input, failure, or an unavailable diacritizer — the caller falls back to the plain transcript, /transcribe never breaks over this.""" if not text or diacritizer is None: return "" try: return diacritizer.diacritize_texts([text])[0] except Exception: logger.warning("Diacritization failed for transcript, returning plain text", exc_info=True) return "" @app.on_event("startup") async def startup(): load_model() load_diacritizer() @app.get("/health") async def health(): """Health check for LiveKit / load balancers.""" return { "status": "ok", "model": "stt_ar_fastconformer_hybrid_large_pcd_v1.0", "diacritizer": diacritizer is not None, } @app.post("/transcribe") async def transcribe(request: Request): """ Transcribe audio to text. Accepts: - Raw PCM: 16kHz, mono, 16-bit signed (Content-Type: application/octet-stream) - WAV file: 16kHz mono (Content-Type: audio/wav or multipart/form-data) Returns: {"text": "...", "is_final": true} """ if asr_model is None: load_model() content_type = request.headers.get("content-type", "") body = await request.body() if not body or len(body) < 1000: raise HTTPException(400, "Audio too short (min ~1s at 16kHz)") wav_path = None try: if "wav" in content_type or body[:4] == b"RIFF": wav_path = _to_16k_wav(body, ".wav") elif "mp3" in content_type or body[:3] == b"ID3" or body[:2] == b"\xff\xfb": wav_path = _to_16k_wav(body, ".mp3") else: wav_path = _pcm_to_wav_temp(body) wav_size = os.path.getsize(wav_path) if wav_path and os.path.exists(wav_path) else 0 logger.info("WAV path=%s size=%d bytes", wav_path, wav_size) output = asr_model.transcribe([str(wav_path)]) logger.info("Transcribe output type=%s len=%s first=%r", type(output).__name__, len(output) if output else 0, output[0] if output else None) if not output: text = "" elif isinstance(output, tuple) and len(output) >= 1: # (best_hypotheses, all_hypotheses) when extract_nbest hyps = output[0] first = hyps[0] if hyps else None if hasattr(first, "text"): text = first.text or "" elif isinstance(first, str): text = first else: text = str(first) if first else "" elif hasattr(output[0], "text"): text = output[0].text or "" elif isinstance(output[0], str): text = output[0] else: text = str(output[0]) if output[0] else "" logger.info("Raw output type: %s, repr: %r", type(output[0]), output[0]) text = text.strip() text_diacritized = _diacritize(text) return JSONResponse({ "text": text_diacritized or text, "text_plain": text, "diacritized": bool(text_diacritized), "is_final": True, }) except Exception as e: logger.exception("Transcription error: %s", e) raise HTTPException(500, str(e)) finally: if wav_path and os.path.exists(wav_path): try: os.unlink(wav_path) except OSError: pass def _pcm_to_wav_temp(pcm_bytes: bytes) -> str: """Convert raw PCM 16kHz mono 16-bit to WAV file.""" import wave with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: wav_path = f.name with wave.open(wav_path, "wb") as wav: wav.setnchannels(1) wav.setsampwidth(2) wav.setframerate(SAMPLE_RATE) wav.writeframes(pcm_bytes) return wav_path def _bytes_to_wav_temp(data: bytes) -> str: """Write bytes to temp WAV file (if already WAV) or try to parse.""" if data[:4] == b"RIFF": with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: f.write(data) return f.name return _pcm_to_wav_temp(data) def _to_16k_wav(audio_bytes: bytes, suffix: str) -> str: """Convert any audio to 16kHz mono WAV via ffmpeg.""" import ffmpeg with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as f: f.write(audio_bytes) tmp_path = f.name wav_path = tempfile.mktemp(suffix=".wav") try: stream = ffmpeg.input(tmp_path) stream = ffmpeg.output( stream, wav_path, acodec="pcm_s16le", ac=1, ar=SAMPLE_RATE, loglevel="error", ) ffmpeg.run(stream, overwrite_output=True) return wav_path except ffmpeg.Error as e: err = (e.stderr or b"").decode(errors="replace") raise RuntimeError(f"FFmpeg conversion failed: {err}") from e finally: if os.path.exists(tmp_path): try: os.unlink(tmp_path) except OSError: pass if __name__ == "__main__": port = int(os.getenv("NEMO_STT_PORT", "3005")) host = os.getenv("NEMO_STT_HOST", "0.0.0.0") logger.info("Starting NeMo STT server on %s:%d", host, port) uvicorn.run(app, host=host, port=port)