Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 import yt_dlp
 import torch
@@ -6,196 +7,73 @@ import whisper
 import threading
 from queue import Queue
 import time
-import os
+
+# Ensure the cache directory exists
+MODEL_CACHE = "./models"
+os.makedirs(MODEL_CACHE, exist_ok=True)
 
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
-# Load Whisper model
-try:
-    processor = AutoProcessor.from_pretrained("openai/whisper-large-v3-turbo")
-    model = AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-large-v3-turbo")
-    if device == "cuda":
-        model.to(device)
-    print("Whisper model loaded successfully")
-except Exception as e:
-    print(f"Error loading Whisper model: {e}")
+# Load Whisper model with cache_dir
+processor = AutoProcessor.from_pretrained("openai/whisper-large-v3-turbo", cache_dir=MODEL_CACHE)
+model = AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-large-v3-turbo", cache_dir=MODEL_CACHE).to(device)
 
 # Load summarization model
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)
 
-class YouTubeHandler:
-    def __init__(self):
-        self.current_video_id = None
-        self.transcript_queue = Queue()
-        self.processing = False
-        self.download_thread = None
-
-    def search_youtube(self, query):
-        """Search YouTube and return video details"""
-        ydl_opts = {
-            "format": "bestaudio/best",
-            "quiet": True,
-            "noplaylist": True,
-            "default_search": "ytsearch1",
-        }
-
-        try:
-            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-                info = ydl.extract_info(f"ytsearch:{query}", download=False)
-                if "entries" not in info or not info["entries"]:
-                    return None, None
-
-                video = info["entries"][0]
-                video_id = video['id']
-                video_url = f"https://www.youtube.com/watch?v={video_id}"
-                return video_id, video['title']
-        except Exception as e:
-            print(f"Error searching YouTube: {e}")
-            return None, None
-
-    def process_audio(self, video_id):
-        """Download and process audio in chunks"""
-        ydl_opts = {
-            "format": "bestaudio/best",
-            "quiet": True,
-            "outtmpl": "temp_audio",
-            "postprocessors": [{
-                "key": "FFmpegExtractAudio",
-                "preferredcodec": "wav",
-            }],
-        }
-
-        try:
-            # Download audio
-            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-                url = f"https://www.youtube.com/watch?v={video_id}"
-                ydl.download([url])
-
-            # Process audio file
-            audio = whisper.load_audio("temp_audio.wav")
-
-            chunk_duration = 10  # seconds
-            sample_rate = 16000
-            chunk_size = chunk_duration * sample_rate
-
-            for i in range(0, len(audio), chunk_size):
-                if not self.processing:
-                    break
-
-                chunk = audio[i:i + chunk_size]
-                if len(chunk) < chunk_size/2:  # Skip very short chunks
-                    continue
-
-                input_features = processor(
-                    chunk,
-                    sampling_rate=16000,
-                    return_tensors="pt"
-                ).input_features
-
-                if device == "cuda":
-                    input_features = input_features.to(device)
-
-                predicted_ids = model.generate(input_features)
-                transcription = processor.batch_decode(
-                    predicted_ids,
-                    skip_special_tokens=True
-                )[0]
+audio_path = "/tmp/temp_audio.wav"  # Use /tmp for temporary files
 
-                if transcription.strip():
-                    self.transcript_queue.put(transcription)
-
-                time.sleep(0.1)  # Prevent overwhelming the system
-
-        except Exception as e:
-            print(f"Error processing audio: {e}")
-        finally:
-            # Cleanup
-            if os.path.exists("temp_audio.wav"):
-                try:
-                    os.remove("temp_audio.wav")
-                except:
-                    pass
-
-    def start_processing(self, video_id):
-        """Start processing a new video"""
-        self.stop_processing()
+def process_audio(video_id):
+    """Download and transcribe YouTube audio"""
+    ydl_opts = {
+        "format": "bestaudio/best",
+        "quiet": True,
+        "outtmpl": "/tmp/temp_audio.%(ext)s",  # Save inside /tmp
+        "postprocessors": [{"key": "FFmpegExtractAudio", "preferredcodec": "wav"}],
+    }
+    try:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([f"https://www.youtube.com/watch?v={video_id}"])
 
