scripts/app.py

"""Streamlit UI for ElevenLabs TTS — Emotion Control, Batch Processing, Voice Cloning."""

import io, os, re, time, zipfile
import requests
import streamlit as st
from pydub import AudioSegment

# ── Constants ─────────────────────────────────────────────────────────────────

ELEVENLABS_BASE = "https://api.elevenlabs.io/v1"
SAMPLE_RATE = 44100
DEFAULT_VOICE = "8uYaUDIcnotZPgF7r41V"
DEFAULT_SEED = 42

MODELS = {
    "Eleven v3 (Expressive)": "eleven_v3",
    "Multilingual v2 (Stable)": "eleven_multilingual_v2",
    "Flash v2.5 (Fast)": "eleven_flash_v2_5",
}

# ElevenLabs v3 audio tags mapped to emotion dimensions
EMOTION_TAGS = {
    "happy":    "[happy]",
    "sad":      "[sad]",
    "angry":    "[angry]",
    "calm":     "[calm]",
    "excited":  "[excited]",
    "serious":  "[serious]",
    "whisper":  "[whisper]",
    "soft":     "[soft]",
    "loud":     "[loud]",
}

EMOTION_PRESETS = {
    "Neutral": {},
    "Happy & Excited": {"happy": 0.7, "excited": 0.8},
    "Calm & Gentle": {"calm": 0.8, "soft": 0.6},
    "Sad & Soft": {"sad": 0.7, "soft": 0.5},
    "Angry & Loud": {"angry": 0.8, "loud": 0.7},
    "Serious & Low": {"serious": 0.7},
    "Whisper": {"whisper": 0.9, "soft": 0.4},
    "Excited & Loud": {"excited": 0.9, "loud": 0.6},
}

ACCENT_TAGS = [
    "None", "American accent", "British accent",
    "slight Japanese accent", "Japanese accent",
    "Australian accent", "Indian accent",
]

# ── Helpers ───────────────────────────────────────────────────────────────────

def get_api_key():
    key = os.environ.get("ELEVENLABS_API_KEY", "")
    if not key:
        env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
        if os.path.exists(env_path):
            with open(env_path) as f:
                for line in f:
                    if line.startswith("ELEVENLABS_API_KEY="):
                        key = line.strip().split("=", 1)[1].strip()
    return key


def build_tagged_text(text, emotions, accent, model):
    """Prepend emotion and accent audio tags to text (v3 only)."""
    if model != "eleven_v3":
        return text
    tags = []
    if accent != "None":
        tags.append(f"[{accent}]")
    # Pick the dominant emotion (highest value > 0.3)
    active = {k: v for k, v in emotions.items() if v > 0.3}
    if active:
        dominant = max(active, key=active.get)
        tags.append(EMOTION_TAGS[dominant])
    return " ".join(tags + [text])


def parse_emotion_text(text):
    """Parse (emotion description) prefixed lines into tagged text.
    Example: '(cheerful and excited) Hello!' -> '[happy] [excited] Hello!'
    """
    EMOTION_MAP = {
        "cheerful": "happy", "joyful": "happy", "happy": "happy",
        "excited": "excited", "enthusiastic": "excited",
        "calm": "calm", "gentle": "calm", "peaceful": "calm",
        "sad": "sad", "melancholic": "sad", "sorrowful": "sad",
        "angry": "angry", "furious": "angry", "mad": "angry",
        "serious": "serious", "solemn": "serious", "grave": "serious",
        "whisper": "whisper", "whispering": "whisper", "hushed": "whisper",
        "soft": "soft", "quiet": "soft", "tender": "soft",
        "loud": "loud", "shouting": "loud", "yelling": "loud",
    }
    match = re.match(r"^\(([^)]+)\)\s*(.+)$", text.strip())
    if not match:
        return text, {}
    desc, content = match.group(1).lower(), match.group(2)
    found = {}
    for word in re.findall(r"[a-z]+", desc):
        if word in EMOTION_MAP:
            found[EMOTION_MAP[word]] = 0.8
    return content, found


def tts_call(api_key, voice_id, text, speed, seed, model, stability, similarity):
    """Call ElevenLabs TTS API. Text should already have tags prepended."""
    resp = requests.post(
        f"{ELEVENLABS_BASE}/text-to-speech/{voice_id}",
        headers={
            "xi-api-key": api_key,
            "Content-Type": "application/json",
            "Accept": "audio/mpeg",
        },
        json={
            "text": text,
            "model_id": model,
            "seed": seed,
            "output_format": "mp3_44100_128",
            "voice_settings": {
                "stability": stability,
                "similarity_boost": similarity,
                "speed": round(speed, 3),
            },
        },
        timeout=60,
    )
    resp.raise_for_status()
    return resp.content


