Update app.py

app.py

@@ -1,132 +1,53 @@
-import gradio as gr
-import torchaudio
-import torchaudio.transforms as T
-from transformers import pipeline
-import requests
-from pydub import AudioSegment
-from pydub.silence import split_on_silence
-import io
-import os
-from bs4 import BeautifulSoup
-import re
-import numpy as np
-from moviepy.video.io.VideoFileClip import VideoFileClip
-import soundfile as sf
-from spellchecker import SpellChecker
-import tempfile
+import torch
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
-# Load the transcription model
-transcription_pipeline = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")
-spell = SpellChecker()
+# Check if CUDA is available and set the device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
-def download_audio_from_url(url):
-    try:
-        if "share" in url:
-            print("Processing shareable link...")
-            response = requests.get(url)
-            soup = BeautifulSoup(response.content, 'html.parser')
-            video_tag = soup.find('video')
-            if video_tag and 'src' in video_tag.attrs:
-                video_url = video_tag['src']
-                print(f"Extracted video URL: {video_url}")
-            else:
-                raise ValueError("Direct video URL not found in the shareable link.")
-        else:
-            video_url = url
-        
-        print(f"Downloading video from URL: {video_url}")
-        response = requests.get(video_url)
-        audio_bytes = response.content
-        print(f"Successfully downloaded {len(audio_bytes)} bytes of data")
-        return audio_bytes
-    except Exception as e:
-        print(f"Error in download_audio_from_url: {str(e)}")
-        raise
-
-def correct_spelling(text):
-    words = text.split()
-    corrected_words = [spell.correction(word) or word for word in words]
-    return ' '.join(corrected_words)
+# Load the Whisper model and processor
+model_name = "openai/whisper-base"
+processor = WhisperProcessor.from_pretrained(model_name)
+model = WhisperForConditionalGeneration.from_pretrained(model_name).to(device)
 
-def format_transcript(transcript):
-    sentences = transcript.split('.')
-    formatted_transcript = []
-    current_speaker = None
-    for sentence in sentences:
-        if ':' in sentence:
-            speaker, content = sentence.split(':', 1)
-            if speaker != current_speaker:
-                formatted_transcript.append(f"\n\n{speaker.strip()}:{content.strip()}.")
-                current_speaker = speaker
-            else:
-                formatted_transcript.append(f"{content.strip()}.")
-        else:
-            formatted_transcript.append(sentence.strip() + '.')
-    return ' '.join(formatted_transcript)
-
-def transcribe_audio(video_bytes):
+def transcribe_audio(audio_file):
     try:
-        with open("temp_video.mp4", "wb") as f:
-            f.write(video_bytes)
-
-        video = VideoFileClip("temp_video.mp4")
-        audio = video.audio
-
-        audio.write_audiofile("temp_audio.wav", fps=16000, nbytes=2, codec='pcm_s16le')
-
-        audio_data, sample_rate = sf.read("temp_audio.wav")
+        # Load and preprocess the audio
+        audio_input, sample_rate = sf.read(audio_file)
+        input_features = processor(audio_input, sampling_rate=sample_rate, return_tensors="pt").input_features.to(device)
 
-        if len(audio_data.shape) > 1:
-            audio_data = audio_data.mean(axis=1)
+        # Generate token ids
+        predicted_ids = model.generate(input_features)
 
-        audio_data = audio_data.astype(np.float32) / np.max(np.abs(audio_data))
+        # Decode token ids to text
+        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
 
-        result = transcription_pipeline(audio_data)
-        transcript = result['text']
-
-        transcript = correct_spelling(transcript)
-        transcript = format_transcript(transcript)
-
-        os.remove("temp_video.mp4")
-        os.remove("temp_audio.wav")
-
-        return transcript
+        return transcription[0]
     except Exception as e:
         print(f"Error in transcribe_audio: {str(e)}")
         raise
 
+# Update the transcribe_video function to use the new transcribe_audio function
 def transcribe_video(url):
     try:
         print(f"Attempting to download audio from URL: {url}")
         audio_bytes = download_audio_from_url(url)
         print(f"Successfully downloaded {len(audio_bytes)} bytes of audio data")
         
+        # Save audio bytes to a temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
+            temp_audio.write(audio_bytes)
+            temp_audio_path = temp_audio.name
+
         print("Starting audio transcription...")
-        transcript = transcribe_audio(audio_bytes)
+        transcript = transcribe_audio(temp_audio_path)
         print("Transcription completed successfully")
         
+        # Clean up the temporary file
+        os.unlink(temp_audio_path)
+
         return transcript
     except Exception as e:
         error_message = f"An error occurred: {str(e)}"
         print(error_message)
-        return error_message
-
-def download_transcript(transcript):
-    with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as temp_file:
-        temp_file.write(transcript)
-        temp_file_path = temp_file.name
-    return temp_file_path
-
-# Create the Gradio interface
-with gr.Blocks(title="Video Transcription") as demo:
-    gr.Markdown("# Video Transcription")
-    video_url = gr.Textbox(label="Video URL")
-    transcribe_button = gr.Button("Transcribe")
-    transcript_output = gr.Textbox(label="Transcript", lines=20)
-    download_button = gr.Button("Download Transcript")
-    download_link = gr.File(label="Download Transcript")
-    
-    transcribe_button.click(fn=transcribe_video, inputs=video_url, outputs=transcript_output)
-    download_button.click(fn=download_transcript, inputs=transcript_output, outputs=download_link)
-
-demo.launch()
+        return error_message