Update app.py

app.py

@@ -1,10 +1,22 @@
+ 
+# Download NLTK data for sentence tokenization
+import nltk
+nltk.download('punkt')
+# Explicitly download 'punkt_tab' as it's often required by sent_tokenize
+nltk.download('punkt_tab')
+
 import random
 import numpy as np
 import torch
-from chatterbox.src.chatterbox.tts import ChatterboxTTS
+from chatterbox.tts import ChatterboxTTS
 import gradio as gr
-import spaces
+import io
+import soundfile as sf
+from pydub import AudioSegment
+from nltk.tokenize import sent_tokenize
+import os # Added for temporary file handling
 
+# Determine the device to run on (GPU if available, otherwise CPU)
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"🚀 Running on device: {DEVICE}")
 
@@ -13,21 +25,26 @@ 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."""
+    and ensures it's on the correct device. This helps avoid reloading
+    the model multiple times which can be slow."""
     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
     return MODEL
 
-# Attempt to load the model at startup.
+# 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:
@@ -42,37 +59,6 @@ def set_seed(seed: int):
     random.seed(seed)
     np.random.seed(seed)
 
-import re
-
-def split_text_smart(text, max_length=2000):
-    """Split text into smart chunks without breaking words mid-way (adapted from Kokoro)."""
-    chunks = []
-    start = 0
-    text_length = len(text)
-
-    while start < text_length:
-        if (text_length - start) <= max_length:
-            chunks.append(text[start:].strip())
-            break
-
-        split_pos = start + max_length
-        split_region = text[start:split_pos]
-
-        # Find whitespace or newlines to split at
-        matches = list(re.finditer(r'[\s\n]', split_region))
-        if matches:
-            split_at = matches[-1].start()
-            split_point = start + split_at
-        else:
-            split_point = split_pos
-
-        chunk = text[start:split_point].strip()
-        chunks.append(chunk)
-        start = split_point + 1
-
-    return chunks
-
-@spaces.GPU
 def generate_tts_audio(
     text_input: str,
     audio_prompt_path_input: str = None,
@@ -80,110 +66,156 @@ def generate_tts_audio(
     temperature_input: float = 0.8,
     seed_num_input: int = 0,
     cfgw_input: float = 0.5
-) -> tuple[int, np.ndarray]:
+) -> str: # Return type changed to str (filepath)
     """
     Generate high-quality speech audio from text using ChatterboxTTS model with optional reference audio styling.
-    Uses smart text chunking to prevent CUDA errors on long texts.
-    
-    This tool synthesizes natural-sounding speech from input text. When a reference audio file 
-    is provided, it captures the speaker's voice characteristics and speaking style. The generated audio 
-    maintains the prosody, tone, and vocal qualities of the reference speaker, or uses default voice if no reference is provided.
+    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
+        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:
-        tuple[int, np.ndarray]: A tuple containing the sample rate (int) and the generated audio waveform (numpy.ndarray)
+        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.")
+        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[:50]}...' (Length: {len(text_input)} chars)")
-    
-    # Handle optional audio prompt
+    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,
     }
-    
+
+    processed_audio_prompt_path = None
     if audio_prompt_path_input:
-        generate_kwargs["audio_prompt_path"] = audio_prompt_path_input
-    
-    # Force chunking for any text over 1000 chars to prevent CUDA errors
-    if len(text_input) > 1000:
-        chunks = split_text_smart(text_input, max_length=1000)
-        print(f"Text chunked into {len(chunks)} smart segments (forced chunking to prevent CUDA errors)")
-    else:
-        chunks = [text_input]
-        print("Processing as single chunk")
-    
-    # Generate audio for each chunk
-    all_audio_segments = []
-    sample_rate = None
-    
-    for i, chunk in enumerate(chunks):
-        print(f"Processing chunk {i+1}/{len(chunks)}: {len(chunk)} chars - '{chunk[:30]}...'")
-        
         try:
-            wav = current_model.generate(
-                chunk,
-                **generate_kwargs
-            )
-            
-            if sample_rate is None:
-                sample_rate = current_model.sr
-                
-            all_audio_segments.append(wav.squeeze(0).numpy())
-            
-        except RuntimeError as e:
-            if "CUDA" in str(e) or "device-side assert" in str(e):
-                print(f"CUDA error on chunk {i+1}, skipping: {e}")
-                continue
-            else:
-                raise e
-    
+            # 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)
+        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)
+
+    # 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()
+
+        buffer = io.BytesIO()
+        sf.write(buffer, wav_numpy, current_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 RuntimeError("All chunks failed to generate - text may be too long or contain invalid characters")
-    
-    # Concatenate all audio segments (no artificial pauses - let natural speech flow)
-    if len(all_audio_segments) == 1:
-        final_audio = all_audio_segments[0]
-    else:
-        final_audio = np.concatenate(all_audio_segments)
-    
-    print("Audio generation complete.")
-    return (sample_rate, final_audio)
+        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]
+
+    # Export to MP3 format
+    output_filename = "combined_chatterbox_output.mp3"
+    combined_audio.export(output_filename, format="mp3")
 
+    print(f"Combined audio generated and saved as {output_filename}")
+
+    # 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}")
+
+    return output_filename # Return the filepath for Gradio
+
+# --- Gradio Interface Definition ---
 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!
         """
     )
     with gr.Row():
         with gr.Column():
             text = gr.Textbox(
-                value="Now let's make my mum's favourite. So three mars bars into the pan. Then we add the tuna and just stir for a bit, just let the chocolate and fish infuse. A sprinkle of olive oil and some tomato ketchup. Now smell that. Oh boy this is going to be incredible.",
-                label="Text to synthesize",  # Removed "max chars 300" from label
-                max_lines=5
+                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"
+                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
@@ -191,16 +223,15 @@ with gr.Blocks() as demo:
             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():
-            audio_output = gr.Audio(label="Output Audio")
+            # 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=[
@@ -213,5 +244,8 @@ with gr.Blocks() as demo:
         ],
         outputs=[audio_output],
     )
- 
-demo.launch(share=True, mcp_server=True)
+
+# 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)