Update app.py

app.py

@@ -15,6 +15,10 @@ import textwrap
 # Device configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Initialize models and tokenizers outside the functions
+tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1").to(device)
+tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-large-v1")
+
 # Translation function
 @spaces.GPU(duration=120)
 def translate(source_text, source_lang, target_lang, batch_size=16):
@@ -53,14 +57,10 @@ def split_text_into_sentences(text):
     sentences = sentence_endings.split(text)
     return [sentence.strip() for sentence in sentences if sentence.strip()]
 
-# Function to generate audio from text
+# Function to generate audio for a single sentence
 @spaces.GPU(duration=120)
-def generate_wav_from_text(prompt, description, output_file_prefix):
-    # Load TTS model and tokenizer
-    tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1").to(device)
-    tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-large-v1")
-
-    input_ids = tts_tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+def generate_single_wav_from_text(sentence, description, output_file_prefix, tts_model, tts_tokenizer):
+    input_ids = tts_tokenizer(sentence, return_tensors="pt").input_ids.to(device)
     prompt_input_ids = tts_tokenizer(description, return_tensors="pt").input_ids.to(device)
     
     generation = tts_model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
@@ -90,33 +90,40 @@ def update_target_lang_options(source_lang):
     }
     return gr.update(choices=options.get(source_lang, []), value=options.get(source_lang, [])[0])
 
-# Main Gradio function
-def process_pdf(pdf_file, translate_checkbox, source_lang, target_lang, description):
-    print("Extracting text from PDF...")
-    text = pdf_to_text(pdf_file.name)
-    
-    # Translate if the translation checkbox is selected
-    if translate_checkbox:
-        print("Translating text...")
-        text = translate(text, source_lang, target_lang)
-    
-    print("Splitting text into sentences...")
-    sentences = split_text_into_sentences(text)
+# Function to process sentences for audio generation
+def process_sentences_for_audio(sentences, description, tts_model, tts_tokenizer):
     audio_files = []
     outputs = []
-    
+
     for i, sentence in enumerate(sentences):
         print(f"Generating audio for sentence {i+1}...")
         output_file_prefix = f"sentence_{i+1}"
-        audio_file = generate_wav_from_text(sentence, description, output_file_prefix)
+        audio_file = generate_single_wav_from_text(sentence, description, output_file_prefix, tts_model, tts_tokenizer)
         audio_files.append(audio_file)
         outputs.append((sentence, audio_file))
-        
-        # Display each sentence and its corresponding audio immediately
+
         print(f"Generated sentence: {sentence}")
         gr.Markdown(f"**Sentence**: {sentence}")
         gr.Audio(value=audio_file, label=sentence)
     
+    return outputs, audio_files
+
+# Main Gradio function
+def process_pdf(pdf_file, translate_checkbox, source_lang, target_lang, description, tts_model, tts_tokenizer):
+    print("Extracting text from PDF...")
+    text = pdf_to_text(pdf_file.name)
+    
+    # Translate if translation checkbox is selected
+    if translate_checkbox:
+        print("Translating text...")
+        text = translate(text, source_lang, target_lang)
+    
+    print("Splitting text into sentences...")
+    sentences = split_text_into_sentences(text)
+    
+    # Process sentences for audio generation
+    outputs, audio_files = process_sentences_for_audio(sentences, description, tts_model, tts_tokenizer)
+
     print("Combining all audio files...")
     combined_output_file = "sentences_combined.wav"
     combine_wav_files(combined_output_file, *audio_files)
@@ -142,7 +149,7 @@ with gr.Blocks() as demo:
             output_group = gr.Group()
 
     def handle_process(pdf_input, translate_checkbox, source_lang, target_lang, description):
-        outputs, combined_output_file = process_pdf(pdf_input, translate_checkbox, source_lang, target_lang, description)
+        outputs, combined_output_file = process_pdf(pdf_input, translate_checkbox, source_lang, target_lang, description, tts_model, tts_tokenizer)
         with output_group:
             for sentence, audio_file in outputs:
                 gr.Markdown(f"**Sentence**: {sentence}")