Update app.py

app.py

@@ -1,5 +1,6 @@
 import nltk
 nltk.download('punkt')
+# Explicitly download 'punkt_tab' as it's often required by sent_tokenize
 nltk.download('punkt_tab')
 
 import random
@@ -11,35 +12,47 @@ import io
 import soundfile as sf
 from pydub import AudioSegment
 from nltk.tokenize import sent_tokenize
-import os
+import os # Added for temporary file handling
 
-# Determine the device
+# 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}")
+print(f"?? Running on device: {DEVICE}")
 
 # --- Global Model Initialization ---
 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."""
     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 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: Model failed to load: {e}")
+    print(f"CRITICAL: Failed to load model on startup. Application may not function. Error: {e}")
 
 def set_seed(seed: int):
+    """Sets the random seed for reproducibility across torch, numpy, and random."""
     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)
@@ -47,78 +60,190 @@ def set_seed(seed: int):
 def generate_tts_audio(
     text_input: str,
     audio_prompt_path_input: str = None,
-    exaggeration_input: float = 0.8,   # Defaulted for 'Elder' gravitas
-    temperature_input: float = 0.7,     # Defaulted for consistency
+    exaggeration_input: float = 0.5,
+    temperature_input: float = 0.8,
     seed_num_input: int = 0,
-    cfgw_input: float = 0.3,           # Defaulted for 'Slow' pace
-    pause_duration: float = 1.0        # NEW: Silence between chunks
-) -> str:
+    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,
     }
 
+    processed_audio_prompt_path = None
     if audio_prompt_path_input:
-        generate_kwargs["audio_prompt_path"] = audio_prompt_path_input
+        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)
-    
-    # Create the 'Silence' segment (in milliseconds)
-    silence_gap = AudioSegment.silent(duration=int(pause_duration * 1000))
+
+    # 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:
-        print(f"Processing: '{sentence[:30]}...'")
-        wav_tensor = current_model.generate(sentence, **generate_kwargs)
+        # 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()
 
-        # Convert to AudioSegment
         buffer = io.BytesIO()
         sf.write(buffer, wav_numpy, current_model.sr, format='WAV')
         buffer.seek(0)
         audio_segment = AudioSegment.from_wav(buffer)
-        
-        # Add the clip + the pause
         all_audio_segments.append(audio_segment)
-        all_audio_segments.append(silence_gap)
 
-    # Combine all
-    combined_audio = AudioSegment.empty()
-    for seg in all_audio_segments:
-        combined_audio += seg
+    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]
 
-    output_filename = "elderly_voice_output.mp3"
+    # Export to MP3 format
+    output_filename = "combined_chatterbox_output.mp3"
     combined_audio.export(output_filename, format="mp3")
-    return output_filename
 
-# --- Gradio Interface ---
-with gr.Blocks(title="Elderly TTS Studio") as demo:
-    gr.Markdown("# 👴 Elderly Wisdom TTS Generator")
-    
+    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(label="Script", lines=8, placeholder="Enter the story here...")
-            ref_wav = gr.Audio(type="filepath", label="Voice Reference (Upload a raspy/old voice)")
-            
-            with gr.Accordion("Character Pacing (Elder Settings)", open=True):
-                cfg_weight = gr.Slider(0.1, 1.0, step=.05, label="Pace (0.3 is slow/elderly)", value=0.3)
-                pause_len = gr.Slider(0, 3.0, step=.1, label="Breath Pause (Seconds between sentences)", value=1.0)
-                exaggeration = gr.Slider(0.2, 2.0, step=.05, label="Gravitas/Emotion", value=0.8)
-                
-            run_btn = gr.Button("Generate 8-Minute Audio", variant="primary")
-            
+            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():
-            audio_output = gr.Audio(label="Final Combined MP3")
+            # 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, gr.State(0.7), gr.State(0), cfg_weight, pause_len],
+        inputs=[
+            text,
+            ref_wav,
+            exaggeration,
+            temp,
+            seed_num,
+            cfg_weight,
+        ],
         outputs=[audio_output],
     )
 
-demo.launch(share=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)