app.py

import gradio as gr
from pydub import AudioSegment
import edge_tts
import os
import asyncio
import uuid
import re
import time
import tempfile
from concurrent.futures import ThreadPoolExecutor
from typing import List, Tuple, Optional, Dict, Any
import math
from dataclasses import dataclass
import multiprocessing
import psutil
import concurrent.futures
import gc
from gradio.themes import Monochrome

class TimingManager:
    def __init__(self):
        self.current_time = 0
        self.segment_gap = 100  # ms gap between segments
    
    def get_timing(self, duration):
        start_time = self.current_time
        end_time = start_time + duration
        self.current_time = end_time + self.segment_gap
        return start_time, end_time

def get_audio_length(audio_file):
    audio = AudioSegment.from_file(audio_file)
    return len(audio) / 1000

def format_time_ms(milliseconds):
    seconds, ms = divmod(int(milliseconds), 1000)
    mins, secs = divmod(seconds, 60)
    hrs, mins = divmod(mins, 60)
    return f"{hrs:02}:{mins:02}:{secs:02},{ms:03}"

@dataclass
class Segment:
    id: int
    text: str
    start_time: int = 0
    end_time: int = 0
    duration: int = 0
    audio: Optional[AudioSegment] = None
    lines: List[str] = None  # Add lines field for display purposes only

class TextProcessor:
    def __init__(self, words_per_line: int, lines_per_segment: int):
        self.words_per_line = words_per_line
        self.lines_per_segment = lines_per_segment
        self.min_segment_words = 3
        self.max_segment_words = words_per_line * lines_per_segment * 1.5  # Allow 50% more for natural breaks
        self.punctuation_weights = {
            '.': 1.0,  # Strong break
            '!': 1.0,
            '?': 1.0,
            ';': 0.8,  # Medium-strong break
            ':': 0.7,
            ',': 0.5,  # Medium break
            '-': 0.3,  # Weak break
            '(': 0.2,
            ')': 0.2
        }
    
    def analyze_sentence_complexity(self, text: str) -> float:
        """Analyze sentence complexity to determine optimal segment length"""
        words = text.split()
        complexity = 1.0
        
        # Adjust for sentence length
        if len(words) > self.words_per_line * 2:
            complexity *= 1.2
        
        # Adjust for punctuation density
        punct_count = sum(text.count(p) for p in self.punctuation_weights.keys())
        complexity *= (1 + (punct_count / len(words)) * 0.5)
        
        return complexity

    def find_natural_breaks(self, text: str) -> List[Tuple[int, float]]:
        """Find natural break points with their weights"""
        breaks = []
        words = text.split()
        
        for i, word in enumerate(words):
            weight = 0
            
            # Check for punctuation
            for punct, punct_weight in self.punctuation_weights.items():
                if word.endswith(punct):
                    weight = max(weight, punct_weight)
            
            # Check for natural phrase boundaries
            phrase_starters = {'however', 'therefore', 'moreover', 'furthermore', 'meanwhile', 'although', 'because'}
            if i < len(words) - 1 and words[i+1].lower() in phrase_starters:
                weight = max(weight, 0.6)
            
            # Check for conjunctions at natural points
            if i > self.min_segment_words:
                conjunctions = {'and', 'but', 'or', 'nor', 'for', 'yet', 'so'}
                if word.lower() in conjunctions:
                    weight = max(weight, 0.4)
            
            if weight > 0:
                breaks.append((i, weight))
        
        return breaks

    def split_into_segments(self, text: str) -> List[Segment]:
        # Normalize text and add proper spacing around punctuation
        text = re.sub(r'\s+', ' ', text.strip())
        text = re.sub(r'([.!?,;:])\s*', r'\1 ', text)
        text = re.sub(r'\s+([.!?,;:])', r'\1', text)
        
        # First, split into major segments by strong punctuation
        segments = []
        current_segment = []
        current_text = ""
        words = text.split()
        
        i = 0
        while i < len(words):
            complexity = self.analyze_sentence_complexity(' '.join(words[i:i + self.words_per_line * 2]))
            breaks = self.find_natural_breaks(' '.join(words[i:i + int(self.max_segment_words * complexity)]))
            
            # Find best break point
            best_break = None
            best_weight = 0
            
            for break_idx, weight in breaks:
                actual_idx = i + break_idx
                if (actual_idx - i >= self.min_segment_words and 
                    actual_idx - i <= self.max_segment_words):
                    if weight > best_weight:
                        best_break = break_idx
                        best_weight = weight
            
            if best_break is None:
                # If no good break found, use maximum length
                best_break = min(self.words_per_line * self.lines_per_segment, len(words) - i)
            
            # Create segment
            segment_words = words[i:i + best_break + 1]
            segment_text = ' '.join(segment_words)
            
            # Split segment into lines
            lines = self.split_into_lines(segment_text)
            final_segment_text = '\n'.join(lines)
            
            segments.append(Segment(
                id=len(segments) + 1,
                text=final_segment_text
            ))
            
            i += best_break + 1
        
        return segments

    def split_into_lines(self, text: str) -> List[str]:
        """Split segment text into natural lines"""
        words = text.split()
        lines = []
        current_line = []
        word_count = 0
        
        for word in words:
            current_line.append(word)
            word_count += 1
            
            # Check for natural line breaks
            is_break = (
                word_count >= self.words_per_line or
                any(word.endswith(p) for p in '.!?') or
                (word_count >= self.words_per_line * 0.7 and
                 any(word.endswith(p) for p in ',;:'))
            )
            
            if is_break:
                lines.append(' '.join(current_line))
                current_line = []
                word_count = 0
        
        if current_line:
            lines.append(' '.join(current_line))
        
        return lines

