Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-from audio_processing import process_audio, print_results
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
+from audio_processing import process_audio, print_results, load_models
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
 import spaces
 import torch
 
@@ -8,14 +8,18 @@ import torch
 cuda_available = torch.cuda.is_available()
 device = "cuda" if cuda_available else "cpu"
 
-# Initialize the summarization and question-answering models
+# 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")
 
 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(duration=120)
+@spaces.GPU
 def transcribe_audio(audio_file, translate, model_size):
     language_segments, final_segments = process_audio(audio_file, translate=translate, model_size=model_size)
     
@@ -38,14 +42,14 @@ def transcribe_audio(audio_file, translate, model_size):
     
     return output, full_text
 
-@spaces.GPU(duration=120)
+@spaces.GPU
 def summarize_text(text):
     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)
     return summary
 
-@spaces.GPU(duration=120)
+@spaces.GPU
 def answer_question(context, question):
     inputs = qa_tokenizer(question, context, return_tensors="pt").to(device)
     outputs = qa_model(**inputs)
@@ -54,13 +58,13 @@ def answer_question(context, question):
     answer = qa_tokenizer.decode(inputs["input_ids"][0][answer_start:answer_end])
     return answer
 
-@spaces.GPU(duration=120)
+@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(duration=120)
+@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)