Update app.py

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
-import whisper
 import moviepy.editor as mp
 from deep_translator import GoogleTranslator
 from pydub import AudioSegment
 import os
 import tempfile
+import torch
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
 def extract_audio(video_path):
     video = mp.VideoFileClip(video_path)
@@ -14,9 +15,23 @@ def extract_audio(video_path):
     return audio_path
 
 def generate_subtitles(audio_path):
-    model = whisper.load_model("base")
-    result = model.transcribe(audio_path)
-    return result["segments"]
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+    model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base").to(device)
+
+    # Load and preprocess the audio
+    audio_input, _ = librosa.load(audio_path, sr=16000)
+    input_features = processor(audio_input, sampling_rate=16000, return_tensors="pt").input_features.to(device)
+
+    # Generate token ids
+    predicted_ids = model.generate(input_features)
+
+    # Decode token ids to text
+    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
+
+    # For simplicity, we're returning a single segment with the full transcription
+    # In a more advanced implementation, you might want to split this into multiple segments
+    return [{"start": 0, "end": len(audio_input) / 16000, "text": transcription[0]}]
 
 def translate_subtitles(subtitles, target_language):
     translator = GoogleTranslator(source='auto', target=target_language)