Update app.py

app.py

@@ -1,249 +1,178 @@
 import nltk
-nltk.download('punkt')
-# Explicitly download 'punkt_tab' as it's often required by sent_tokenize
-nltk.download('punkt_tab')
+nltk.download("punkt")
 
 import random
 import numpy as np
 import torch
-from chatterbox.src.chatterbox.tts import ChatterboxTTS
-import gradio as gr
 import io
+import os
 import soundfile as sf
-from pydub import AudioSegment
 from nltk.tokenize import sent_tokenize
-import os # Added for temporary file handling
+from pydub import AudioSegment
+import gradio as gr
 
-# Determine the device to run on (GPU if available, otherwise CPU)
+from chatterbox.src.chatterbox.tts import ChatterboxTTS
+
+# ===============================
+# DEVICE
+# ===============================
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"?? Running on device: {DEVICE}")
+print(f"Running on: {DEVICE}")
 
-# --- Global Model Initialization ---
+# ===============================
+# LOAD MODEL ONCE
+# ===============================
 MODEL = None
 
-def get_or_load_model():
-    """Loads the ChatterboxTTS model if it hasn't been loaded already,
-    and ensures it's on the correct device. This helps avoid reloading
-    the model multiple times which can be slow."""
+def get_model():
     global MODEL
     if MODEL is None:
-        print("Model not loaded, initializing...")
-        try:
-            # Load the model and move it to the determined device (CPU or CUDA)
-            MODEL = ChatterboxTTS.from_pretrained(DEVICE)
-            # Ensure the model is explicitly on the correct device after loading
-            if hasattr(MODEL, 'to') and str(MODEL.device) != DEVICE:
-                MODEL.to(DEVICE)
-            print(f"Model loaded successfully. Internal device: {getattr(MODEL, 'device', 'N/A')}")
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            # Re-raise the exception to indicate a critical failure
-            raise
+        print("Loading Chatterbox model...")
+        MODEL = ChatterboxTTS.from_pretrained(DEVICE)
+        if hasattr(MODEL, "to"):
+            MODEL.to(DEVICE)
+        print("Model ready.")
     return MODEL
 
-# Attempt to load the model at startup of the script.
-# This ensures the model is ready when the Gradio interface starts.
-try:
-    get_or_load_model()
-except Exception as e:
-    print(f"CRITICAL: Failed to load model on startup. Application may not function. Error: {e}")
+get_model()
 
-def set_seed(seed: int):
-    """Sets the random seed for reproducibility across torch, numpy, and random."""
+# ===============================
+# SEED
+# ===============================
+def set_seed(seed):
     torch.manual_seed(seed)
     if DEVICE == "cuda":
-        torch.cuda.manual_seed(seed)
         torch.cuda.manual_seed_all(seed)
     random.seed(seed)
     np.random.seed(seed)
 
