fix: revert back to old version

gradio_app.py

@@ -1,255 +1,281 @@
+# filepath: /Users/udaylunawat/Downloads/Data-Science-Projects/NotebookLM_clone/gradio_app.py
 import os
 import tempfile
-import shutil
-import ast
-import numpy as np
+import gradio as gr
+from notebook_lm_kokoro import generate_podcast_script, KPipeline
 import soundfile as sf
+import numpy as np
+import ast
+import shutil
 import warnings
-import multiprocessing
-import concurrent.futures
-import urllib.request
-import pathlib
-
-try:
-    from moshi.models.tts import TTSModel
-except ImportError:
-    print("Moshi TTSModel not available — install Kyutai’s version via pip.")
-    TTSModel = None
-
-from notebook_lm_kokoro import (
-    generate_podcast_script,
-    generate_audio_from_script,
-    generate_audio_kyutai,
-    KPipeline,
-)
-
-import sys
-
-# Diagnostic: where is ~/.cache pointing?
-print(f"[DEBUG] HOME = {os.environ.get('HOME')}")
-print(f"[DEBUG] XDG_CACHE_HOME = {os.environ.get('XDG_CACHE_HOME')}")
-print(f"[DEBUG] Trying to create /.cache/test.txt")
-
-try:
-    os.makedirs("/.cache", exist_ok=True)
-    with open("/.cache/test.txt", "w") as f:
-        f.write("test")
-    print("[DEBUG] Successfully wrote to /.cache")
-except Exception as e:
-    print(f"[DEBUG] ❌ Failed to write to /.cache: {e}")
-
-# Set cache dirs BEFORE importing torch, transformers, or moshi
-os.environ["HF_HOME"] = "/tmp/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
-os.environ["XDG_CACHE_HOME"] = "/tmp/huggingface"
-os.environ["TORCH_HOME"] = "/tmp/torch"
-os.environ["MOSHI_CACHE_DIR"] = "/tmp/moshi"
-
-# Explicitly override ~/.cache
-os.environ["HOME"] = "/tmp/home"
-os.makedirs("/tmp/home", exist_ok=True)
-
-for path in [
-    "/tmp/.cache",
-    "/tmp/huggingface",
-    "/tmp/huggingface/transformers",
-    "/tmp/torch",
-    "/tmp/moshi",
-]:
-    os.makedirs(path, exist_ok=True)
-
-if not os.path.exists("/.cache"):
-    try:
-        os.symlink("/tmp/.cache", "/.cache")
-        print("[DEBUG] Symlinked /.cache to /tmp/.cache")
-    except Exception as e:
-        print(f"[DEBUG] Couldn't symlink /.cache: {e}")
-
+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")
-NUM_WORKERS = multiprocessing.cpu_count()
 
-def ensure_gradio_frpc():
-    """
-    Ensures the frpc binary is present in the location Gradio expects.
-    Avoids /.cache symlinks (which are not writable in HF Spaces).
-    """
-    gradio_temp_dir = os.environ.get("GRADIO_TEMP_DIR", "/tmp/gradio")
-    target_dir = os.path.join(gradio_temp_dir, "frpc")
-    os.makedirs(target_dir, exist_ok=True)
-
-    frpc_file = os.path.join(target_dir, "frpc_linux_amd64_v0.3")
-
-    if not os.path.exists(frpc_file):
-        print(f"[INFO] Downloading frpc binary to: {frpc_file}")
-        try:
-            url = "https://cdn-media.huggingface.co/frpc-gradio-0.3/frpc_linux_amd64"
-            urllib.request.urlretrieve(url, frpc_file)
-            os.chmod(frpc_file, 0o755)  # Make it executable
-            print("[SUCCESS] frpc binary downloaded and made executable.")
-        except Exception as e:
-            print(f"[ERROR] Failed to download frpc binary: {e}")
-    else:
-        print("[INFO] frpc binary already exists at expected path.")
+# 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
+    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)
-    return np.concatenate([audio for _, _, audio in generator], axis=0) if generator else None
+    
+    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):
-    print("[DEBUG] Raw transcript string:")
-    print(script)
-    
     voice_map = {"Speaker 1": speaker1_voice, "Speaker 2": speaker2_voice}
