kokorotts-server/app.py

import os
import subprocess
import tempfile
import re
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename

app = Flask(__name__)

# Configuration from environment variables
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "output")
TTS_COMMAND = os.environ.get("TTS_COMMAND", "kokoro-tts")
PORT = int(os.environ.get("PORT", 5012))

# Ensure output directory exists
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Available voices
VOICES = ["bm_fable", "bm_lewis", "bm_george"]


def parse_speaker_text(text):
    """Parse text into speaker: complete lines pairs for full conversation"""
    # Split by newlines and filter empty ones
    lines = [line.strip() for line in text.strip().split("\n") if line.strip()]

    speakers = {}
    current_speaker = None
    speaker_content = []
    tracks = []

    for line in lines:
        # Check if this is a speaker line (format: "SpeakerName: Some text")
        match = re.match(r"^(.*?): (.*)$", line)
        if match:
            # Handle previous speaker content
            if current_speaker and speaker_content:
                speakers[current_speaker] = "\n".join(speaker_content)

            # Start new speaker
            current_speaker = match.group(1).strip()
            speaker_content = [
                match.group(2).strip()
            ]  # First part after colon as initial content

            tracks += [(current_speaker, speaker_content[0])]
        else:
            # This is a continuation line for the current speaker
            if current_speaker and line.strip():
                speaker_content.append(line.strip())

    # Save final speaker's content
    if current_speaker and speaker_content:
        speakers[current_speaker] = "\n".join(speaker_content)

    return (speakers, tracks)


def assign_voices_to_speakers(speakers):
    """Assign voices to speakers (same voice for each speaker throughout conversation)"""
    speaker_voice_map = {}
    used_voices = set()

    for speaker in speakers.keys():
        # If we haven't assigned a voice yet, pick an unused one
        if not speaker_voice_map.get(speaker):
            available_voices = [voice for voice in VOICES if voice not in used_voices]
            if available_voices:
                selected_voice = available_voices[0]
            else:
                # If all voices are used, pick randomly (this shouldn't happen with few speakers)
                selected_voice = VOICES[0]  # Fallback

            speaker_voice_map[speaker] = selected_voice
            used_voices.add(selected_voice)

    return speaker_voice_map


def generate_wav_files(speakers, speaker_voice_map, title, tracks):
    """Generate individual WAV files for each speaker"""
    wav_files = []
    turn = 0

    for speaker, text in tracks:
        # Create a unique filename for this speaker's contribution
        safe_speaker = secure_filename(speaker)
        filepath = os.path.join(OUTPUT_DIR, f"{title}_{turn}.wav")
        turn += 1

        # Create temporary input file with the text
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False
        ) as temp_file:
            temp_file.write(text)
            temp_input_path = temp_file.name

        try:
            # Run the kokoro-tts command
            cmd = [
                TTS_COMMAND,
                temp_input_path,
                filepath,
                "--voice",
                speaker_voice_map[speaker],
            ]
            subprocess.run(cmd, check=True, capture_output=True)

            # Clean up temporary file
            os.unlink(temp_input_path)

            wav_files.append(filepath)
        except Exception as e:
            print(f"Error generating audio for {speaker}: {e}")
            try:
                os.unlink(temp_input_path)
            except:
                pass
            return None

    return wav_files


def merge_wav_files(wav_files, final_output_path):
    """Merge multiple WAV files into a single file using sox"""
    try:
        # If no files provided, return None
        if not wav_files:
            return None

        # If only one file, copy it and return
        if len(wav_files) == 1:
            import shutil

            shutil.copy2(wav_files[0], final_output_path)
            return final_output_path

        # Use sox to concatenate all WAV files
        cmd = ["sox"] + wav_files + [final_output_path]
        subprocess.run(cmd, check=True, capture_output=True)

        # Clean up all files
        for wav_file in wav_files:
            try:
                os.unlink(wav_file)
            except OSError:
                pass

        return final_output_path

    except Exception as e:
        print(f"Error merging WAV files: {e}")
        # Fallback to returning the last file if merger fails
        if wav_files and len(wav_files) > 0:
            import shutil

            try:
                shutil.copy2(wav_files[-1], final_output_path)
                return final_output_path
            except:
                pass
        return None


@app.route("/tts", methods=["POST"])
def text_to_speech():
    """Endpoint to convert full conversation to audio podcast"""
    data = request.get_json()

    if not data:
        return jsonify({"error": "No JSON data provided"}), 400

    text = data.get("text")
    title = data.get("title")

    if not text or not title:
        return jsonify({"error": "Both 'text' and 'title' fields are required"}), 400

    # Sanitize title to prevent path traversal attacks
    safe_title = secure_filename(title)
    if not safe_title:
        return jsonify({"error": "Invalid title format"}), 400

    # Parse the text to identify speakers and their complete conversations
    (speakers, tracks) = parse_speaker_text(text)

    if not speakers:
        return jsonify({"error": "No speaker lines found in text"}), 400

    # Assign voices to speakers
    speaker_voice_map = assign_voices_to_speakers(speakers)

    # Generate individual WAV files for each speaker
    wav_files = generate_wav_files(speakers, speaker_voice_map, safe_title, tracks)

    if not wav_files:
        return jsonify({"error": "Failed to generate audio files"}), 500

    # Merge all generated WAV files into a single podcast file
    final_wav_path = os.path.join(OUTPUT_DIR, f"{safe_title}.wav")

    merged_file = merge_wav_files(wav_files, final_wav_path)
    if not merged_file:
        return jsonify({"error": "Failed to merge audio files"}), 500

    # Return the location of the complete podcast file
    return jsonify(
        {
            "file_path": final_wav_path,
            "message": "Full conversation podcast generated successfully",
            "speakers": list(speakers.keys()),
            "total_speakers": len(speakers),
            "total_lines": len(text.split("\n")),
        }
    )


@app.route("/health", methods=["GET"])
def health_check():
    """Health check endpoint"""
    return jsonify({"status": "healthy"})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=PORT, debug=False)