-def generate_tts_audio(
-    text_input: str,
-    audio_prompt_path_input: str = None,
-    exaggeration_input: float = 0.5,
-    temperature_input: float = 0.8,
-    seed_num_input: int = 0,
-    cfgw_input: float = 0.5
-) -> str: # Return type changed to str (filepath)
-    """
-    Generate high-quality speech audio from text using ChatterboxTTS model with optional reference audio styling.
-    Handles long scripts by chunking text, generating audio for each chunk, and combining them into an MP3.
-
-    Args:
-        text_input (str): The text to synthesize into speech.
-        audio_prompt_path_input (str, optional): File path or URL to the reference audio file that defines the target voice style. Defaults to None.
-        exaggeration_input (float, optional): Controls speech expressiveness (0.25-2.0, neutral=0.5, extreme values may be unstable). Defaults to 0.5.
-        temperature_input (float, optional): Controls randomness in generation (0.05-5.0, higher=more varied). Defaults to 0.8.
-        seed_num_input (int, optional): Random seed for reproducible results (0 for random generation). Defaults to 0.
-        cfgw_input (float, optional): CFG/Pace weight controlling generation guidance (0.2-1.0). Defaults to 0.5.
-    Returns:
-        str: Filepath to the generated combined MP3 audio waveform.
-    """
-    current_model = get_or_load_model()
-    if current_model is None:
-        raise RuntimeError("TTS model is not loaded. Please check the startup logs for errors.")
-
-    if seed_num_input != 0:
-        set_seed(int(seed_num_input))
-
-    print(f"Generating audio for text: '{text_input[:100]}...' (first 100 chars)")
-    print(f"Audio prompt path received: {audio_prompt_path_input}") # Debug print for the received path
-
-    generate_kwargs = {
-        "exaggeration": exaggeration_input,
-        "temperature": temperature_input,
-        "cfg_weight": cfgw_input,
+# ===============================
+# PODCAST SAFE SETTINGS
+# ===============================
+MAX_CHARS = 220                  # stable for chatterbox
+SILENCE_MS = 350                 # natural pause
+FADE_IN = 30
+FADE_OUT = 60
+
+# ===============================
+# MAIN TTS FUNCTION
+# ===============================
+def generate_tts(
+    text,
+    ref_audio=None,
+    exaggeration=0.4,
+    temperature=0.7,
+    seed=0,
+    cfg_weight=0.6,
+):
+
+    model = get_model()
+
+    if seed != 0:
+        set_seed(int(seed))
+
+    kwargs = {
+        "exaggeration": exaggeration,
+        "temperature": temperature,
+        "cfg_weight": cfg_weight,
     }
 
-    processed_audio_prompt_path = None
-    if audio_prompt_path_input:
+    # --------------------------------
+    # Handle reference voice
+    # --------------------------------
+    temp_prompt = None
+    if ref_audio:
         try:
-            # Load the input audio using pydub
-            audio = AudioSegment.from_file(audio_prompt_path_input)
-            # Create a temporary WAV file to ensure compatibility with ChatterboxTTS
-            temp_wav_path = "temp_prompt.wav"
-            audio.export(temp_wav_path, format="wav")
-            processed_audio_prompt_path = temp_wav_path
-            print(f"Converted audio prompt to temporary WAV: {processed_audio_prompt_path}")
-        except Exception as e:
-            print(f"Warning: Could not process audio prompt file '{audio_prompt_path_input}'. Error: {e}")
-            print("Proceeding without audio prompt (will use default voice).")
-            # If conversion fails, ensure the audio prompt path is not used
-            processed_audio_prompt_path = None
-
-    if processed_audio_prompt_path:
-        generate_kwargs["audio_prompt_path"] = processed_audio_prompt_path
-
-    all_audio_segments = []
-    # Split text into sentences for more natural chunking
-    sentences = sent_tokenize(text_input)
-
-    # Chatterbox model has an implicit input limit, typically around 300 characters.
-    # We'll chunk sentences to stay within this limit.
-    MAX_CHARS_PER_MODEL_INPUT = 300
-
-    current_chunk_sentences = []
-    current_chunk_char_count = 0
-
-    for sentence in sentences:
-        # If adding the current sentence exceeds the max chars, process the current chunk
-        # and ensure current_chunk_sentences is not empty to avoid creating empty chunks
-        if current_chunk_char_count + len(sentence) + 1 > MAX_CHARS_PER_MODEL_INPUT and current_chunk_sentences: # +1 for space
-            chunk_text = " ".join(current_chunk_sentences)
-            print(f"Processing chunk (chars: {len(chunk_text)}): '{chunk_text[:50]}...'")
-            wav_tensor = current_model.generate(chunk_text, **generate_kwargs)
-            wav_numpy = wav_tensor.squeeze(0).cpu().numpy()
-
-            # Convert numpy array to AudioSegment via an in-memory WAV buffer
-            buffer = io.BytesIO()
-            sf.write(buffer, wav_numpy, current_model.sr, format='WAV')
-            buffer.seek(0) # Rewind the buffer to the beginning
-            audio_segment = AudioSegment.from_wav(buffer)
-            all_audio_segments.append(audio_segment)
-
-            # Start a new chunk with the current sentence
-            current_chunk_sentences = [sentence]
-            current_chunk_char_count = len(sentence)
+            audio = AudioSegment.from_file(ref_audio)
+            temp_prompt = "voice_prompt.wav"
+            audio.export(temp_prompt, format="wav")
+            kwargs["audio_prompt_path"] = temp_prompt
+        except:
+            print("Reference audio failed — using default voice.")
+
+    # --------------------------------
+    # Sentence chunking
+    # --------------------------------
+    sentences = sent_tokenize(text)
+
+    chunks = []
+    current = ""
+
+    for s in sentences:
+        if len(current) + len(s) < MAX_CHARS:
+            current += " " + s
         else:
-            current_chunk_sentences.append(sentence)
-            # Add 1 for space between sentences, but only if it's not the very first sentence in a chunk
-            current_chunk_char_count += len(sentence) + (1 if current_chunk_sentences else 0)
+            chunks.append(current.strip())
+            current = s
+
+    if current.strip():
+        chunks.append(current.strip())
+
+    print(f"Total chunks: {len(chunks)}")
 
-    # Process the last remaining chunk
-    if current_chunk_sentences:
-        chunk_text = " ".join(current_chunk_sentences)
-        print(f"Processing final chunk (chars: {len(chunk_text)}): '{chunk_text[:50]}...'")
-        wav_tensor = current_model.generate(chunk_text, **generate_kwargs)
-        wav_numpy = wav_tensor.squeeze(0).cpu().numpy()
+    # --------------------------------
+    # Generate audio per chunk
+    # --------------------------------
+    final_audio = AudioSegment.empty()
+    silence = AudioSegment.silent(duration=SILENCE_MS)
+
+    for i, chunk in enumerate(chunks):
+        print(f"Generating chunk {i+1}/{len(chunks)}")
+
+        wav = model.generate(chunk, **kwargs)
+        wav_np = wav.squeeze(0).cpu().numpy()
 
         buffer = io.BytesIO()
-        sf.write(buffer, wav_numpy, current_model.sr, format='WAV')
+        sf.write(buffer, wav_np, model.sr, format="WAV")
         buffer.seek(0)
-        audio_segment = AudioSegment.from_wav(buffer)
-        all_audio_segments.append(audio_segment)
-
-    if not all_audio_segments:
-        raise ValueError("No audio segments were generated. Please ensure the input text is not empty or too short.")
 
-    # Concatenate all audio segments
-    combined_audio = all_audio_segments[0]
-    for i in range(1, len(all_audio_segments)):
-        combined_audio += all_audio_segments[i]
+        segment = AudioSegment.from_wav(buffer)
+        segment = segment.fade_in(FADE_IN).fade_out(FADE_OUT)
 
-    # Export to MP3 format
-    output_filename = "combined_chatterbox_output.mp3"
-    combined_audio.export(output_filename, format="mp3")
+        final_audio += segment + silence
 
-    print(f"Combined audio generated and saved as {output_filename}")
+    # --------------------------------
+    # Export
+    # --------------------------------
+    output_path = "story_voice.mp3"
+    final_audio.export(output_path, format="mp3", bitrate="192k")
 
-    # Clean up the temporary WAV file if it was created
-    if processed_audio_prompt_path and os.path.exists(processed_audio_prompt_path):
-        os.remove(processed_audio_prompt_path)
-        print(f"Cleaned up temporary prompt file: {processed_audio_prompt_path}")
+    if temp_prompt and os.path.exists(temp_prompt):
+        os.remove(temp_prompt)
 
-    return output_filename # Return the filepath for Gradio
+    return output_path
 
-# --- Gradio Interface Definition ---
+# ===============================
+# GRADIO UI
+# ===============================
 with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # Chatterbox TTS Demo
-        Generate high-quality speech from text with reference audio styling.
-        Now supports longer scripts and MP3 output!
-        """
+    gr.Markdown("## 🎙️ Storyteller / Podcast Chatterbox TTS")
+
+    text = gr.Textbox(
+        label="Story Text",
+        lines=12,
+        placeholder="Paste your full story here..."
+    )
+
+    ref = gr.Audio(
+        sources=["upload", "microphone"],
+        type="filepath",
+        label="Reference Voice (optional)"
     )
-    with gr.Row():
-        with gr.Column():
-            text = gr.Textbox(
-                value="""
-                The quick brown fox jumps over the lazy dog. This is a common pangram used to display all letters of the alphabet.
-                Now, let's try a slightly longer passage to test the new chunking functionality.
-                This should demonstrate how the system handles multiple sentences and combines them seamlessly.
-                We are aiming for a natural flow, even with extended input.
-                The sun dipped below the horizon, painting the sky in hues of orange and purple.
-                A gentle breeze rustled through the leaves, carrying the scent of night-blooming jasmine.
-                Soon, the stars would emerge, tiny pinpricks of light in the vast, dark canvas above.
-                """,
-                label="Text to synthesize (can be long)",
-                lines=10 # Increased lines for longer text input
-            )
-            # Gradio's Audio component handles file uploads directly.
-            # The 'value' here is a placeholder for the demo.
-            ref_wav = gr.Audio(
-                sources=["upload", "microphone"],
-                type="filepath",
-                label="Reference Audio File (Optional)",
-                value="https://storage.googleapis.com/chatterbox-demo-samples/prompts/female_shadowheart4.flac" # Default example audio
-            )
-            exaggeration = gr.Slider(
-                0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5, extreme values can be unstable)", value=.5
-            )
-            cfg_weight = gr.Slider(
-                0.2, 1, step=.05, label="CFG/Pace", value=0.5
-            )
-            with gr.Accordion("More options", open=False):
-                seed_num = gr.Number(value=0, label="Random seed (0 for random)")
-                temp = gr.Slider(0.05, 5, step=.05, label="Temperature", value=.8)
-            run_btn = gr.Button("Generate", variant="primary")
-        with gr.Column():
-            # Output type is now implicitly a filepath to the MP3
-            audio_output = gr.Audio(label="Output Audio (MP3)")
-
-    # Define the action when the "Generate" button is clicked
-    run_btn.click(
-        fn=generate_tts_audio,
-        inputs=[
-            text,
-            ref_wav,
-            exaggeration,
-            temp,
-            seed_num,
-            cfg_weight,
-        ],
-        outputs=[audio_output],
+
+    exaggeration = gr.Slider(0.25, 1.0, value=0.4, step=0.05, label="Emotion")
+    temperature = gr.Slider(0.3, 1.2, value=0.7, step=0.05, label="Variation")
+    cfg = gr.Slider(0.3, 1.0, value=0.6, step=0.05, label="Voice Stability")
+
+    seed = gr.Number(value=0, label="Seed (0 = random)")
+
+    btn = gr.Button("Generate Voice")
+    out = gr.Audio(label="Final Audio")
+
+    btn.click(
+        fn=generate_tts,
+        inputs=[text, ref, exaggeration, temperature, seed, cfg],
+        outputs=out
     )
 
-# Launch the Gradio interface.
-# Use share=True to get a public URL for the app, essential for Colab.
-# debug=True can be useful for seeing more detailed error messages in the Colab output.
-demo.launch(share=True, debug=True)
+demo.launch(share=True)