Spaces:
Running
Running
| """LFM3-Audio-2B as an optional second engine of the Kokoro Space (for A/B tests). | |
| Off unless the Space variable LFM3_ENABLED=1. When on, the routes below sit | |
| beside the Kokoro ones with the same contract (a client only changes the path): | |
| POST /lfm3/tts/stream {"text", "voice", "sample_rate"} -> raw PCM16 mono | |
| POST /lfm3/tts/bench {"text", "voice"} -> seconds per piece | |
| GET /lfm3/tts/health | |
| The model (about 8 GB in float32 on CPU) loads on the first request, not at | |
| start-up, so Kokoro boots and answers exactly as before; transformers is not | |
| imported until then. One LFM3 synthesis at a time. "speed" is ignored (LFM3 has | |
| no speed control; X-Speed-Ignored says so). Voices: "female", "male" or a | |
| Kokoro id such as "af_heart" (see pieces.parse_voice). | |
| LiquidAI/LFM3-Audio-2B is under the LFM Open License v1.0: commercial use only | |
| below US$10M annual revenue. | |
| """ | |
| import asyncio | |
| import os | |
| import threading | |
| import time | |
| from typing import Callable | |
| import numpy as np | |
| from fastapi import APIRouter, HTTPException, Request | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| from pieces import parse_voice, system_prompt | |
| MODEL_ID = os.getenv("LFM3_MODEL_ID", "LiquidAI/LFM3-Audio-2B") | |
| MAX_NEW_TOKENS = int(os.getenv("LFM3_MAX_NEW_TOKENS", "256")) # audio frames per piece | |
| SAMPLE_RATE = 24000 # LFM3 codec output | |
| def enabled() -> bool: | |
| return os.getenv("LFM3_ENABLED", "").strip().lower() in ("1", "true", "yes") | |
| class Lfm3Engine: | |
| """Lazy, thread-safe LFM3 text-to-speech.""" | |
| def __init__(self, threads: int) -> None: | |
| self.threads = threads | |
| self.device = "cpu" | |
| self.dtype = "float32" | |
| self._model = None | |
| self._processor = None | |
| self._codec = None | |
| self._load_lock = threading.Lock() | |
| self._synth_lock = threading.Lock() | |
| def loaded(self) -> bool: | |
| return self._model is not None | |
| def _load(self) -> None: | |
| with self._load_lock: | |
| if self._model is not None: | |
| return | |
| import torch | |
| from transformers import AutoModelForMultimodalLM, AutoProcessor | |
| started = time.perf_counter() | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 | |
| self.dtype = str(dtype).replace("torch.", "") | |
| model = AutoModelForMultimodalLM.from_pretrained( | |
| MODEL_ID, dtype=dtype, device_map=self.device, trust_remote_code=True | |
| ).eval() | |
| processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| self._codec = processor.audio_tokenizer.to(device=self.device).eval() | |
| self._processor = processor | |
| self._model = model # set last: `loaded` means everything is ready | |
| print(f"[LFM3] loaded {MODEL_ID} on {self.device} ({self.dtype}) in {time.perf_counter() - started:.1f}s") | |
| def synth(self, text: str, language: str, gender: str) -> np.ndarray: | |
| """Float32 mono audio at 24 kHz for one piece of text.""" | |
| import torch | |
| self._load() | |
| conversation = [ | |
| {"role": "system", "content": system_prompt(language, gender)}, | |
| {"role": "user", "content": text}, | |
| {"role": "assistant", "content": "<|audio_start|>"}, | |
| ] | |
| with self._synth_lock, torch.no_grad(): | |
| inputs = self._processor.apply_chat_template( | |
| conversation, | |
| continue_final_message=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ).to(self.device) | |
| outputs = self._model.generate( | |
| **inputs, | |
| max_new_tokens=MAX_NEW_TOKENS, | |
| audio_temperature=1.0, | |
| audio_top_k=128, | |
| audio_repetition_penalty=1.1, | |
| audio_repetition_window=4096, | |
| return_dict_in_generate=True, | |
| ) | |
| waveform = self._codec.decode( | |
| outputs.audio_codes[:1, :, : outputs.audio_lengths[0]] | |
| ).audio_values[0] | |
| return waveform.float().cpu().numpy().reshape(-1) | |
| def build_router( | |
| *, | |
| engine: Lfm3Engine, | |
| check_key: Callable[[Request], None], | |
| split: Callable[[str], list], | |
| to_pcm16: Callable[[np.ndarray, int], bytes], | |
| max_chars: int, | |
| ) -> APIRouter: | |
| """Routes that share the Kokoro Space's API-key check, text splitting and | |
| PCM conversion.""" | |
| router = APIRouter(prefix="/lfm3/tts") | |
| async def read_body(request: Request) -> dict: | |
| try: | |
| body = await request.json() | |
| except Exception: | |
| raise HTTPException(status_code=400, detail="json body required") | |
| if not isinstance(body, dict): | |
| raise HTTPException(status_code=400, detail="json object required") | |
| return body | |
| def voice_or_400(body: dict): | |
| parsed = parse_voice(str(body.get("voice") or "female")) | |
| if parsed is None: | |
| raise HTTPException( | |
| status_code=400, detail="unknown voice: use female, male or an a?_/b?_ Kokoro id" | |
| ) | |
| return parsed | |
| def health(): | |
| return {"ok": True, "engine": "lfm3", "model": MODEL_ID, "loaded": engine.loaded, | |
| "device": engine.device, "dtype": engine.dtype, "threads": engine.threads} | |
| async def stream(request: Request): | |
| check_key(request) | |
| body = await read_body(request) | |
| text = str(body.get("text") or "").strip()[:max_chars] | |
| language, gender = voice_or_400(body) | |
| try: | |
| rate = int(body.get("sample_rate") or 16000) | |
| except (TypeError, ValueError): | |
| raise HTTPException(status_code=400, detail="invalid sample_rate") | |
| if rate not in (16000, 24000): | |
| raise HTTPException(status_code=400, detail="sample_rate must be 16000 or 24000") | |
| if not text: | |
| return JSONResponse({"detail": "empty text"}, status_code=400) | |
| pieces = split(text) | |
| def generate(): | |
| started = time.perf_counter() | |
| for i, piece in enumerate(pieces): | |
| pcm = to_pcm16(engine.synth(piece, language, gender), rate) | |
| if i == 0: | |
| print(f"[LFM3 stream] first piece in {time.perf_counter() - started:.2f}s ({len(piece)} chars)") | |
| yield pcm | |
| return StreamingResponse( | |
| generate(), # sync generator: Starlette runs it in a worker thread | |
| media_type="application/octet-stream", | |
| headers={"X-Sample-Rate": str(rate), "X-Pieces": str(len(pieces)), | |
| "X-Speed-Ignored": "1", "Cache-Control": "no-store"}, | |
| ) | |
| async def bench(request: Request): | |
| """Seconds spent synthesizing against seconds of audio, per piece | |
| (rtf below 1 is faster than real time). The first call also loads the | |
| model, so run it twice and read the second.""" | |
| check_key(request) | |
| body = await read_body(request) | |
| text = str(body.get("text") or "").strip()[:max_chars] | |
| language, gender = voice_or_400(body) | |
| if not text: | |
| return JSONResponse({"detail": "empty text"}, status_code=400) | |
| def run(): | |
| rows = [] | |
| for piece in split(text): | |
| started = time.perf_counter() | |
| audio = engine.synth(piece, language, gender) | |
| took = time.perf_counter() - started | |
| seconds = audio.size / SAMPLE_RATE | |
| rows.append({"chars": len(piece), "synth_seconds": round(took, 3), | |
| "audio_seconds": round(seconds, 3), | |
| "rtf": round(took / seconds, 3) if seconds else None}) | |
| return rows | |
| rows = await asyncio.to_thread(run) | |
| return {"device": engine.device, "threads": engine.threads, "pieces": rows, | |
| "first_piece_seconds": rows[0]["synth_seconds"] if rows else None} | |
| return router | |