# IMPROVEMENT 1: Enhanced Error Handling
class TTSError(Exception):
    """Custom exception for TTS processing errors"""
    pass

class ResourceOptimizer:
    @staticmethod
    def get_optimal_workers():
        cpu_count = multiprocessing.cpu_count()
        return max(cpu_count - 1, 1)  # Leave one core for system
    
    @staticmethod
    def get_memory_limit():
        # Use up to 70% of available RAM
        return int(psutil.virtual_memory().available * 0.7)
    
    @staticmethod
    def get_batch_size(total_segments):
        # Calculate optimal batch size based on CPU cores
        return min(total_segments, ResourceOptimizer.get_optimal_workers() * 2)

async def process_segment_with_timing(segment: Segment, voice: str, rate: str, pitch: str) -> Segment:
    """Process a complete segment as a single TTS unit with improved error handling"""
    # Pre-allocate memory for audio processing
    gc.collect()  # Force garbage collection before processing
    
    audio_file = os.path.join(tempfile.gettempdir(), f"temp_segment_{segment.id}_{uuid.uuid4()}.wav")
    try:
        # Process the entire segment text as one unit, replacing newlines with spaces
        segment_text = ' '.join(segment.text.split('\n'))
        tts = edge_tts.Communicate(segment_text, voice, rate=rate, pitch=pitch)
        
        try:
            await tts.save(audio_file)
        except Exception as e:
            raise TTSError(f"Failed to generate audio for segment {segment.id}: {str(e)}")
        
        if not os.path.exists(audio_file) or os.path.getsize(audio_file) == 0:
            raise TTSError(f"Generated audio file is empty or missing for segment {segment.id}")
        
        try:
            segment.audio = AudioSegment.from_file(audio_file)
            # Optimize memory usage for audio processing
            segment.audio = segment.audio.set_channels(1)  # Convert to mono for memory efficiency
            silence = AudioSegment.silent(duration=30)
            segment.audio = silence + segment.audio + silence
            segment.duration = len(segment.audio)
        except Exception as e:
            raise TTSError(f"Failed to process audio file for segment {segment.id}: {str(e)}")
        
        return segment
    finally:
        if os.path.exists(audio_file):
            try:
                os.remove(audio_file)
            except Exception:
                pass

