whisper_api_v2.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. import asyncio
  2. import numpy as np
  3. import soundfile as sf
  4. import torch
  5. import websockets
  6. import whisper
  7. from io import BytesIO
  8. # Whisper 모델 로드
  9. model = whisper.load_model("large") # 원하는 모델 크기 지정
  10. async def process_audio(websocket):
  11. print("클라이언가 연결되었습니다.")
  12. transcription = ['']
  13. try:
  14. async for message in websocket:
  15. # 브라우저에서 받은 오디오 데이터 (webm 형식)
  16. audio_data = message
  17. # webm 데이터를 메모리에서 처리
  18. with BytesIO(audio_data) as audio_file:
  19. # soundfile로 webm을 읽고 PCM 데이터로 변환
  20. audio, sample_rate = sf.read(audio_file, dtype='float32')
  21. # Whisper로 텍스트 변환
  22. result = model.transcribe(audio, language="ko", fp16=torch.cuda.is_available())
  23. text = result['text'].strip()
  24. # 변환된 텍스트를 리스트에 추가
  25. transcription.append(text)
  26. print(text) # 콘솔에 출력
  27. # 클라이언트로 텍스트 전송
  28. await websocket.send(text)
  29. except websockets.ConnectionClosed:
  30. print("클라이언트 연결이 종료되었습니다.")
  31. except Exception as e:
  32. print(f"오류 발생: {e}")
  33. finally:
  34. print("WebSocket 연결 정리 완료")
  35. async def main():
  36. # WebSocket 서버 시작
  37. server = await websockets.serve(
  38. process_audio,
  39. "localhost",
  40. 8765,
  41. ping_interval=20, # 클라이언트와의 연결 상태 확인
  42. ping_timeout=10 # 타임아웃 설정
  43. )
  44. print("WebSocket 서버가 localhost:8765에서 실행 중...")
  45. await server.wait_closed()
  46. if __name__ == "__main__":
  47. asyncio.run(main())