Update app.py

app.py

@@ -3,23 +3,40 @@ from audio_processing import process_audio, load_models
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
 import spaces
 import torch
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Check if CUDA is available
 cuda_available = torch.cuda.is_available()
 device = "cuda" if cuda_available else "cpu"
+logger.info(f"Using device: {device}")
 
 # Load models globally
 print("Loading models...")
-load_models()  # Load Whisper and diarization models
-
-summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn").to(device)
-summarizer_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
+try:
+    load_models()  # Load Whisper and diarization models
+except Exception as e:
+    logger.error(f"Error loading Whisper and diarization models: {str(e)}")
+    raise
+
+try:
+    summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn").to(device)
+    summarizer_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
+except Exception as e:
+    logger.error(f"Error loading summarization model: {str(e)}")
+    raise
+
+try:
+    qa_model = AutoModelForQuestionAnswering.from_pretrained("distilbert-base-cased-distilled-squad").to(device)
+    qa_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased-distilled-squad")
+except Exception as e:
+    logger.error(f"Error loading QA model: {str(e)}")
+    raise
 
-qa_model = AutoModelForQuestionAnswering.from_pretrained("distilbert-base-cased-distilled-squad").to(device)
-qa_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased-distilled-squad")
 print("Models loaded successfully.")
 
-
 @spaces.GPU
 def transcribe_audio(audio_file, translate, model_size):
     language_segments, final_segments = process_audio(audio_file, translate=translate, model_size=model_size)
@@ -43,7 +60,6 @@ def transcribe_audio(audio_file, translate, model_size):
     
     return output, full_text
 
-
 @spaces.GPU
 def summarize_text(text):
     inputs = summarizer_tokenizer(text, max_length=1024, truncation=True, return_tensors="pt").to(device)
@@ -51,7 +67,6 @@ def summarize_text(text):
     summary = summarizer_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return summary
 
-
 @spaces.GPU
 def answer_question(context, question):
     inputs = qa_tokenizer(question, context, return_tensors="pt").to(device)
@@ -61,21 +76,18 @@ def answer_question(context, question):
     answer = qa_tokenizer.decode(inputs["input_ids"][0][answer_start:answer_end])
     return answer
 
-
 @spaces.GPU
 def process_and_summarize(audio_file, translate, model_size):
     transcription, full_text = transcribe_audio(audio_file, translate, model_size)
     summary = summarize_text(full_text)
     return transcription, summary
 
-
 @spaces.GPU
 def qa_interface(audio_file, translate, model_size, question):
     _, full_text = transcribe_audio(audio_file, translate, model_size)
     answer = answer_question(full_text, question)
     return answer
 
-    
 # Main interface
 with gr.Blocks() as iface:
     gr.Markdown("# WhisperX Audio Transcription, Translation, Summarization, and QA (with ZeroGPU support)")
@@ -109,10 +121,14 @@ with gr.Blocks() as iface:
         )
 
     gr.Markdown(
-        """
+        f"""
+        ## System Information
+        - Device: {device}
+        - CUDA Available: {"Yes" if cuda_available else "No"}
+        
         ## ZeroGPU Support
         This application supports ZeroGPU for Hugging Face Spaces pro users. 
-        GPU-intensive tasks are automatically optimized for better performance.
+        GPU-intensive tasks are automatically optimized for better performance when available.
         """
     )