# IMPROVEMENT 2: Better File Management with cleanup
class FileManager:
    """Manages temporary and output files with cleanup capabilities"""
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp(prefix="tts_app_")
        self.output_files = []
        self.max_files_to_keep = 5  # Keep only the 5 most recent output pairs
        
    def get_temp_path(self, prefix):
        """Get a path for a temporary file"""
        return os.path.join(self.temp_dir, f"{prefix}_{uuid.uuid4()}")
    
    def create_output_paths(self):
        """Create paths for output files"""
        unique_id = str(uuid.uuid4())
        audio_path = os.path.join(self.temp_dir, f"final_audio_{unique_id}.mp3")
        srt_path = os.path.join(self.temp_dir, f"final_subtitles_{unique_id}.srt")
        
        self.output_files.append((srt_path, audio_path))
        self.cleanup_old_files()
        
        return srt_path, audio_path
    
    def cleanup_old_files(self):
        """Clean up old output files, keeping only the most recent ones"""
        if len(self.output_files) > self.max_files_to_keep:
            old_files = self.output_files[:-self.max_files_to_keep]
            for srt_path, audio_path in old_files:
                try:
                    if os.path.exists(srt_path):
                        os.remove(srt_path)
                    if os.path.exists(audio_path):
                        os.remove(audio_path)
                except Exception:
                    pass  # Ignore deletion errors
            
            # Update the list to only include files we're keeping
            self.output_files = self.output_files[-self.max_files_to_keep:]
    
    def cleanup_all(self):
        """Clean up all managed files"""
        for srt_path, audio_path in self.output_files:
            try:
                if os.path.exists(srt_path):
                    os.remove(srt_path)
                if os.path.exists(audio_path):
                    os.remove(audio_path)
            except Exception:
                pass  # Ignore deletion errors
        
        try:
            os.rmdir(self.temp_dir)
        except Exception:
            pass  # Ignore if directory isn't empty or can't be removed

# Create global file manager
file_manager = FileManager()

# IMPROVEMENT 3: Parallel Processing for Segments
async def generate_accurate_srt(
    text: str, 
    voice: str, 
    rate: str, 
    pitch: str, 
    words_per_line: int, 
    lines_per_segment: int,
    progress_callback=None,
    parallel: bool = True,
    max_workers: Optional[int] = None
) -> Tuple[str, str]:
    """Generate accurate SRT with optimized resource utilization"""
    processor = TextProcessor(words_per_line, lines_per_segment)
    segments = processor.split_into_segments(text)
    total_segments = len(segments)
    
    # Optimize worker count based on system resources
    if max_workers is None:
        max_workers = ResourceOptimizer.get_optimal_workers()
    
    if parallel and total_segments > 1:
        # Enhanced parallel processing with resource optimization
        batch_size = ResourceOptimizer.get_batch_size(total_segments)
        semaphore = asyncio.Semaphore(max_workers)
        processed_segments = []
        processed_count = 0
        
        # Process in batches for better resource utilization
        for i in range(0, total_segments, batch_size):
            batch = segments[i:i + batch_size]
            batch_tasks = []
            
            for segment in batch:
                batch_tasks.append(
                    process_with_semaphore(segment, voice, rate, pitch, semaphore)
                )
            
            # Process batch with maximum resource utilization
            batch_results = await asyncio.gather(*batch_tasks)
            processed_segments.extend(batch_results)
            
            # Force garbage collection between batches
            gc.collect()
            
            if progress_callback:
                processed_count += len(batch)
                progress = 0.1 + (0.8 * processed_count / total_segments)
                progress_callback(progress, f"Processed {processed_count}/{total_segments} segments")
    else:
        # Process segments sequentially (original method)
        for i, segment in enumerate(segments):
            try:
                processed_segment = await process_segment_with_timing(segment, voice, rate, pitch)
                processed_segments.append(processed_segment)
                
                if progress_callback:
                    progress = 0.1 + (0.8 * (i + 1) / total_segments)
                    progress_callback(progress, f"Processed {i + 1}/{total_segments} segments")
            except Exception as e:
                if progress_callback:
                    progress_callback(0.9, f"Error processing segment {segment.id}: {str(e)}")
                raise TTSError(f"Failed to process segment {segment.id}: {str(e)}")
    
    # Sort segments by ID to ensure correct order
    processed_segments.sort(key=lambda s: s.id)
    
    if progress_callback:
        progress_callback(0.9, "Finalizing audio and subtitles")
    
    # Now combine the segments in the correct order
    current_time = 0
    final_audio = AudioSegment.empty()
    srt_content = ""
    
    for segment in processed_segments:
        # Calculate precise timing
        segment.start_time = current_time
        segment.end_time = current_time + segment.duration
        
        # Add to SRT with precise timing
        srt_content += (
            f"{segment.id}\n"
            f"{format_time_ms(segment.start_time)} --> {format_time_ms(segment.end_time)}\n"
            f"{segment.text}\n\n"
        )
        
        # Add to final audio with precise positioning
        final_audio = final_audio.append(segment.audio, crossfade=0)
        
        # Update timing with precise gap
        current_time = segment.end_time
    
    # Export with high precision
    srt_path, audio_path = file_manager.create_output_paths()
    
    try:
        # Export with optimized quality settings and compression
        export_params = {
            'format': 'mp3',
            'bitrate': '192k',  # Reduced from 320k but still high quality
            'parameters': [
                '-ar', '44100',  # Standard sample rate
                '-ac', '2',      # Stereo
                '-compression_level', '0',  # Best compression
                '-qscale:a', '2'  # High quality VBR encoding
            ]
        }
        final_audio.export(audio_path, **export_params)
        
        with open(srt_path, "w", encoding='utf-8') as f:
            f.write(srt_content)
    except Exception as e:
        if progress_callback:
            progress_callback(1.0, f"Error exporting final files: {str(e)}")
        raise TTSError(f"Failed to export final files: {str(e)}")
    
    if progress_callback:
        progress_callback(1.0, "Complete!")
    
    return srt_path, audio_path

