app.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. import tempfile
  2. import os
  3. import numpy as np
  4. import time
  5. import logging
  6. import os
  7. import sys
  8. from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect
  9. from faster_whisper import WhisperModel
  10. def env_bool(name: str, default: bool = False) -> bool:
  11. return os.getenv(name, str(default)).lower() in ("1", "true", "yes", "on")
  12. LOG_LEVEL = os.getenv("WHISPER_LOG_LEVEL", "INFO").upper()
  13. VAD_ENABLED = env_bool("WHISPER_VAD", True)
  14. MODEL_SIZE = str(os.getenv("MODEL_SIZE", "small").lower())
  15. MIN_AUDIO_SEC = float(os.getenv("MIN_AUDIO_SEC", "2.0"))
  16. MAX_AUDIO_SEC = float(os.getenv("MAX_AUDIO_SEC", "10.0"))
  17. logging.basicConfig(
  18. level=getattr(logging, LOG_LEVEL, logging.INFO),
  19. format="%(asctime)s [%(levelname)s] [%(name)s] %(message)s",
  20. stream=sys.stdout,
  21. force=True
  22. )
  23. logger = logging.getLogger("whisper-ws")
  24. logger.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))
  25. logger.info(f"Log level = {LOG_LEVEL}")
  26. logger.info(f"VAD enabled = {VAD_ENABLED}")
  27. app = FastAPI()
  28. # =========================
  29. # Whisper model (shared)
  30. # =========================
  31. model = WhisperModel(MODEL_SIZE, device="cuda", compute_type="float32")
  32. # =========================
  33. # REST API (Swagger / tests)
  34. # =========================
  35. @app.post("/transcribe")
  36. async def transcribe(file: UploadFile = File(...)):
  37. with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
  38. tmp.write(await file.read())
  39. tmp_path = tmp.name
  40. segments, info = model.transcribe(
  41. tmp_path,
  42. language="fr",
  43. vad_filter=True
  44. )
  45. os.remove(tmp_path)
  46. text = " ".join(seg.text for seg in segments)
  47. return {
  48. "text": text,
  49. "language": info.language
  50. }
  51. # =========================
  52. # WebSocket STT (streaming)
  53. # =========================
  54. SAMPLE_RATE = 16000
  55. @app.websocket("/ws/transcribe")
  56. async def ws_transcribe(ws: WebSocket):
  57. await ws.accept()
  58. logger.info("WebSocket connecté !")
  59. try:
  60. while True:
  61. data = await ws.receive_bytes()
  62. logger.debug(f"Chunk reçu: {len(data)} bytes")
  63. pcm = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
  64. energy = float(np.abs(pcm).mean())
  65. logger.debug(f"Énergie audio = {energy:.6f}")
  66. # silence gate léger
  67. if np.abs(pcm).mean() < 0.00005:
  68. logger.debug("Chunk silencieux ignoré")
  69. continue
  70. duration = len(pcm) / SAMPLE_RATE
  71. logger.info(f"Transcription chunk {duration:.2f}s")
  72. segments, _ = model.transcribe(
  73. pcm,
  74. language="fr",
  75. vad_filter=VAD_ENABLED
  76. )
  77. text = " ".join(seg.text for seg in segments).strip()
  78. logger.info(f"Texte: '{text}'")
  79. if text:
  80. await ws.send_json({
  81. "type": "final",
  82. "text": text
  83. })
  84. except WebSocketDisconnect:
  85. logger.info("WebSocket déconnecté")
  86. except Exception:
  87. logger.exception("Erreur WebSocket")
  88. await ws.close()