# ── Page Config ───────────────────────────────────────────────────────────────

st.set_page_config(page_title="ElevenLabs TTS Studio", layout="wide")
st.title("ElevenLabs TTS Studio")

api_key = get_api_key()
if not api_key:
    st.error("ELEVENLABS_API_KEY not found. Add it to your .env file.")
    st.stop()

# ── Sidebar: Voice & Model ────────────────────────────────────────────────────

with st.sidebar:
    st.header("Voice & Model")

    voice_id = st.text_input("Voice ID", value=DEFAULT_VOICE)

    model_name = st.selectbox("Model", list(MODELS.keys()))
    model_id = MODELS[model_name]

    is_v3 = model_id == "eleven_v3"

    st.divider()
    st.header("Voice Settings")

    stability = st.slider("Stability", 0.0, 1.0, 0.85, 0.05,
                           help="Higher = more consistent, less responsive to tags")
    similarity = st.slider("Similarity Boost", 0.0, 1.0, 0.90, 0.05,
                            help="Higher = closer to original voice clone")
    speed = st.slider("Speed", 0.7, 1.5, 1.0, 0.05)

    st.divider()
    st.header("Determinism")
    seed = st.number_input("Seed", min_value=0, max_value=4294967295, value=DEFAULT_SEED,
                            help="Fixed seed = consistent output across runs")

    st.divider()
    st.header("Accent (v3 only)")
    accent = st.selectbox("Accent Tag", ACCENT_TAGS, disabled=not is_v3)
    if not is_v3 and accent != "None":
        accent = "None"

# ── Main Area: Tabs ───────────────────────────────────────────────────────────

tab_single, tab_batch = st.tabs(["Single", "Batch"])

# ══════════════════════════════════════════════════════════════════════════════
# SINGLE MODE
# ══════════════════════════════════════════════════════════════════════════════

with tab_single:
    st.subheader("Text Input")
    st.caption("Supports emotion-tagged text: `(cheerful and excited) Hello!`")
    text_input = st.text_area("Enter text", height=150, key="single_text",
                               placeholder="Type text here...\nor (happy) This is great!\nor (calm and gentle) Take your time.")

    # Emotion Control
    st.subheader("Emotion Control" + (" (v3 only)" if not is_v3 else ""))

    col_preset, col_custom = st.columns([1, 2])

    with col_preset:
        preset = st.selectbox("Preset", list(EMOTION_PRESETS.keys()))
        preset_vals = EMOTION_PRESETS[preset]

    with col_custom:
        st.caption("Fine-tune each dimension (0.0 = off, 1.0 = max)")

    # Emotion sliders in 2 rows of 4
    emotions = {}
    emotion_keys = list(EMOTION_TAGS.keys())
    row1 = st.columns(5)
    row2 = st.columns(4)
    all_cols = row1 + row2
    for i, emo in enumerate(emotion_keys):
        with all_cols[i]:
            default = preset_vals.get(emo, 0.0)
            emotions[emo] = st.slider(emo.capitalize(), 0.0, 1.0, default, 0.1,
                                       key=f"emo_{emo}", disabled=not is_v3)

    # Generate button
    if st.button("Generate", type="primary", key="single_gen"):
        if not text_input.strip():
            st.error("Enter some text.")
        else:
            with st.spinner("Generating..."):
                try:
                    # Check for emotion-tagged text format
                    clean_text, parsed_emo = parse_emotion_text(text_input.strip())
                    # Merge: parsed emotions from text override slider if present
                    merged_emo = {**emotions, **parsed_emo} if parsed_emo else emotions
                    final_text = build_tagged_text(clean_text, merged_emo, accent, model_id)

                    mp3 = tts_call(api_key, voice_id, final_text, speed, seed, model_id, stability, similarity)
                    audio = AudioSegment.from_mp3(io.BytesIO(mp3)).set_frame_rate(SAMPLE_RATE).set_channels(1)

                    wav_buf = io.BytesIO()
                    audio.export(wav_buf, format="wav")

                    st.session_state["s_audio"] = wav_buf.getvalue()
                    st.session_state["s_mp3"] = mp3
                    st.session_state["s_dur"] = len(audio) / 1000
                    st.session_state["s_tags"] = final_text
                except Exception as e:
                    st.error(f"API Error: {e}")

    # Results
    if "s_audio" in st.session_state:
        st.divider()
        c1, c2 = st.columns([1, 3])
        with c1:
            st.metric("Duration", f"{st.session_state['s_dur']:.2f}s")
        with c2:
            with st.expander("Tags sent to API"):
                st.code(st.session_state["s_tags"])
        st.audio(st.session_state["s_audio"], format="audio/wav")
        col1, col2 = st.columns(2)
        with col1:
            st.download_button("Download WAV", data=st.session_state["s_audio"],
                                file_name="output.wav", mime="audio/wav", key="dl_wav_s")
        with col2:
            st.download_button("Download MP3", data=st.session_state["s_mp3"],
                                file_name="output.mp3", mime="audio/mpeg", key="dl_mp3_s")