async def process_with_semaphore(segment, voice, rate, pitch, semaphore):
    async with semaphore:
        return await process_segment_with_timing(segment, voice, rate, pitch)

# IMPROVEMENT 4: Progress Reporting with proper error handling for older Gradio versions
async def process_text_with_progress(
    text, 
    pitch, 
    rate, 
    voice, 
    words_per_line, 
    lines_per_segment, 
    parallel_processing,
    progress=gr.Progress()
):
    # Input validation
    if not text or text.strip() == "":
        return None, None, None, True, "Please enter some text to convert to speech."
    
    # Format pitch and rate strings
    pitch_str = f"{pitch:+d}Hz" if pitch != 0 else "+0Hz"
    rate_str = f"{rate:+d}%" if rate != 0 else "+0%"
    
    try:
        # Start progress tracking
        progress(0, "Preparing text...")
        
        def update_progress(value, status):
            progress(value, status)
        
        srt_path, audio_path = await generate_accurate_srt(
            text,
            voice_options[voice],
            rate_str,
            pitch_str,
            words_per_line,
            lines_per_segment,
            progress_callback=update_progress,
            parallel=parallel_processing
        )
        
        # If successful, return results and hide error
        return srt_path, audio_path, audio_path, False, ""
    except TTSError as e:
        # Return specific TTS error
        return None, None, None, True, f"TTS Error: {str(e)}"
    except Exception as e:
        # Return any other error
        return None, None, None, True, f"Unexpected error: {str(e)}"

