Update app.py

app.py

@@ -1,27 +1,40 @@
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
+import whisper
 import gradio as gr
 
 # Load the Whisper model
-model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
-processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+model = whisper.load_model("base")  # Choose 'tiny', 'base', 'small', 'medium', or 'large'
 
 def transcribe_video(video_file):
-    audio_input = processor(audio_file=video_file.name, return_tensors="pt", sampling_rate=16000)
-    # Perform the transcription
-    with torch.no_grad():
-        predicted_ids = model.generate(**audio_input)
-    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-    
-    return transcription
-
-# Define Gradio interface
+    # Transcribe the video to generate subtitles
+    result = model.transcribe(video_file.name)
+
+    # Prepare the subtitles in SRT format
+    srt_file = "generated_subtitles.srt"
+
+    def write_srt(transcription, output_file):
+        with open(output_file, "w") as f:
+            for i, segment in enumerate(transcription['segments']):
+                start = segment['start']
+                end = segment['end']
+                text = segment['text']
+                # Format timestamps for SRT
+                start_time = whisper.utils.format_timestamp(start)
+                end_time = whisper.utils.format_timestamp(end)
+                f.write(f"{i + 1}\n")
+                f.write(f"{start_time} --> {end_time}\n")
+                f.write(f"{text.strip()}\n\n")
+
+    # Write the transcription as subtitles
+    write_srt(result, srt_file)
+
+    return srt_file  # Return the generated SRT file
+
+# Create a Gradio interface
 iface = gr.Interface(
     fn=transcribe_video,
-    inputs=gr.Video(label="Upload Video"),
-    outputs="text",
-    title="Video Transcription",
-    description="Upload a video file, and this tool will transcribe it."
+    inputs=gr.inputs.File(label="Upload Video"),
+    outputs=gr.outputs.File(label="Download Subtitles")
 )
 
-# Launch the Gradio app
-iface.launch()
+# Launch the interface
+iface.launch(share=True)