Optimized app.py with on-demand model loading and lighter models

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from audio_processing import process_audio, load_models
+from audio_processing import process_audio
 from transformers import pipeline
 import spaces
 import torch
@@ -7,41 +7,27 @@ import logging
 import traceback
 import sys
 
-# Set up logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
     handlers=[
-        logging.StreamHandler(sys.stdout),
-        logging.FileHandler('app.log')
+        logging.StreamHandler(sys.stdout)
     ]
 )
 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 Whisper model
-# print("Loading Whisper model...")
-# try:
-#     load_models()  # Load Whisper model
-# except Exception as e:
-#     logger.error(f"Error loading Whisper model: {str(e)}")
-#     raise
-
-print("Whisper model loaded successfully.")
-
 def load_summarization_model():
     logger.info("Loading summarization model...")
     try:
+        cuda_available = torch.cuda.is_available()
         summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=0 if cuda_available else -1)
+        logger.info(f"Summarization model loaded successfully on {'GPU' if cuda_available else 'CPU'}")
+        return summarizer
     except Exception as e:
         logger.warning(f"Failed to load summarization model on GPU. Falling back to CPU. Error: {str(e)}")
         summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=-1)
-    logger.info("Summarization model loaded.")
-    return summarizer
+        logger.info("Summarization model loaded successfully on CPU")
+        return summarizer
 
 def process_with_fallback(func, *args, **kwargs):
     try:
@@ -51,7 +37,6 @@ def process_with_fallback(func, *args, **kwargs):
         logger.error(traceback.format_exc())
         if "CUDA" in str(e) or "GPU" in str(e):
             logger.info("Falling back to CPU processing...")
-            # Modify kwargs to force CPU processing
             kwargs['use_gpu'] = False
             return func(*args, **kwargs)
         else:
@@ -59,24 +44,58 @@ def process_with_fallback(func, *args, **kwargs):
 
 @spaces.GPU(duration=60)
 def transcribe_audio(audio_file, translate, model_size, use_diarization):
-    return process_with_fallback(process_audio, audio_file, translate=translate, model_size=model_size, use_diarization=use_diarization)
+    logger.info(f"Starting transcription: translate={translate}, model_size={model_size}, use_diarization={use_diarization}")
+    try:
+        result = process_with_fallback(process_audio, audio_file, translate=translate, model_size=model_size, use_diarization=use_diarization)
+        logger.info("Transcription completed successfully")
+        return result
+    except Exception as e:
+        logger.error(f"Transcription failed: {str(e)}")
+        raise gr.Error(f"Transcription failed: {str(e)}")
 
 @spaces.GPU(duration=60)
 def summarize_text(text):
-    summarizer = load_summarization_model()
+    logger.info("Starting text summarization")
     try:
+        summarizer = load_summarization_model()
         summary = summarizer(text, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
+        logger.info("Summarization completed successfully")
+        return summary
     except Exception as e:
-        logger.error(f"Error during summarization: {str(e)}")
+        logger.error(f"Summarization failed: {str(e)}")
         logger.error(traceback.format_exc())
-        summary = "Error occurred during summarization. Please try again."
-    return summary
+        return "Error occurred during summarization. Please try again."
 
 @spaces.GPU(duration=60)
 def process_and_summarize(audio_file, translate, model_size, use_diarization, do_summarize):
-    transcription, full_text = transcribe_audio(audio_file, translate, model_size, use_diarization)
-    summary = summarize_text(full_text) if do_summarize else ""
-    return transcription, summary
+    logger.info(f"Starting process_and_summarize: translate={translate}, model_size={model_size}, use_diarization={use_diarization}, do_summarize={do_summarize}")
+    try:
+        language_segments, final_segments = transcribe_audio(audio_file, translate, model_size, use_diarization)
+        
+        transcription = "Detected language changes:\n\n"
+        for segment in language_segments:
+            transcription += f"Language: {segment['language']}\n"
+            transcription += f"Time: {segment['start']:.2f}s - {segment['end']:.2f}s\n\n"
+
+        transcription += f"Transcription with language detection and speaker diarization (using {model_size} model):\n\n"
+        full_text = ""
+        for segment in final_segments:
+            transcription += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}) {segment['speaker']}:\n"
+            transcription += f"Original: {segment['text']}\n"
+            if translate:
+                transcription += f"Translated: {segment['translated']}\n"
+                full_text += segment['translated'] + " "
+            else:
+                full_text += segment['text'] + " "
+            transcription += "\n"
+
+        summary = summarize_text(full_text) if do_summarize else ""
+        logger.info("Process and summarize completed successfully")
+        return transcription, summary
+    except Exception as e:
+        logger.error(f"Process and summarize failed: {str(e)}")
+        logger.error(traceback.format_exc())
+        raise gr.Error(f"Processing failed: {str(e)}")
 
 # Main interface
 with gr.Blocks() as iface:
@@ -105,8 +124,8 @@ with gr.Blocks() as iface:
     gr.Markdown(
         f"""
         ## System Information
-        - Device: {device}
-        - CUDA Available: {"Yes" if cuda_available else "No"}
+        - Device: {"CUDA" if torch.cuda.is_available() else "CPU"}
+        - CUDA Available: {"Yes" if torch.cuda.is_available() else "No"}
         
         ## ZeroGPU Support
         This application supports ZeroGPU for Hugging Face Spaces pro users.