feat: kyutai moshi

gradio_app.py

@@ -1,281 +1,187 @@
-# filepath: /Users/udaylunawat/Downloads/Data-Science-Projects/NotebookLM_clone/gradio_app.py
 import os
 import tempfile
 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 ast
+import numpy as np
+import soundfile as sf
 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")
+import concurrent.futures
 
-# Define number of workers based on CPU cores
-NUM_WORKERS = multiprocessing.cpu_count()  # Gets total CPU cores
+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,
+)
+
+warnings.filterwarnings("ignore")
+NUM_WORKERS = multiprocessing.cpu_count()
 
 def process_segment(entry_and_voice_map):
-    entry, voice_map = entry_and_voice_map  # Unpack the tuple
+    entry, voice_map = entry_and_voice_map
     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
+    return np.concatenate([audio for _, _, audio in generator], axis=0) if generator else 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}
+    print("[DEBUG] Raw transcript string:")
+    print(script)
     
-    # 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
-
+    voice_map = {"Speaker 1": speaker1_voice, "Speaker 2": speaker2_voice}
     try:
         transcript_list = ast.literal_eval(script)
         if not isinstance(transcript_list, list):
             raise ValueError("Transcript is not a list")
 
-        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")
+        entries = [(entry, voice_map) for entry in transcript_list if isinstance(entry, tuple) and len(entry) == 2]
+        with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
+            results = [r for r in executor.map(process_segment, entries) if r is not None]
+        if not results:
             return None
-
-        # Add a pause between segments
         sample_rate = 24000
         pause = np.zeros(sample_rate, dtype=np.float32)
-        final_audio = all_audio_segments[0]
-        for seg in all_audio_segments[1:]:
+        final_audio = results[0]
+        for seg in results[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"Error processing transcript: {e}")
+        print(f"Transcript parse error: {e}")
         return None
 
-
-def process_pdf(pdf_file, speaker1_voice, speaker2_voice, provider, api_key, openrouter_base=None):
-    """Process the uploaded PDF file and generate audio"""
+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:
-    
-        # Set API configuration based on provider
-        if provider == "openai":
-            os.environ["OPENAI_API_KEY"] = api_key
+        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 ""
             os.environ["OPENROUTER_API_BASE"] = "https://api.openai.com/v1"
-        else:
-            os.environ["OPENAI_API_KEY"] = api_key
+        if provider in ["openrouter", "kyutai"]:
+            os.environ["OPENAI_API_KEY"] = openrouter_key or ""
             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
-            
-        # 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 "Error generating transcript", None
+        tmp_path = pdf_file.name
 
-        # 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
-        # )
+        script_provider = "openrouter" if provider == "kyutai" and openrouter_key else provider
+        transcript, _ = generate_podcast_script(pdf_file.name, provider=script_provider)
 
-        # 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 transcript is None:
+            return "Transcript generation failed: got None", None
+        if not transcript.strip().startswith("["):
+            return f"Malformed transcript:\n{transcript}", None
 
-    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
+        audio_path = os.path.join(os.path.dirname(tmp_path), f"audio_{os.path.basename(tmp_path).replace('.pdf', '.wav')}")
 
-    except Exception as e:
-        print(f"Error in process_pdf: {str(e)}")
-        return f"Error processing file: {str(e)}", None
+        if tts_engine == "kyutai":
+            result = generate_audio_kyutai(transcript, kyutai_voice1, kyutai_voice2, audio_path)
+        else:
+            with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
+                result = executor.submit(
+                    generate_audio_from_script_with_voices,
+                    transcript, speaker1_voice, speaker2_voice, audio_path
+                ).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"),
+    ]
 
 def create_gradio_app():
-    # Add CSS for better styling
-    css = """
-    .gradio-container {max-width: 900px !important}
-    """
-    
+    css = ".gradio-container {max-width: 900px !important}"
     with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
-        gr.Markdown(
-            """
-            # 📚 NotebookLM-Kokoro TTS App
-            Upload a PDF, choose voices, and generate conversational audio using Kokoro TTS.
-            """
-        )
-        
+        gr.Markdown("# 🎧 PDF to Podcast — NotebookLM + Kokoro/Kyutai")
+
         with gr.Row():
-            with gr.Column(scale=2):
-                pdf_input = gr.File(
-                    label="Upload PDF Document",
-                    file_types=[".pdf"],
-                    type="filepath"
+            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.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..."
-                )
-                audio_output = gr.Audio(
-                    label="Generated Audio",
-                    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
                 )
-        
-        # # 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
-            """
-        )
-    
+
+                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.
+        """)
+
     return app
 
 if __name__ == "__main__":
-    demo = create_gradio_app()
-    demo.queue().launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True,
-        debug=True,
-        pwa=True
-    )
+    create_gradio_app().queue().launch(server_name="0.0.0.0", server_port=7860, share=True, debug=True, pwa=True)

notebook_lm_kokoro.py

@@ -23,6 +23,14 @@ 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
@@ -154,7 +162,8 @@ 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}'...")
 
-            pipeline = KPipeline(lang_code="a")
+            # 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 = []
@@ -186,6 +195,67 @@ 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")

requirements.txt

@@ -5,4 +5,6 @@ PyPDF2
 numpy
 openai
 ipython
-gradio>=4.0.0
+gradio>=4.0.0
+moshi>=0.2.4
+sphn>=0.2.0