Update app.py

app.py

@@ -4,21 +4,23 @@ import yt_dlp
 import os
 import subprocess
 import json
-from threading import Thread
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import spaces
 import time
 import langdetect
 import uuid
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
-print("Starting the program...")
+# Load Hugging Face Token
+HF_TOKEN = os.getenv("HF_TOKEN")
 
+print("Starting the program...")
 model_path = "Qwen/Qwen2.5-7B-Instruct"
-print(f"Loading model {model_path}...")
+
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True).cuda()
-model = model.eval()
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True).to(device).eval()
 print("Model successfully loaded.")
 
 def generate_unique_filename(extension):
@@ -33,6 +35,7 @@ def cleanup_files(*files):
 def download_youtube_audio(url):
     print(f"Downloading audio from YouTube: {url}")
     output_path = generate_unique_filename(".wav")
+    
     ydl_opts = {
         'format': 'bestaudio/best',
         'postprocessors': [{
@@ -40,10 +43,13 @@ def download_youtube_audio(url):
             'preferredcodec': 'wav',
         }],
         'outtmpl': output_path,
-        'keepvideo': True,
     }
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        ydl.download([url])
+    
+    try:
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([url])
+    except Exception as e:
+        return f"Error downloading audio: {str(e)}"
     
     if os.path.exists(output_path + ".wav"):
         os.rename(output_path + ".wav", output_path)
@@ -53,6 +59,7 @@ def download_youtube_audio(url):
 def transcribe_audio(file_path):
     print(f"Starting transcription of file: {file_path}")
     temp_audio = None
+    
     if file_path.endswith(('.mp4', '.avi', '.mov', '.flv')):
         print("Video file detected. Extracting audio using ffmpeg...")
         temp_audio = generate_unique_filename(".wav")
@@ -70,7 +77,11 @@ def transcribe_audio(file_path):
         "--timestamp", "chunk",
         "--transcript-path", output_file
     ]
-    subprocess.run(command, check=True)
+    
+    try:
+        subprocess.run(command, check=True)
+    except Exception as e:
+        return f"Error in transcription: {str(e)}"
     
     with open(output_file, "r") as f:
         transcription = json.load(f)
@@ -88,13 +99,20 @@ def generate_summary_stream(transcription):
     prompt = f"""Summarize the following video transcription in 150-300 words in {detected_language}:
     {transcription[:300000]}..."""
     
-    response, history = model.chat(tokenizer, prompt, history=[])
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    output_ids = model.generate(input_ids, max_length=500)
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    
     return response
 
 def process_youtube(url):
     if not url:
         return "Please enter a YouTube URL.", None
+    
     audio_file = download_youtube_audio(url)
+    if "Error" in audio_file:
+        return audio_file, None
+    
     transcription = transcribe_audio(audio_file)
     cleanup_files(audio_file)
     return transcription, None
@@ -126,4 +144,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     url_button.click(process_youtube, inputs=[url_input], outputs=[transcription_output, summary_output])
     summary_button.click(generate_summary_stream, inputs=[transcription_output], outputs=[summary_output])
 
-demo.launch()
+demo.launch()