+    
+    # Clean up the script string if needed
+    script = script.strip()
+    if not script.startswith("[") or not script.endswith("]"):
+        print("Invalid transcript format. Expected a list of tuples.")
+        return None
+
     try:
         transcript_list = ast.literal_eval(script)
         if not isinstance(transcript_list, list):
             raise ValueError("Transcript is not a list")
 
-        results = []
-        for entry in transcript_list:
-            audio = process_segment((entry, voice_map))
-            if audio is not None:
-                results.append(audio)
-        if not results:
+        all_audio_segments = []
+        # 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")
             return None
+
+        # Add a pause between segments
         sample_rate = 24000
         pause = np.zeros(sample_rate, dtype=np.float32)
-        final_audio = results[0]
-        for seg in results[1:]:
+        final_audio = all_audio_segments[0]
+        for seg in all_audio_segments[1:]:
             final_audio = np.concatenate((final_audio, pause, seg), axis=0)
+
         sf.write(output_file, final_audio, sample_rate)
+        print(f"Saved final audio as {output_file}")
         return output_file
+
     except Exception as e:
-        print(f"Transcript parse error: {e}")
+        print(f"Error processing transcript: {e}")
         return None
 
-def process_pdf(pdf_file, speaker1_voice, speaker2_voice, kyutai_voice1, kyutai_voice2,
-                provider, openai_key=None, openrouter_key=None, openrouter_base=None, tts_engine=None):
-    try:
-        if provider == "openai" and not openai_key:
-            return "OpenAI API key is required", None
-        if provider == "openrouter" and not openrouter_key:
-            return "OpenRouter API key is required", None
 
-        if provider in ["openai", "kyutai"]:
-            os.environ["OPENAI_API_KEY"] = openai_key or ""
+def process_pdf(pdf_file, speaker1_voice, speaker2_voice, provider, api_key, openrouter_base=None):
+    """Process the uploaded PDF file and generate audio"""
+    try:
+    
+        # Set API configuration based on provider
+        if provider == "openai":
+            os.environ["OPENAI_API_KEY"] = api_key
             os.environ["OPENROUTER_API_BASE"] = "https://api.openai.com/v1"
-        if provider in ["openrouter", "kyutai"]:
-            os.environ["OPENAI_API_KEY"] = openrouter_key or ""
+        else:
+            os.environ["OPENAI_API_KEY"] = api_key
             os.environ["OPENROUTER_API_BASE"] = openrouter_base or "https://openrouter.ai/api/v1"
-
+        # Check if we received a valid file
         if pdf_file is None:
             return "No file uploaded", None
-
-        tmp_path = pdf_file.name
-
-        script_provider = "openrouter" if provider == "kyutai" and openrouter_key else provider
-        transcript, _ = generate_podcast_script(pdf_file.name, provider=script_provider)
-
+            
+        # Create a temporary file with .pdf extension
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+            # For Gradio uploads, we need to copy the file
+            shutil.copy2(pdf_file.name, tmp.name)
+            tmp_path = tmp.name
+            
+        print(f"Uploaded PDF saved at {tmp_path}")
+
+        # Generate transcript using your existing function
+        transcript, transcript_path = generate_podcast_script(tmp_path, provider=provider)
         if transcript is None:
