app.py

import gradio as gr
from pydub import AudioSegment
import edge_tts
import os
import asyncio
import uuid
import re

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}"

def smart_text_split(text, words_per_line, lines_per_segment):
    # First split by major punctuation (periods, exclamation marks, question marks)
    sentences = re.split(r'([.!?]+)', text)
    
    # Recombine sentences with their punctuation
    sentences = [''.join(i) for i in zip(sentences[::2], sentences[1::2] + [''])]
    
    segments = []
    current_segment = []
    current_line = []
    
    for sentence in sentences:
        # Split sentence into words
        words = sentence.strip().split()
        
        for word in words:
            current_line.append(word)
            
            # Check if current line has reached words_per_line
            if len(current_line) >= words_per_line:
                current_segment.append(' '.join(current_line))
                current_line = []
                
                # Check if current segment has reached lines_per_segment
                if len(current_segment) >= lines_per_segment:
                    segments.append('\n'.join(current_segment))
                    current_segment = []
        
        # If there are words in current_line, add them as a line
        if current_line:
            current_segment.append(' '.join(current_line))
            current_line = []
            
            # Check if we should start a new segment at sentence boundary
            if len(current_segment) >= lines_per_segment:
                segments.append('\n'.join(current_segment))
                current_segment = []
    
    # Add any remaining lines
    if current_segment:
        segments.append('\n'.join(current_segment))
    
    return segments

async def generate_accurate_srt(text, voice, rate, pitch, words_per_line, lines_per_segment):
    segments = smart_text_split(text, words_per_line, lines_per_segment)
    
    srt_content = ""
    combined_audio = AudioSegment.empty()
    current_time = 0
    
    for idx, segment in enumerate(segments, 1):
        # Generate audio for this segment
        audio_file = f"temp_segment_{idx}.wav"
        tts = edge_tts.Communicate(segment, voice, rate=rate, pitch=pitch)
        await tts.save(audio_file)
        
        # Get segment duration
        segment_audio = AudioSegment.from_file(audio_file)
        segment_duration = len(segment_audio)
        
        # Add to SRT content with precise timing
        srt_content += f"{idx}\n"
        srt_content += f"{format_time_ms(current_time)} --> {format_time_ms(current_time + segment_duration)}\n"
        srt_content += segment + "\n\n"
        
        # Update timing and combine audio
        current_time += segment_duration
        combined_audio += segment_audio
        
        # Cleanup
        os.remove(audio_file)
    
    # Export final files
    unique_id = uuid.uuid4()
    audio_path = f"final_audio_{unique_id}.mp3"
    srt_path = f"final_subtitles_{unique_id}.srt"
    
    combined_audio.export(audio_path, format="mp3", bitrate="320k")
    with open(srt_path, "w", encoding='utf-8') as f:
        f.write(srt_content)
    
    return srt_path, audio_path

async def process_text(text, pitch, rate, voice, words_per_line, lines_per_segment):
    pitch_str = f"{pitch}Hz" if pitch != 0 else "0Hz"
    rate_str = f"{'+' if rate > 0 else ''}{rate}%"
    
    srt_path, audio_path = await generate_accurate_srt(
        text, 
        voice_options[voice],
        rate_str,
        pitch_str,
        words_per_line,
        lines_per_segment
    )
    
    return srt_path, audio_path, audio_path

# Voice options dictionary (same as before)
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 here...
}

# Create Gradio interface
app = gr.Interface(
    fn=process_text,
    inputs=[
        gr.Textbox(label="Enter Text", lines=10),
        gr.Slider(label="Pitch Adjustment (Hz)", minimum=-20, maximum=20, value=0, step=1),
        gr.Slider(label="Rate Adjustment (%)", minimum=-50, maximum=50, value=0, step=1),
        gr.Dropdown(label="Select Voice", choices=list(voice_options.keys()), value="Jenny Female"),
        gr.Slider(label="Words per Line", minimum=1, maximum=15, value=8, step=1),
        gr.Slider(label="Lines per Segment", minimum=1, maximum=5, value=2, step=1)
    ],
    outputs=[
        gr.File(label="Download SRT"),
        gr.File(label="Download Audio"),
        gr.Audio(label="Preview Audio")
    ],
    title="Advanced TTS with Configurable SRT Generation",
    description="Generate perfectly synchronized audio and subtitles with custom segmentation control."
)

app.launch()