# Voice options dictionary
voice_options = {
    "Andrew Male": "en-US-AndrewNeural",
    "Jenny Female": "en-US-JennyNeural",
    "Guy Male": "en-US-GuyNeural",
    "Ana Female": "en-US-AnaNeural",
    "Aria Female": "en-US-AriaNeural",
    "Brian Male": "en-US-BrianNeural",
    "Christopher Male": "en-US-ChristopherNeural",
    "Eric Male": "en-US-EricNeural",
    "Michelle Male": "en-US-MichelleNeural",
    "Roger Male": "en-US-RogerNeural",
    "Natasha Female": "en-AU-NatashaNeural",
    "William Male": "en-AU-WilliamNeural",
    "Clara Female": "en-CA-ClaraNeural",
    "Liam Female ": "en-CA-LiamNeural",
    "Libby Female": "en-GB-LibbyNeural",
    "Maisie": "en-GB-MaisieNeural",
    "Ryan": "en-GB-RyanNeural",
    "Sonia": "en-GB-SoniaNeural",
    "Thomas": "en-GB-ThomasNeural",
    "Sam": "en-HK-SamNeural",
    "Yan": "en-HK-YanNeural",
    "Connor": "en-IE-ConnorNeural",
    "Emily": "en-IE-EmilyNeural",
    "Neerja": "en-IN-NeerjaNeural",
    "Prabhat": "en-IN-PrabhatNeural",
    "Asilia": "en-KE-AsiliaNeural",
    "Chilemba": "en-KE-ChilembaNeural",
    "Abeo": "en-NG-AbeoNeural",
    "Ezinne": "en-NG-EzinneNeural",
    "Mitchell": "en-NZ-MitchellNeural",
    "James": "en-PH-JamesNeural",
    "Rosa": "en-PH-RosaNeural",
    "Luna": "en-SG-LunaNeural",
    "Wayne": "en-SG-WayneNeural",
    "Elimu": "en-TZ-ElimuNeural",
    "Imani": "en-TZ-ImaniNeural",
    "Leah": "en-ZA-LeahNeural",
    "Luke": "en-ZA-LukeNeural"
    # Add other voices as needed
}

# Register cleanup on exit
import atexit
atexit.register(file_manager.cleanup_all)

# Create custom theme
theme = gr.themes.Monochrome(
    primary_hue="blue",
    secondary_hue="slate",
    neutral_hue="zinc",
    radius_size=gr.themes.sizes.radius_sm,
    font=("Inter", "system-ui", "sans-serif"),
    font_mono=("IBM Plex Mono", "monospace")
)

