change whisper

serve.py

@@ -7,13 +7,17 @@ import io
 from pydub import AudioSegment
 import time
 import logging
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-model = load_model(device)
+# model = load_model(device)
+
+processor = WhisperProcessor.from_pretrained("Ivydata/whisper-small-japanese").to(device)
+model = WhisperForConditionalGeneration.from_pretrained("Ivydata/whisper-small-japanese").to(device)
 
 def transcribe_audio(audio_data_bytes):
     try:
@@ -33,6 +37,35 @@ def transcribe_audio(audio_data_bytes):
         return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+    
+def transcribe_whisper(audio_data_bytes):
+    try:
+        start_time = time.time()
+        audio_segment = AudioSegment.from_mp3(io.BytesIO(audio_data_bytes))
+        
+        # Get audio data as numpy array
+        audio_data_int16 = np.array(audio_segment.get_array_of_samples())
+        # Convert to float32 normalized to [-1, 1]
+        audio_data_float32 = audio_data_int16.astype(np.float32) / 32768.0
+        
+        # Process with whisper
+        input_features = processor(audio=audio_data_float32, 
+                                sampling_rate=audio_segment.frame_rate, 
+                                return_tensors="pt").input_features.to(device)
+                                
+        predicted_ids = model.generate(input_features=input_features)
+
+        result = processor.batch_decode(predicted_ids, skip_special_tokens=True)
+        resultText = result[0] if isinstance(result, list) and len(result) > 0 else str(result)
+        result = {
+            "text": resultText
+        }
+        end_time = time.time()
+        print(f"Time taken: {end_time - start_time} seconds")
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    
 
 app = FastAPI()