Update model/wav2vec2.py

model/wav2vec2.py

@@ -1,8 +1,9 @@
 import torch
-import librosa
+import torchaudio
 import numpy as np
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 import warnings
+import io
 
 warnings.filterwarnings("ignore")
 
@@ -25,13 +26,22 @@ class Wav2Vec2:
             
         self.model.eval()
         
-    def preprocess_audio(self, audio_data: np.ndarray, original_sr: int) -> np.ndarray:
+    def preprocess_audio(self, audio_data: torch.Tensor, original_sr: int) -> np.ndarray:
         """오디오 데이터 전처리"""
         # 샘플링 레이트 변환
         if original_sr != self.sampling_rate:
-            audio_data = librosa.resample(audio_data, orig_sr=original_sr, target_sr=self.sampling_rate)
+            resampler = torchaudio.transforms.Resample(original_sr, self.sampling_rate)
+            audio_data = resampler(audio_data)
         
-        # 정규화
+        # numpy로 변환
+        if isinstance(audio_data, torch.Tensor):
+            audio_data = audio_data.numpy()
+            
+        # 스테레오를 모노로 변환 (필요한 경우)
+        if len(audio_data.shape) > 1:
+            audio_data = np.mean(audio_data, axis=0)
+        
+        # float32로 변환
         if audio_data.dtype != np.float32:
             audio_data = audio_data.astype(np.float32)
             
@@ -45,7 +55,7 @@ class Wav2Vec2:
         """오디오 파일을 텍스트로 변환"""
         try:
             # 오디오 파일 로드
-            audio_data, sample_rate = librosa.load(audio_file_path, sr=None)
+            audio_data, sample_rate = torchaudio.load(audio_file_path)
             
             # 전처리
             audio_data = self.preprocess_audio(audio_data, sample_rate)