app.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  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. # silence gate léger
  65. if np.abs(pcm).mean() < 0.001:
  66. logger.debug("Chunk silencieux ignoré")
  67. continue
  68. duration = len(pcm) / SAMPLE_RATE
  69. logger.info(f"Transcription chunk {duration:.2f}s")
  70. segments, _ = model.transcribe(
  71. pcm,
  72. language="fr",
  73. vad_filter=VAD_ENABLED
  74. )
  75. text = " ".join(seg.text for seg in segments).strip()
  76. logger.info(f"Texte: '{text}'")
  77. if text:
  78. await ws.send_json({
  79. "type": "final",
  80. "text": text
  81. })
  82. except WebSocketDisconnect:
  83. logger.info("WebSocket déconnecté")
  84. except Exception:
  85. logger.exception("Erreur WebSocket")
  86. await ws.close()