-        self.current_video_id = video_id
-        self.processing = True
-        self.download_thread = threading.Thread(
-            target=self.process_audio,
-            args=(video_id,)
-        )
-        self.download_thread.start()
-        return f"Processing video {video_id}"
+        audio = whisper.load_audio(audio_path)
+        input_features = processor(audio, sampling_rate=16000, return_tensors="pt").input_features.to(device)
+        predicted_ids = model.generate(input_features)
+        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
 
-    def stop_processing(self):
-        """Stop current processing"""
-        if self.processing:
-            self.processing = False
-            if self.download_thread:
-                self.download_thread.join()
-            return "Processing stopped"
-        return "No video processing"
+        return transcription
+    except Exception as e:
+        return f"Error processing audio: {e}"
+    finally:
+        if os.path.exists(audio_path):
+            os.remove(audio_path)  # Clean up
 
 def summarize_text(text):
     """Summarize transcript into a short version."""
     if len(text.split()) < 10:
         return "Transcript too short for summarization."
-
     try:
         summary = summarizer(text, max_length=50, min_length=10, do_sample=False)
         return summary[0]['summary_text']
     except Exception as e:
-        print(f"Error summarizing text: {e}")
         return "Summarization failed"
 
-# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# 🎥 YouTube Real-Time Transcriber")
-
-    youtube_handler = YouTubeHandler()
-
-    with gr.Row():
-        word_input = gr.Textbox(label="Enter a Word or Phrase")
-        search_button = gr.Button("Search Video")
-
+    
+    video_input = gr.Textbox(label="Enter YouTube Video ID")
+    process_button = gr.Button("Transcribe")
     video_embed = gr.HTML()
-    status_output = gr.Textbox(label="Status", interactive=False)
     transcript_output = gr.Textbox(label="Live Transcript", interactive=False)
     summary_output = gr.Textbox(label="Summary", interactive=False)
 
-    stop_button = gr.Button("Stop Processing")
-
-    def handle_video(word):
-        video_id, title = youtube_handler.search_youtube(word)
-        if not video_id:
-            return "<p>No video found.</p>", "No video found.", "", ""
-
-        embed_html = f'''<iframe width="560" height="315" src="https://www.youtube.com/embed/{video_id}" frameborder="0" allowfullscreen></iframe>'''
-        status = youtube_handler.start_processing(video_id)
-        return embed_html, status, "", ""
-
-    def update_transcript():
-        transcript = ""
-        while not youtube_handler.transcript_queue.empty():
-            transcript += " " + youtube_handler.transcript_queue.get()
-
-        if transcript.strip():
-            summary = summarize_text(transcript)
-            transcript_output.update(value=transcript)
-            summary_output.update(value=summary)
-
-    def background_update():
-        while True:
-            update_transcript()
-            time.sleep(1)  # Update every second
-
-    threading.Thread(target=background_update, daemon=True).start()
+    def handle_video(video_id):
+        embed_html = f"""<iframe width='560' height='315' src='https://www.youtube.com/embed/{video_id}' frameborder='0' allowfullscreen></iframe>"""
+        transcript = process_audio(video_id)
+        summary = summarize_text(transcript) if transcript else ""
+        return embed_html, transcript, summary
 
-    search_button.click(handle_video, inputs=word_input, outputs=[video_embed, status_output, transcript_output, summary_output])
-    stop_button.click(youtube_handler.stop_processing, outputs=status_output)
+    process_button.click(handle_video, inputs=[video_input], outputs=[video_embed, transcript_output, summary_output])
 
 demo.launch(debug=True)