Update app.py

app.py

@@ -1,132 +1,75 @@
-import torch
 import gradio as gr
-import yt_dlp as youtube_dl
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
-import tempfile
 import os
-import time
-
-# Constants
-MODEL_NAME = "openai/whisper-large-v3"
-BATCH_SIZE = 8
-FILE_LIMIT_MB = 25  # File size limit in MB
-YT_LENGTH_LIMIT_S = 3600  # 1 hour YouTube file limit
-
-# Device configuration (CUDA if available)
-device = 0 if torch.cuda.is_available() else "cpu"
-
-# Load Whisper model and processor
-processor = WhisperProcessor.from_pretrained(MODEL_NAME)
-model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+import numpy as np
+import librosa
+
+# Initialize Whisper model
+processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
+
+# Set light green theme
+theme = gr.themes.Base(
+    primary_hue="emerald",
+    secondary_hue="emerald",
+    neutral_hue="gray",
+)
 
-def transcribe_audio(inputs):
-    """Transcribe audio using Whisper model."""
-    if inputs is None:
-        raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
+def validate_file(file):
+    # Check file size (25 MB limit)
+    file_size_mb = os.path.getsize(file) / (1024 * 1024)
+    if file_size_mb > 25:
+        return False, f"File size is {file_size_mb:.2f} MB. Please upload a file smaller than 25 MB."
     
-    # Check file size (max 25MB)
-    if os.path.getsize(inputs) > FILE_LIMIT_MB * 1024 * 1024:
-        raise gr.Error(f"File size exceeds {FILE_LIMIT_MB}MB limit.")
+    # Check file extension
+    file_extension = os.path.splitext(file)[1].lower()
+    if file_extension not in ['.mp3', '.wav']:
+        return False, "Only .mp3 and .wav formats are supported."
     
-    # Preprocess audio input
-    audio_input = processor(inputs, return_tensors="pt", sampling_rate=16000).to(device)
-    
-    # Generate transcription
-    predicted_ids = model.generate(audio_input.input_values, max_length=448)
-    
-    # Decode the transcription output
-    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-    return transcription
-
-def _return_yt_html_embed(yt_url):
-    """Return YouTube embed HTML for display."""
-    video_id = yt_url.split("?v=")[-1]
-    html_embed = f'<center><iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"></iframe></center>'
-    return html_embed
+    return True, "File is valid."
 
-def download_yt_audio(yt_url, filename):
-    """Download audio from a YouTube URL."""
-    info_loader = youtube_dl.YoutubeDL()
+def transcribe_audio(audio_file):
+    # Validate the file first
+    is_valid, message = validate_file(audio_file)
+    if not is_valid:
+        return message
     
     try:
-        info = info_loader.extract_info(yt_url, download=False)
-    except youtube_dl.utils.DownloadError as err:
-        raise gr.Error(f"Download error: {str(err)}")
+        # Load audio file
+        speech_array, sampling_rate = librosa.load(audio_file, sr=16000)
+        
+        # Process the audio file
+        input_features = processor(speech_array, sampling_rate=16000, return_tensors="pt").input_features
+        
+        # Generate token ids
+        predicted_ids = model.generate(input_features)
+        
+        # Decode token ids to text
+        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+        
+        return transcription
+    except Exception as e:
+        return f"An error occurred during transcription: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(theme=theme) as demo:
+    gr.Markdown("# Audio Transcription with Whisper")
+    gr.Markdown("Upload an audio file (.mp3 or .wav) of maximum 25MB to get the transcription.")
     
-    # Check video length
-    file_length_s = int(info.get("duration", 0))
+    with gr.Row():
+        with gr.Column():
+            audio_input = gr.Audio(type="filepath", label="Upload Audio File")
+            submit_btn = gr.Button("Transcribe", variant="primary")
+        
+        with gr.Column():
+            output = gr.Textbox(label="Transcription Result", lines=10)
     
-    if file_length_s > YT_LENGTH_LIMIT_S:
-        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
-        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
-        raise gr.Error(f"Maximum YouTube video length is {yt_length_limit_hms}, but video is {file_length_hms}.")
+    submit_btn.click(fn=transcribe_audio, inputs=audio_input, outputs=output)
     
-    # Download the video
-    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
-    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
-        try:
-            ydl.download([yt_url])
-        except youtube_dl.utils.ExtractorError as err:
-            raise gr.Error(f"Error while downloading video: {str(err)}")
-
-def yt_transcribe(yt_url):
-    """Transcribe YouTube video using Whisper model."""
-    html_embed = _return_yt_html_embed(yt_url)
-
-    with tempfile.TemporaryDirectory() as tmpdirname:
-        filepath = os.path.join(tmpdirname, "video.mp4")
-        download_yt_audio(yt_url, filepath)
-
-        with open(filepath, "rb") as file:
-            audio_input = file.read()
-
-    # Process and transcribe
-    transcription = transcribe_audio(audio_input)
-    return html_embed, transcription
-
-# Create Gradio interface
-demo = gr.Blocks()
-
-# Microphone transcription interface
-mf_transcribe = gr.Interface(
-    fn=transcribe_audio,
-    inputs=[
-        gr.inputs.Audio(source="microphone", type="filepath", optional=True),
-    ],
-    outputs="text",
-    layout="horizontal",
-    theme="huggingface",
-    title="Whisper Transcription (Microphone)",
-    description="Transcribe audio from your microphone. File size limit is 25MB."
-)
-
-# File upload transcription interface
-file_transcribe = gr.Interface(
-    fn=transcribe_audio,
-    inputs=[
-        gr.inputs.Audio(source="upload", type="filepath", optional=True, label="Audio file"),
-    ],
-    outputs="text",
-    layout="horizontal",
-    theme="huggingface",
-    title="Whisper Transcription (File)",
-    description="Upload an audio file to transcribe. File size limit is 25MB."
-)
-
-# YouTube video transcription interface
-yt_transcribe = gr.Interface(
-    fn=yt_transcribe,
-    inputs=[
-        gr.inputs.Textbox(lines=1, placeholder="Paste YouTube URL", label="YouTube URL"),
-    ],
-    outputs=["html", "text"],
-    layout="horizontal",
-    theme="huggingface",
-    title="Free Transcript Maker",
-    description="Upload an audio file (WAV, MP3, etc.) up to 25MB to get its transcription. The transcript will be displayed and available for download. Please use responsibly."
-)
-
-with demo:
-    gr.TabbedInterface([mf_transcribe, file_transcribe, yt_transcribe], ["Microphone", "Audio file", "YouTube"])
+    gr.Markdown("### Limitations")
+    gr.Markdown("- Maximum file size: 25 MB")
+    gr.Markdown("- Supported formats: .mp3 and .wav")
+    gr.Markdown("- Uses the Whisper base model which works best with clear audio")
 
-demo.launch(enable_queue=True)
+# Launch the app
+demo.launch()