# Create Gradio interface with modern UI
with gr.Blocks(
    title="Text to Speech Studio", 
    theme=theme,
    css="""
        .container { max-width: 1200px; margin: auto; padding: 2rem; }
        .title { text-align: center; margin-bottom: 2.5rem; }
        .title h1 { font-size: 2.5rem; font-weight: 700; margin-bottom: 0.5rem; }
        .title h3 { font-size: 1.2rem; font-weight: 400; opacity: 0.8; }
        .input-group { margin-bottom: 1.5rem; border-radius: 8px; }
        .help-text { font-size: 0.9rem; opacity: 0.8; padding: 0.5rem 0; }
        .status-area { margin: 1.5rem 0; padding: 1rem; border-radius: 8px; }
        .error-message { color: #dc2626; }
        .preview-audio { margin: 1rem 0; }
        .download-file { padding: 1rem; }
        button.primary { transform: scale(1); transition: transform 0.2s; }
        button.primary:hover { transform: scale(1.02); }
        button.secondary:hover { opacity: 0.9; }
    """
) as app:
    with gr.Group(elem_classes="container"):
        gr.Markdown(
            """
            # 🎙️ Text to Speech Studio
            ### Generate professional quality audio with synchronized subtitles
            """
        , elem_classes="title")
        
        with gr.Tabs():
            with gr.TabItem("📝 Text Input"):
                with gr.Row():
                    with gr.Column(scale=3):
                        text_input = gr.Textbox(
                            label="Your Text",
                            lines=10,
                            placeholder="Enter your text here. The AI will automatically segment it into natural phrases...",
                            elem_classes="input-group"
                        )
                        gr.Markdown(
                            "💡 **Tip:** For best results, ensure proper punctuation in your text.",
                            elem_classes="help-text"
                        )
                    
                    with gr.Column(scale=2):
                        with gr.Group():
                            gr.Markdown("### Voice Settings")
                            voice_dropdown = gr.Dropdown(
                                label="Voice",
                                choices=list(voice_options.keys()),
                                value="Jenny Female",
                                elem_classes="input-group"
                            )
                            with gr.Row():
                                with gr.Column():
                                    pitch_slider = gr.Slider(
                                        label="Pitch",
                                        minimum=-10,
                                        maximum=10,
                                        value=0,
                                        step=1,
                                        elem_classes="input-group"
                                    )
                                with gr.Column():
                                    rate_slider = gr.Slider(
                                        label="Speed",
                                        minimum=-25,
                                        maximum=25,
                                        value=0,
                                        step=1,
                                        elem_classes="input-group"
                                    )
            
            with gr.TabItem("⚙️ Advanced Settings"):
                with gr.Row():
                    with gr.Column():
                        words_per_line = gr.Slider(
                            label="Words per Line",
                            minimum=3,
                            maximum=12,
                            value=6,
                            step=1,
                            info="📝 Controls subtitle line length",
                            elem_classes="input-group"
                        )
                    with gr.Column():
                        lines_per_segment = gr.Slider(
                            label="Lines per Segment",
                            minimum=1,
                            maximum=4,
                            value=2,
                            step=1,
                            info="📑 Controls subtitle block size",
                            elem_classes="input-group"
                        )
                    with gr.Column():
                        parallel_processing = gr.Checkbox(
                            label="Parallel Processing",
                            value=True,
                            info="⚡ Faster processing for longer texts",
                            elem_classes="input-group"
                        )
        
        with gr.Row():
            with gr.Column(scale=2):
                submit_btn = gr.Button(
                    "🎯 Generate Audio & Subtitles",
                    variant="primary",
                    scale=2
                )
            with gr.Column():
                clear_btn = gr.Button("🔄 Clear All", variant="secondary")
        
        with gr.Group(elem_classes="status-area"):
            error_output = gr.Textbox(
                label="Status",
                visible=False,
                elem_classes="error-message"
            )
            
        with gr.Tabs():
            with gr.TabItem("🎧 Preview"):
                audio_output = gr.Audio(
                    label="Generated Audio",
                    elem_classes="preview-audio"
                )
            with gr.TabItem("📥 Downloads"):
                with gr.Row():
                    with gr.Column():
                        srt_file = gr.File(
                            label="📄 Subtitle File (SRT)",
                            elem_classes="download-file"
                        )
                    with gr.Column():
                        audio_file = gr.File(
                            label="🎵 Audio File (MP3)",
                            elem_classes="download-file"
                        )

        gr.Markdown(
            """
            ### 📌 Features
            - Professional-quality text-to-speech conversion
            - Automatic natural speech segmentation
            - Perfectly synchronized subtitles
            - Multiple voice options and customization
            """,
            elem_classes="help-text"
        )

    # Clear button functionality
    def clear_inputs():
        return {
            text_input: "",
            pitch_slider: 0,
            rate_slider: 0,
            voice_dropdown: "Jenny Female",
            words_per_line: 6,
            lines_per_segment: 2,
            parallel_processing: True,
            error_output: gr.update(visible=False),
            audio_output: None,
            srt_file: None,
            audio_file: None
        }
    
    clear_btn.click(
        fn=clear_inputs,
        inputs=[],
        outputs=[
            text_input, pitch_slider, rate_slider, voice_dropdown,
            words_per_line, lines_per_segment, parallel_processing,
            error_output, audio_output, srt_file, audio_file
        ]
    )

    # Existing button click handler
    submit_btn.click(
        fn=process_text_with_progress,
        inputs=[
            text_input, pitch_slider, rate_slider, voice_dropdown,
            words_per_line, lines_per_segment, parallel_processing
        ],
        outputs=[
            srt_file, audio_file, audio_output, error_output, error_output
        ],
        api_name="generate"
    )

if __name__ == "__main__":
    # Set process priority to high
    p = psutil.Process()
    try:
        p.nice(psutil.BELOW_NORMAL_PRIORITY_CLASS if os.name == 'nt' else 10)
    except Exception:
        pass
    app.launch()