-            return "Transcript generation failed: got None", None
-        if not transcript.strip().startswith("["):
-            return f"Malformed transcript:\n{transcript}", None
-
-        audio_path = os.path.join(os.path.dirname(tmp_path), f"audio_{os.path.basename(tmp_path).replace('.pdf', '.wav')}")
+            return "Error generating transcript", None
+
+        # Define an output file path for the generated audio
+        audio_output_path = os.path.join(
+            os.path.dirname(tmp_path),
+            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
+        # )
+
+        # 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
 
-        if tts_engine == "kyutai":
-            result = generate_audio_kyutai(transcript, kyutai_voice1, kyutai_voice2, audio_path)
-        else:
-            result = generate_audio_from_script_with_voices(transcript, speaker1_voice, speaker2_voice, audio_path)
+    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
+        
+        return "Process complete!", result
 
-        return ("Process complete!", result) if result else ("Error generating audio", None)
     except Exception as e:
-        print(f"process_pdf error: {e}")
-        return f"Error: {e}", None
-
-def update_ui(provider, tts_engine):
-    return [
-        gr.update(visible=tts_engine == "kokoro"),
-        gr.update(visible=tts_engine == "kokoro"),
-        gr.update(visible=tts_engine == "kyutai"),
-        gr.update(visible=tts_engine == "kyutai"),
-        gr.update(visible=provider in ["openai", "kyutai"]),
-        gr.update(visible=provider in ["openrouter", "kyutai"]),
-        gr.update(visible=provider == "openrouter"),
-    ]
+        print(f"Error in process_pdf: {str(e)}")
+        return f"Error processing file: {str(e)}", None
+
 
 def create_gradio_app():
-    css = ".gradio-container {max-width: 900px !important}"
+    # Add CSS for better styling
+    css = """
+    .gradio-container {max-width: 900px !important}
+    """
+    
     with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
-        gr.Markdown("# 🎧 PDF to Podcast — NotebookLM + Kokoro/Kyutai")
-
+        gr.Markdown(
+            """
+            # 📚 NotebookLM-Kokoro TTS App
+            Upload a PDF, choose voices, and generate conversational audio using Kokoro TTS.
+            """
+        )
+        
         with gr.Row():
-            with gr.Column(scale=1.5):
-                pdf_input = gr.File(file_types=[".pdf"], type="filepath", label="📄 Upload your PDF")
-                provider = gr.Radio(["openai", "openrouter"], value="openrouter", label="🧠 API Provider")
-                tts_engine = gr.Radio(["kokoro", "kyutai"], value="kokoro", label="🎤 TTS Engine")
-
-                speaker1_voice = gr.Dropdown(["af_heart","af_bella","hf_beta"], value="af_heart", label="Speaker 1 Voice", visible=True)
-                speaker2_voice = gr.Dropdown(["af_nicole","af_heart","bf_emma"], value="bf_emma", label="Speaker 2 Voice", visible=True)
-                kyutai_voice1 = gr.Dropdown(
-                    [
-                        "expresso/ex03-ex01_happy_001_channel1_334s.wav",
-                        "expresso/ex03-ex02_narration_001_channel1_674s.wav",
-                        "vctk/p226_023_mic1.wav"
-                    ],
-                    value="expresso/ex03-ex01_happy_001_channel1_334s.wav",
-                    label="Kyutai Voice 1",
-                    visible=True
+            with gr.Column(scale=2):
+                pdf_input = gr.File(
+                    label="Upload PDF Document",
+                    file_types=[".pdf"],
+                    type="filepath"
                 )
-
-                kyutai_voice2 = gr.Dropdown(
-                    [
-                        "expresso/ex03-ex01_happy_001_channel1_334s.wav",
-                        "expresso/ex03-ex02_narration_001_channel1_674s.wav",
-                        "vctk/p225_023_mic1.wav"
-                    ],
-                    value="expresso/ex03-ex02_narration_001_channel1_674s.wav",
-                    label="Kyutai Voice 2",
-                    visible=True
+                
+                with gr.Row():
+                    speaker1_voice = gr.Dropdown(
+                        choices=["af_heart", "af_bella", "hf_beta"],
+                        value="af_heart",
+                        label="Speaker 1 Voice"
+                    )
+                    speaker2_voice = gr.Dropdown(
+                        choices=["af_nicole", "af_heart", "bf_emma"],
+                        value="bf_emma",
+                        label="Speaker 2 Voice"
+                    )
+                
+
+                with gr.Group():
+                    provider = gr.Radio(
+                        choices=["openai", "openrouter"],
+                        value="openrouter",
+                        label="API Provider"
+                    )
+                    
+                    api_key = gr.Textbox(
+                        label="API Key",
+                        placeholder="Enter your API key here...",
+                        type="password",
+                        elem_classes="api-input"
+                    )
+                    
+                    openrouter_base = gr.Textbox(
+                        label="OpenRouter Base URL (optional)",
+                        placeholder="https://openrouter.ai/api/v1",
+                        visible=False,
+                        elem_classes="api-input"
+                    )
+
+                    # Show/hide OpenRouter base URL based on provider selection
+                    def toggle_openrouter_base(provider_choice):
+                        return gr.update(visible=provider_choice == "openrouter")
+                    
+                    provider.change(
+                        fn=toggle_openrouter_base,
+                        inputs=[provider],
+                        outputs=[openrouter_base]
+                    )
+                
+                submit_btn = gr.Button("🎙️ Generate Audio", variant="primary")
+            
+            with gr.Column(scale=2):
+                status_output = gr.Textbox(
+                    label="Status",
+                    placeholder="Processing status will appear here..."
                 )
-
-                with gr.Accordion("🔐 API Keys", open=True):
-                    openai_key = gr.Textbox(type="password", label="OpenAI Key", show_label=True, visible=True)
-                    openrouter_key = gr.Textbox(type="password", label="OpenRouter Key", show_label=True, visible=True)
-                    openrouter_base = gr.Textbox(placeholder="https://openrouter.ai/api/v1", label="OpenRouter Base URL", visible=True)
-
-                submit_btn = gr.Button("🎙️ Generate Podcast", variant="primary")
-
-            with gr.Column(scale=1):
-                status_output = gr.Textbox(label="📝 Status", interactive=False)
-                audio_output = gr.Audio(type="filepath", label="🎵 Your Podcast")
-
-            submit_btn.click(
-                process_pdf,
-                inputs=[pdf_input, speaker1_voice, speaker2_voice, kyutai_voice1, kyutai_voice2,
-                        provider, openai_key, openrouter_key, openrouter_base, tts_engine],
-                outputs=[status_output, audio_output]
-            )
-
-            provider.change(update_ui, [provider, tts_engine], 
-                            [speaker1_voice, speaker2_voice, kyutai_voice1, kyutai_voice2,
-                             openai_key, openrouter_key, openrouter_base])
-            tts_engine.change(update_ui, [provider, tts_engine], 
-                              [speaker1_voice, speaker2_voice, kyutai_voice1, kyutai_voice2,
-                               openai_key, openrouter_key, openrouter_base])
-
-        gr.Markdown("""
-        **📌 Tips**
-        - Pick your API provider and then set appropriate keys.
-        - Choose **TTS Engine** (Kokoro/Kyutai) to reveal relevant voice options.
-        - Works well with clean, structured PDFs.
-        """)
-
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    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,
+        # )
+        
+        submit_btn.click(
+            fn=process_pdf,
+            inputs=[
+                pdf_input, 
+                speaker1_voice, 
+                speaker2_voice, 
+                provider,
+                api_key,
+                openrouter_base
+            ],
+            outputs=[status_output, audio_output],
+            api_name="generate"
+        )
+        
+        gr.Markdown(
+            """
+            ### 📝 Notes
+            - Make sure your PDF is readable and contains text (not scanned images)
+            - Processing large PDFs may take a few minutes
+            - You need a valid OpenAI/OpenRouter API key set as environment variable
+            """
+        )
+    
     return app
 
-
-ensure_gradio_frpc()
 if __name__ == "__main__":
-    create_gradio_app().queue().launch(server_name="0.0.0.0", server_port=7860, share=True, debug=True, pwa=True)
+    demo = create_gradio_app()
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        debug=True,
+        pwa=True
+    )

notebook_lm_kokoro.py

@@ -23,14 +23,6 @@ import asyncio
 import ast
 import json
 import warnings
-import torch
-import time
-try:
-    from moshi.models.loaders import CheckpointInfo
-    from moshi.models.tts import DEFAULT_DSM_TTS_REPO, DEFAULT_DSM_TTS_VOICE_REPO, TTSModel
-except ImportError:
-    CheckpointInfo = None
-    TTSModel = None
 warnings.filterwarnings("ignore")
 
 # Set your OpenAI (or OpenRouter) API key from the environment
@@ -38,17 +30,6 @@ openai.api_key = os.getenv("OPENAI_API_KEY")
 # For OpenRouter compatibility, set the API base if provided.
 openai.api_base = os.getenv("OPENROUTER_API_BASE", "https://api.openai.com/v1")
 
-# Set cache dirs BEFORE importing torch, transformers, or moshi
-os.environ["HF_HOME"] = "/tmp/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
-os.environ["XDG_CACHE_HOME"] = "/tmp/huggingface"
-os.environ["TORCH_HOME"] = "/tmp/torch"
-os.environ["MOSHI_CACHE_DIR"] = "/tmp/moshi"
-
-# Explicitly override ~/.cache
-os.environ["HOME"] = "/tmp/home"
-os.makedirs("/tmp/home", exist_ok=True)
-
 pdf = "1706.03762v7.pdf"
 
 
@@ -173,8 +154,7 @@ def generate_audio_from_script(script, output_file="podcast_audio.wav"):
             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")
+            pipeline = KPipeline(lang_code="a")
             generator = pipeline(dialogue, voice=chosen_voice)
 
             segment_audio = []
@@ -206,67 +186,6 @@ def generate_audio_from_script(script, output_file="podcast_audio.wav"):
         print(f"Error processing transcript: {e}")
         return
 
-def generate_audio_kyutai(script, speaker1_voice=None, speaker2_voice=None, output_file="kyutai_audio.wav"):
-    if TTSModel is None:
-        print("Moshi is not installed.")
-        return None
-
-    try:
-        print(f"[INFO] Requested Kyutai voices: {speaker1_voice=}, {speaker2_voice=}")
-        # Reject absolute/local paths
-        if os.path.isabs(speaker1_voice) or os.path.isfile(speaker1_voice):
-            raise ValueError(f"❌ Invalid voice path for speaker1: {speaker1_voice}")
-        if os.path.isabs(speaker2_voice) or os.path.isfile(speaker2_voice):
-            raise ValueError(f"❌ Invalid voice path for speaker2: {speaker2_voice}")
-
-        transcript_list = ast.literal_eval(script)
-
-        # Load TTS model
-        checkpoint_info = CheckpointInfo.from_hf_repo(DEFAULT_DSM_TTS_REPO)
-        tts_model = TTSModel.from_checkpoint_info(
-            checkpoint_info,
-            n_q=32,
-            temp=0.6,
-            device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        )
-
-        # Use voice names directly from dropdown
-        print("[INFO] Resolving voice paths...")
-
-        start = time.time()
-        voice1_path = tts_model.get_voice_path(speaker1_voice)
-        print(f"[INFO] Got voice1_path in {time.time() - start:.2f}s")
-
-        start = time.time()
-        voice2_path = tts_model.get_voice_path(speaker2_voice)
-        print(f"[INFO] Got voice2_path in {time.time() - start:.2f}s")
-
-        texts = [dialogue for _, dialogue in transcript_list]
-        entries = tts_model.prepare_script(texts, padding_between=1)
-
-        condition_attributes = tts_model.make_condition_attributes([voice1_path, voice2_path], cfg_coef=2.0)
-
-        pcms = []
-        def _on_frame(frame):
-            if (frame != -1).all():
-                pcm = tts_model.mimi.decode(frame[:, 1:, :]).cpu().numpy()
-                pcms.append(np.clip(pcm[0, 0], -1, 1))
-
-        with tts_model.mimi.streaming(1):
-            tts_model.generate([entries], [condition_attributes], on_frame=_on_frame)
-
-        if pcms:
-            audio = np.concatenate(pcms, axis=-1)
-            sf.write(output_file, audio, tts_model.mimi.sample_rate)
-            print(f"[SUCCESS] Audio saved to: {output_file}")
-            return output_file
-
-        print("[WARNING] No audio segments were produced.")
-        return None
-
-    except Exception as e:
-        print(f"[ERROR] Kyutai TTS error: {e}")
-        return None
 
 def generate_tts():
     pipeline = KPipeline(lang_code="a")
@@ -303,16 +222,15 @@ def generate_podcast_script(
     Set provider="openrouter" to use OpenRouter, otherwise uses OpenAI.
     """
     pdf_basename = os.path.splitext(os.path.basename(pdf_path))[0]
-    folder = os.path.join("/tmp", pdf_basename)
+    folder = os.path.join(os.getcwd(), pdf_basename)
     os.makedirs(folder, exist_ok=True)
 
     destination_pdf = os.path.join(folder, os.path.basename(pdf_path))
-    try:
+    if not os.path.exists(destination_pdf):
         shutil.copy(pdf_path, destination_pdf)
         print(f"Copied {pdf_path} to {destination_pdf}")
-    except PermissionError:
-        print(f"[WARNING] Cannot copy PDF to {destination_pdf}, using original path.")
-        destination_pdf = pdf_path  # fallback
+    else:
+        print(f"PDF already copied at {destination_pdf}")
 
     transcript_path = os.path.join(folder, output_file)
     # If transcript exists, load and return it without calling the API.