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

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from audio_processing import process_audio, load_models
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering, pipeline
 import spaces
 import torch
 import logging
@@ -13,44 +13,54 @@ 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...")
-try:
-    load_models()  # Load Whisper and diarization models
-except Exception as e:
-    logger.error(f"Error loading Whisper and diarization models: {str(e)}")
-    raise
+# Initialize model variables
+summarizer_model = None
+summarizer_tokenizer = None
+qa_model = None
+qa_tokenizer = None
 
+# Load Whisper model
+print("Loading Whisper model...")
 try:
-    summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn").to(device)
-    summarizer_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
+    load_models()  # Load Whisper model
 except Exception as e:
-    logger.error(f"Error loading summarization model: {str(e)}")
+    logger.error(f"Error loading Whisper 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
+print("Whisper model loaded successfully.")
+
+def load_summarization_model():
+    global summarizer_model, summarizer_tokenizer
+    if summarizer_model is None:
+        logger.info("Loading summarization model...")
+        summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6").to(device)
+        summarizer_tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
+        logger.info("Summarization model loaded.")
 
-print("Models loaded successfully.")
+def load_qa_model():
+    global qa_model, qa_tokenizer
+    if qa_model is None:
+        logger.info("Loading QA model...")
+        qa_model = AutoModelForQuestionAnswering.from_pretrained("distilbert-base-cased-distilled-squad").to(device)
+        qa_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased-distilled-squad")
+        logger.info("QA model loaded.")
 
 @spaces.GPU
-def transcribe_audio(audio_file, translate, model_size):
-    language_segments, final_segments = process_audio(audio_file, translate=translate, model_size=model_size)
+def transcribe_audio(audio_file, translate, model_size, use_diarization):
+    language_segments, final_segments = process_audio(audio_file, translate=translate, model_size=model_size, use_diarization=use_diarization)
     
     output = "Detected language changes:\n\n"
     for segment in language_segments:
         output += f"Language: {segment['language']}\n"
         output += f"Time: {segment['start']:.2f}s - {segment['end']:.2f}s\n\n"
 
-    output += f"Transcription with language detection and speaker diarization (using {model_size} model):\n\n"
+    output += f"Transcription with language detection {f'and speaker diarization' if use_diarization else ''} (using {model_size} model):\n\n"
     full_text = ""
     for segment in final_segments:
-        output += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}) {segment['speaker']}:\n"
-        output += f"Original: {segment['text']}\n"
+        output += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']})"
+        if use_diarization:
+            output += f" {segment['speaker']}:"
+        output += f"\nOriginal: {segment['text']}\n"
         if translate:
             output += f"Translated: {segment['translated']}\n"
             full_text += segment['translated'] + " "
@@ -62,6 +72,7 @@ def transcribe_audio(audio_file, translate, model_size):
 
 @spaces.GPU
 def summarize_text(text):
+    load_summarization_model()
     inputs = summarizer_tokenizer(text, max_length=1024, truncation=True, return_tensors="pt").to(device)
     summary_ids = summarizer_model.generate(inputs["input_ids"], max_length=150, min_length=50, do_sample=False)
     summary = summarizer_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
@@ -69,6 +80,7 @@ def summarize_text(text):
 
 @spaces.GPU
 def answer_question(context, question):
+    load_qa_model()
     inputs = qa_tokenizer(question, context, return_tensors="pt").to(device)
     outputs = qa_model(**inputs)
     answer_start = torch.argmax(outputs.start_logits)
@@ -77,14 +89,14 @@ def answer_question(context, question):
     return answer
 
 @spaces.GPU
-def process_and_summarize(audio_file, translate, model_size):
-    transcription, full_text = transcribe_audio(audio_file, translate, model_size)
+def process_and_summarize(audio_file, translate, model_size, use_diarization):
+    transcription, full_text = transcribe_audio(audio_file, translate, model_size, use_diarization)
     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)
+def qa_interface(audio_file, translate, model_size, use_diarization, question):
+    _, full_text = transcribe_audio(audio_file, translate, model_size, use_diarization)
     answer = answer_question(full_text, question)
     return answer
 
@@ -96,13 +108,14 @@ with gr.Blocks() as iface:
         audio_input = gr.Audio(type="filepath")
         translate_checkbox = gr.Checkbox(label="Enable Translation")
         model_dropdown = gr.Dropdown(choices=["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"], label="Whisper Model Size", value="small")
+        diarization_checkbox = gr.Checkbox(label="Enable Speaker Diarization")
         transcribe_button = gr.Button("Transcribe and Summarize")
         transcription_output = gr.Textbox(label="Transcription")
         summary_output = gr.Textbox(label="Summary")
         
         transcribe_button.click(
             process_and_summarize,
-            inputs=[audio_input, translate_checkbox, model_dropdown],
+            inputs=[audio_input, translate_checkbox, model_dropdown, diarization_checkbox],
             outputs=[transcription_output, summary_output]
         )
     
@@ -110,13 +123,14 @@ with gr.Blocks() as iface:
         qa_audio_input = gr.Audio(type="filepath")
         qa_translate_checkbox = gr.Checkbox(label="Enable Translation")
         qa_model_dropdown = gr.Dropdown(choices=["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"], label="Whisper Model Size", value="small")
+        qa_diarization_checkbox = gr.Checkbox(label="Enable Speaker Diarization")
         question_input = gr.Textbox(label="Ask a question about the audio")
         qa_button = gr.Button("Get Answer")
         answer_output = gr.Textbox(label="Answer")
         
         qa_button.click(
             qa_interface,
-            inputs=[qa_audio_input, qa_translate_checkbox, qa_model_dropdown, question_input],
+            inputs=[qa_audio_input, qa_translate_checkbox, qa_model_dropdown, qa_diarization_checkbox, question_input],
             outputs=answer_output
         )