Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/bin/_server_app.py

@@ -0,0 +1,298 @@
+"""FunASR Server — unified vLLM-based inference service.
+
+Provides OpenAI-compatible API (/v1/audio/transcriptions) and REST API (/asr).
+Uses vLLM for Fun-ASR-Nano (GPU) or falls back to AutoModel for non-LLM models (SenseVoice/Paraformer).
+"""
+
+import io
+import os
+import re
+import time
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np
+import soundfile as sf
+
+try:
+    from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+    from fastapi.responses import JSONResponse
+except ImportError:
+    raise ImportError(
+        "funasr-server requires additional packages. Install with: pip install vllm fastapi uvicorn python-multipart"
+    )
+
+logger = logging.getLogger("funasr.server")
+
+
+def prepare_audio_for_inference(audio_data, sr, target_sr=16000):
+    """Return mono float32 audio at target_sr for ASR inference."""
+    audio_data = np.asarray(audio_data)
+    if audio_data.ndim > 1:
+        channel_axis = -1 if audio_data.shape[-1] <= audio_data.shape[0] else 0
+        audio_data = audio_data.mean(axis=channel_axis)
+
+    if sr != target_sr:
+        import librosa
+        audio_data = librosa.resample(audio_data, orig_sr=sr, target_sr=target_sr)
+        sr = target_sr
+
+    return audio_data.astype(np.float32), sr
+
+def create_app(device: str = "cuda", preload_model: str = "auto") -> FastAPI:
+    if preload_model == "auto":
+        preload_model = "fun-asr-nano" if device.startswith("cuda") else "sensevoice"
+
+    app = FastAPI(title="FunASR Server", version="1.3.6")
+    app.state.device = device
+    app.state.engine = None
+    app.state.vad_model = None
+    app.state.fallback_models = {}
+
+    # Non-LLM model configs (use AutoModel, no vLLM)
+    FALLBACK_CONFIGS = {
+        "sensevoice": {
+            "model": "iic/SenseVoiceSmall",
+            "vad_model": "fsmn-vad",
+            "vad_kwargs": {"max_single_segment_time": 30000},
+        },
+        "paraformer": {
+            "model": "paraformer-zh",
+            "vad_model": "fsmn-vad",
+            "punc_model": "ct-punc",
+        },
+    }
+
+    def _load_vllm_engine():
+        """Load Fun-ASR-Nano vLLM engine. Falls back to AutoModel if vLLM unavailable."""
+        if app.state.engine is not None:
+            return
+        try:
+            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+            from funasr import AutoModel as _AutoModel
+
+            logger.info("Loading Fun-ASR-Nano vLLM engine...")
+            t0 = time.time()
+            app.state.engine = FunASRNanoVLLM.from_pretrained(
+                model="FunAudioLLM/Fun-ASR-Nano-2512",
+                hub="hf",
+                device=device,
+                dtype="bf16",
+                max_model_len=4096,
+                gpu_memory_utilization=0.5,
+            )
+            logger.info(f"vLLM engine ready in {time.time()-t0:.1f}s")
+            app.state.use_vllm = True
+
+            logger.info("Loading VAD model...")
+            app.state.vad_model = _AutoModel(model="fsmn-vad", device=device, disable_update=True)
+            logger.info("VAD ready.")
+        except Exception as e:
+            logger.warning(f"vLLM failed ({e}), falling back to AutoModel for fun-asr-nano")
+            app.state.use_vllm = False
+            from funasr import AutoModel
+            cfg = {
+                "model": "FunAudioLLM/Fun-ASR-Nano-2512",
+                "hub": "hf",
+                "trust_remote_code": True,
+                "vad_model": "fsmn-vad",
+                "vad_kwargs": {"max_single_segment_time": 30000},
+                "device": device,
+                "disable_update": True,
+            }
+            app.state.fallback_models["fun-asr-nano"] = AutoModel(**cfg)
+            logger.info("Fallback AutoModel loaded for fun-asr-nano.")
+
+    def _load_fallback(name: str):
+        """Load non-LLM model via AutoModel."""
+        if name in app.state.fallback_models:
+            return app.state.fallback_models[name]
+        if name not in FALLBACK_CONFIGS:
+            return None
+        from funasr import AutoModel
+        cfg = FALLBACK_CONFIGS[name].copy()
+        cfg["device"] = device
+        cfg["disable_update"] = True
+        logger.info(f"Loading fallback model '{name}'...")
+        model = AutoModel(**cfg)
+        app.state.fallback_models[name] = model
+        return model
+
+    def _process_vllm(audio_data, sr, language=None, hotwords=None, use_spk=False):
+        """Process audio with vLLM engine (Fun-ASR-Nano)."""
+        audio_data, sr = prepare_audio_for_inference(audio_data, sr)
+
+        # VAD
+        vad_res = app.state.vad_model.generate(input=audio_data, fs=sr)
+        segments = vad_res[0]["value"] if vad_res and vad_res[0].get("value") else [[0, int(len(audio_data)*1000/sr)]]
+
+        seg_audios = []
+        seg_times = []
+        for seg in segments:
+            s0 = int(seg[0] * sr / 1000)
+            s1 = int(seg[1] * sr / 1000)
+            seg_audio = audio_data[s0:s1]
+            if len(seg_audio) > sr * 0.3:
+                seg_audios.append(seg_audio)
+                seg_times.append((seg[0], seg[1]))
+
+        if not seg_audios:
+            return {"text": "", "segments": [], "duration": len(audio_data)/sr}
+
+        # repetition_penalty is left at the neutral 1.0: the Fun-ASR-Nano vLLM
+        # engine runs in prompt-embeds mode, where any other value crashes the
+        # CUDA kernel (see issue #2948 and fun_asr_nano.vllm_utils).
+        gen_kwargs = {"max_new_tokens": 500, "repetition_penalty": 1.0}
+        if language:
+            gen_kwargs["language"] = language
+        if hotwords:
+            gen_kwargs["hotwords"] = hotwords
+
+        results = app.state.engine.generate(inputs=seg_audios, **gen_kwargs)
+
+        output_segments = []
+        full_text_parts = []
+        for r, (start_ms, end_ms) in zip(results, seg_times):
+            text = r["text"]
+            seg_info = {"text": text, "start": start_ms/1000, "end": end_ms/1000}
+            if "timestamps" in r:
+                offset = start_ms / 1000
+                seg_info["words"] = [
+                    {"word": ts["token"], "start": ts["start_time"]+offset, "end": ts["end_time"]+offset}
+                    for ts in r["timestamps"]
+                ]
+            output_segments.append(seg_info)
+            full_text_parts.append(text)
+
+        return {
+            "text": "".join(full_text_parts),
+            "segments": output_segments,
+            "duration": len(audio_data) / sr,
+        }
+
+    def _process_fallback(model_name, audio_path, language=None):
+        """Process with non-LLM model (SenseVoice/Paraformer)."""
+        model = _load_fallback(model_name)
+        kwargs = {"input": audio_path, "batch_size": 1}
+        if language:
+            kwargs["language"] = language
+        result = model.generate(**kwargs)
+        text = re.sub(r'<\|[^|]*\|>', '', result[0]["text"]).strip()
+        segments = []
+        if "sentence_info" in result[0]:
+            for s in result[0]["sentence_info"]:
+                segments.append({
+                    "start": s.get("start", 0)/1000,
+                    "end": s.get("end", 0)/1000,
+                    "text": re.sub(r'<\|[^|]*\|>', '', s.get("text", "")).strip(),
+                    "speaker": s.get("spk"),
+                })
+        return {"text": text, "segments": segments}
+
+    # Pre-load
+    if preload_model == "fun-asr-nano":
+        _load_vllm_engine()
+    else:
+        _load_fallback(preload_model)
+
+    @app.post("/v1/audio/transcriptions")
+    async def transcribe(
+        file: UploadFile = File(...),
+        model: str = Form(default="fun-asr-nano"),
+        language: Optional[str] = Form(default=None),
+        response_format: Optional[str] = Form(default="json"),
+        spk: bool = Form(default=False),
+    ):
+        content = await file.read()
+        t0 = time.perf_counter()
+
+        if model == "fun-asr-nano":
+            _load_vllm_engine()
+            if app.state.use_vllm:
+                audio_data, sr = sf.read(io.BytesIO(content))
+                result = _process_vllm(audio_data, sr, language=language, use_spk=spk)
+            else:
+                suffix = os.path.splitext(file.filename)[1] if file.filename else ".wav"
+                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                    tmp.write(content)
+                    tmp_path = tmp.name
+                try:
+                    result = _process_fallback("fun-asr-nano", tmp_path, language=language)
+                finally:
+                    os.unlink(tmp_path)
+        elif model in FALLBACK_CONFIGS:
+            suffix = os.path.splitext(file.filename)[1] if file.filename else ".wav"
+            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                tmp.write(content)
+                tmp_path = tmp.name
+            try:
+                result = _process_fallback(model, tmp_path, language=language)
+            finally:
+                os.unlink(tmp_path)
+        else:
+            raise HTTPException(400, f"Unknown model '{model}'. Available: fun-asr-nano, {', '.join(FALLBACK_CONFIGS.keys())}")
+
+        t1 = time.perf_counter()
+
+        if response_format == "verbose_json":
+            return JSONResponse({
+                "task": "transcribe",
+                "language": language or "zh",
+                "duration": result.get("duration", 0),
+                "text": result["text"],
+                "segments": [
+                    {"id": i, "start": s["start"], "end": s["end"], "text": s["text"], "words": s.get("words", [])}
+                    for i, s in enumerate(result["segments"])
+                ],
+            })
+        elif response_format == "text":
+            return JSONResponse(result["text"])
+        else:
+            return JSONResponse({"text": result["text"]})
+
+    @app.post("/asr")
+    async def asr_endpoint(
+        file: UploadFile = File(...),
+        language: Optional[str] = Form(default=None),
+        hotwords: str = Form(default=""),
+        spk: bool = Form(default=False),
+    ):
+        """Full-featured ASR endpoint with timestamps and speaker diarization."""
+        content = await file.read()
+        _load_vllm_engine()
+        hw_list = [w.strip() for w in hotwords.split(",") if w.strip()] if hotwords else None
+
+        t0 = time.perf_counter()
+        if app.state.use_vllm:
+            audio_data, sr = sf.read(io.BytesIO(content))
+            result = _process_vllm(audio_data, sr, language=language, hotwords=hw_list, use_spk=spk)
+        else:
+            suffix = ".wav"
+            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                tmp.write(content)
+                tmp_path = tmp.name
+            try:
+                result = _process_fallback("fun-asr-nano", tmp_path, language=language)
+            finally:
+                os.unlink(tmp_path)
+        t1 = time.perf_counter()
+
+        result["processing_time"] = round(t1 - t0, 3)
+        result["rtf"] = round((t1 - t0) / result["duration"], 4) if result.get("duration", 0) > 0 else 0
+        return JSONResponse(result)
+
+    @app.get("/v1/models")
+    async def list_models():
+        all_models = ["fun-asr-nano"] + list(FALLBACK_CONFIGS.keys())
+        return JSONResponse({"object": "list", "data": [{"id": n, "object": "model"} for n in all_models]})
+
+    @app.get("/health")
+    async def health():
+        loaded = []
+        if app.state.engine is not None:
+            loaded.append("fun-asr-nano (vLLM)")
+        loaded.extend(app.state.fallback_models.keys())
+        return {"status": "ok", "device": device, "models_loaded": loaded}
+
+    return app