# ══════════════════════════════════════════════════════════════════════════════
# BATCH MODE
# ══════════════════════════════════════════════════════════════════════════════

with tab_batch:
    st.subheader("Batch Text Input")
    st.caption("One line per segment. Supports emotion tags: `(happy) Line text`")
    batch_text = st.text_area("Enter lines (one per segment)", height=250, key="batch_text",
                               placeholder="(cheerful) Hello everyone!\n(calm) Welcome to the session.\n(serious) Let's begin.\nPlain text line without emotion.")

    if st.button("Generate Batch", type="primary", key="batch_gen"):
        lines = [l.strip() for l in batch_text.strip().split("\n") if l.strip()]
        if not lines:
            st.error("Enter at least one line.")
        else:
            progress = st.progress(0, text="Starting...")
            results = []

            for i, line in enumerate(lines):
                progress.progress((i + 1) / len(lines), text=f"Segment {i+1}/{len(lines)}")
                try:
                    clean_text, parsed_emo = parse_emotion_text(line)
                    merged_emo = parsed_emo if parsed_emo else emotions
                    final_text = build_tagged_text(clean_text, merged_emo, accent, model_id)

                    mp3 = tts_call(api_key, voice_id, final_text, speed, seed, model_id, stability, similarity)
                    audio = AudioSegment.from_mp3(io.BytesIO(mp3)).set_frame_rate(SAMPLE_RATE).set_channels(1)
                    dur = len(audio) / 1000

                    wav_buf = io.BytesIO()
                    audio.export(wav_buf, format="wav")

                    results.append({
                        "idx": i + 1, "text": clean_text, "tags": final_text,
                        "duration": dur, "wav": wav_buf.getvalue(), "mp3": mp3, "status": "OK",
                    })
                except Exception as e:
                    results.append({
                        "idx": i + 1, "text": line, "tags": "",
                        "duration": 0, "wav": b"", "mp3": b"", "status": f"ERROR: {e}",
                    })
                time.sleep(0.15)

            progress.progress(1.0, text="Done!")
            st.session_state["batch_results"] = results

    # Batch Results
    if "batch_results" in st.session_state:
        results = st.session_state["batch_results"]
        st.divider()

        # Summary
        ok = sum(1 for r in results if r["status"] == "OK")
        err = len(results) - ok
        total_dur = sum(r["duration"] for r in results)
        c1, c2, c3 = st.columns(3)
        c1.metric("Segments", f"{ok}/{len(results)}")
        c2.metric("Errors", err)
        c3.metric("Total Duration", f"{total_dur:.1f}s")

        # Report table
        st.dataframe([{
            "#": r["idx"], "Text": r["text"][:60],
            "Duration": f"{r['duration']:.2f}s", "Status": r["status"],
        } for r in results], use_container_width=True)

        # Per-segment playback
        with st.expander("Preview segments"):
            for r in results:
                if r["wav"]:
                    st.caption(f"#{r['idx']}: {r['text'][:80]}")
                    st.audio(r["wav"], format="audio/wav")

        # Downloads
        col1, col2 = st.columns(2)
        with col1:
            # Stitched output
            combined = AudioSegment.empty()
            for r in results:
                if r["wav"]:
                    combined += AudioSegment.from_wav(io.BytesIO(r["wav"]))
                    combined += AudioSegment.silent(duration=300)  # 300ms gap between segments
            buf = io.BytesIO()
            combined.export(buf, format="wav")
            st.download_button("Download Combined WAV", data=buf.getvalue(),
                                file_name="batch_combined.wav", mime="audio/wav", key="dl_batch_wav")

        with col2:
            zip_buf = io.BytesIO()
            with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf:
                for r in results:
                    if r["wav"]:
                        zf.writestr(f"{r['idx']:04d}.wav", r["wav"])
            st.download_button("Download All Segments (.zip)", data=zip_buf.getvalue(),
                                file_name="batch_segments.zip", mime="application/zip", key="dl_batch_zip")