added multiprocessing capabilities

gradio_app.py

@@ -10,10 +10,31 @@ import shutil
 import warnings
 import os
 import gradio as gr
+import concurrent.futures
+import multiprocessing
 from notebook_lm_kokoro import generate_podcast_script, generate_audio_from_script
 warnings.filterwarnings("ignore")
 
-# A modified version of generate_audio_from_script to accept voice mapping
+# Define number of workers based on CPU cores
+NUM_WORKERS = multiprocessing.cpu_count()  # Gets total CPU cores
+
+def process_segment(entry_and_voice_map):
+    entry, voice_map = entry_and_voice_map  # Unpack the tuple
+    speaker, dialogue = entry
+    chosen_voice = voice_map.get(speaker, "af_heart")
+    print(f"Generating audio for {speaker} with voice '{chosen_voice}'...")
+    
+    pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M")
+    generator = pipeline(dialogue, voice=chosen_voice)
+    
+    segment_audio = []
+    for _, _, audio in generator:
+        segment_audio.append(audio)
+        
+    if segment_audio:
+        return np.concatenate(segment_audio, axis=0)
+    return None
+
 def generate_audio_from_script_with_voices(script, speaker1_voice, speaker2_voice, output_file):
     voice_map = {"Speaker 1": speaker1_voice, "Speaker 2": speaker2_voice}
     
@@ -29,31 +50,24 @@ def generate_audio_from_script_with_voices(script, speaker1_voice, speaker2_voic
             raise ValueError("Transcript is not a list")
 
         all_audio_segments = []
-        # Process each dialogue entry
-        for i, entry in enumerate(transcript_list):
-            if not isinstance(entry, tuple) or len(entry) != 2:
-                print(f"Skipping invalid entry {i}: {entry}")
-                continue
-
-            speaker, dialogue = entry
-            chosen_voice = voice_map.get(speaker, "af_heart")
-            print(f"Generating audio for {speaker} with voice '{chosen_voice}'...")
-
-            # Updated KPipeline initialization with explicit repo_id
-            pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M")
-            generator = pipeline(dialogue, voice=chosen_voice)
-
-            segment_audio = []
-            for j, (gs, ps, audio) in enumerate(generator):
-                # print(f"{speaker} - Segment {j}: Global Step = {gs}, Partial Step = {ps}")
-                segment_audio.append(audio)
-
-            if segment_audio:
-                segment_full = np.concatenate(segment_audio, axis=0)
-                all_audio_segments.append(segment_full)
+        # Prepare input data with voice_map for each entry
+        entries_with_voice_map = [(entry, voice_map) for entry in transcript_list]
 
+        try:
+            # Process segments in parallel
+            with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
+                # Map the processing function across all dialogue entries
+                results = list(executor.map(process_segment, entries_with_voice_map))
+                
+                # Filter out None results and combine audio segments
+                all_audio_segments = [r for r in results if r is not None]
+        
+        except Exception as e:
+            print(f"Error during audio generation: {e}")
+            return None
+        
         if not all_audio_segments:
-            print("No audio segments were generated.")
+            print("No audio segments were generated")
             return None
 
         # Add a pause between segments
@@ -106,12 +120,31 @@ def process_pdf(pdf_file, speaker1_voice, speaker2_voice, provider, api_key, ope
             f"audio_{os.path.basename(tmp_path).replace('.pdf', '.wav')}"
         )
         
-        result = generate_audio_from_script_with_voices(
-            transcript, 
-            speaker1_voice, 
-            speaker2_voice, 
-            output_file=audio_output_path
-        )
+        # result = generate_audio_from_script_with_voices(
+        #     transcript, 
+        #     speaker1_voice, 
+        #     speaker2_voice, 
+        #     output_file=audio_output_path
+        # )
+
+        # Use ProcessPoolExecutor with explicit number of workers
+        with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
+            print(f"Processing with {NUM_WORKERS} CPU cores")
+            # Submit audio generation task to the executor
+            future = executor.submit(
+                generate_audio_from_script_with_voices,
+                transcript, speaker1_voice, speaker2_voice, audio_output_path
+            )
+            result = future.result()
+            
+            if result is None:
+                return "Error generating audio", None
+            
+            return "Process complete!", result
+
+    except Exception as e:
+        print(f"Error in process_pdf: {str(e)}")
+        return f"Error processing file: {str(e)}", None
         
         if result is None:
             return "Error generating audio", None
@@ -201,16 +234,16 @@ def create_gradio_app():
                     type="filepath"
                 )
         
-        # Examples section
-        gr.Examples(
-            examples=[
-                ["sample.pdf", "af_heart", "af_nicole", "openrouter", "your-api-key-here", "https://openrouter.ai/api/v1"],
-            ],
-            inputs=[pdf_input, speaker1_voice, speaker2_voice, provider, api_key, openrouter_base],
-            outputs=[status_output, audio_output],
-            fn=process_pdf,
-            cache_examples=True,
-        )
+        # # Examples section
+        # gr.Examples(
+        #     examples=[
+        #         ["sample.pdf", "af_heart", "af_nicole", "openrouter", "your-api-key-here", "https://openrouter.ai/api/v1"],
+        #     ],
+        #     inputs=[pdf_input, speaker1_voice, speaker2_voice, provider, api_key, openrouter_base],
+        #     outputs=[status_output, audio_output],
+        #     fn=process_pdf,
+        #     cache_examples=True,
+        # )
         
         submit_btn.click(
             fn=process_pdf,