app.py

from __future__ import annotations

import os
import threading
from functools import lru_cache
from pathlib import Path
from typing import Optional, Tuple

import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from PackedTTS import PackedTTS


MODEL_REPO_ID = os.getenv("PACKEDTTS_MODEL_REPO_ID", "HiMind/Packed-TTS")
BUNDLE_FILENAME = os.getenv("PACKEDTTS_BUNDLE_FILENAME", "tts.pt")
LOCAL_BUNDLE_PATH = Path(BUNDLE_FILENAME)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_MODEL_LOCK = threading.Lock()
_MODEL: Optional[PackedTTS] = None
_MODEL_ERROR: Optional[str] = None
_MODEL_PATH: Optional[Path] = None


# ---------------------------
# MODEL LOADING
# ---------------------------
def _resolve_bundle_path() -> Path:
    if LOCAL_BUNDLE_PATH.exists():
        return LOCAL_BUNDLE_PATH

    return Path(
        hf_hub_download(
            repo_id=MODEL_REPO_ID,
            repo_type="model",
            filename=BUNDLE_FILENAME,
        )
    )


def _load_model() -> PackedTTS:
    global _MODEL, _MODEL_ERROR, _MODEL_PATH

    with _MODEL_LOCK:
        if _MODEL is not None:
            return _MODEL

        bundle_path = _resolve_bundle_path()
        _MODEL_PATH = bundle_path

        _MODEL = PackedTTS.load(bundle_path, device=DEVICE)
        _MODEL_ERROR = None
        return _MODEL


def _safe_model():
    try:
        return _load_model(), None
    except Exception:
        return None, _MODEL_ERROR or "Model failed to load."


# ---------------------------
# CATALOG
# ---------------------------
@lru_cache(maxsize=1)
def _catalog() -> dict:
    model, err = _safe_model()

    if model is None:
        return {
            "voices": [],
            "emotions": [],
            "error": err,
            "bundle_path": None,
        }

    emotions = model.list_emotions()
    return {
        "voices": model.list_voices(),
        "emotions": list(emotions.keys()),
        "emotion_counts": emotions,
        "error": None,
        "bundle_path": str(_MODEL_PATH),
    }


def _choice(items):
    return ["Auto / default", *items]


def _norm(x):
    if not x or str(x).lower() in {"auto", "default", "auto / default"}:
        return None
    return str(x)


# ---------------------------
# MODEL CALL
# ---------------------------
def _generate(text, voice, emotion, voice_ref, emo_ref,
              cfg_weight, temperature, exaggeration, seed):

    try:
        model = _load_model()

        sr, audio, meta = model.generate(
            text=text,
            voice=_norm(voice),
            emotion=_norm(emotion),
            voice_ref=voice_ref or None,
            emo_ref=emo_ref or None,
            cfg_weight=float(cfg_weight),
            temperature=float(temperature),
            exaggeration=float(exaggeration),
            seed=int(seed),
        )

        audio = np.nan_to_num(np.asarray(audio, dtype=np.float32))

        return (sr, audio), (
            f"Voice: {meta.get('voice')} | "
            f"Emotion: {meta.get('emotion')} | "
            f"SR: {sr}"
        )

    except Exception as e:
        return None, f"❌ {type(e).__name__}: {e}"


# ---------------------------
# REFRESH UI
# ---------------------------
def _refresh():
    c = _catalog()

    voices = _choice(c["voices"])
    emotions = _choice(c["emotions"])

    status = (
        f"❌ {c['error']}" if c["error"]
        else f"✅ Loaded | Voices: {len(c['voices'])} | Emotions: {len(c['emotions'])}"
    )

    return (
        gr.update(choices=voices, value="Auto / default"),
        gr.update(choices=emotions, value="Auto / default"),
        status,
    )


# ---------------------------
# PRESETS
# ---------------------------
def p_basic():
    return "Hello world test", "Sarah", "Disgust", 0.5, 0.8, 0.5, 42

def p_voice():
    return "Voice only test", "Sarah", "Auto / default", 0.5, 0.8, 0.5, 42

def p_emotion():
    return "Emotion only test", "Auto / default", "Happy", 0.5, 0.8, 0.5, 42

def p_default():
    return "Default settings test", "Auto / default", "Auto / default", 0.5, 0.8, 0.5, 42

def p_expressive():
    return "Very expressive speech test", "Sarah", "Angry", 0.7, 0.9, 0.6, 123

def p_refs():
    return "Reference audio mode test", "Auto / default", "Auto / default", 0.5, 0.8, 0.5, 42


# ---------------------------
# UI
# ---------------------------
with gr.Blocks(title="PackedTTS", theme=gr.themes.Soft()) as demo:

    gr.Markdown("# 🎙️ PackedTTS — Voice + Emotion Synthesis")

    status = gr.Markdown("Loading model...")

    with gr.Row():
        text = gr.Textbox(label="Text", lines=4)

    with gr.Row():
        voice = gr.Dropdown(label="Voice", choices=["Auto / default"], value="Auto / default")
        emotion = gr.Dropdown(label="Emotion", choices=["Auto / default"], value="Auto / default")

    with gr.Row():
        voice_ref = gr.Audio(label="Voice reference", type="filepath", sources=["upload"])
        emo_ref = gr.Audio(label="Emotion reference", type="filepath", sources=["upload"])

    with gr.Row():
        cfg = gr.Slider(0, 1.5, value=0.5, label="CFG")
        temp = gr.Slider(0.1, 2.0, value=0.8, label="Temperature")
        expo = gr.Slider(0, 1.0, value=0.5, label="Exaggeration")
        seed = gr.Number(value=42, precision=0, label="Seed")

    with gr.Row():
        btn_gen = gr.Button("Generate", variant="primary")
        btn_refresh = gr.Button("Refresh")

    audio = gr.Audio(label="Output", type="numpy")
    meta = gr.Textbox(label="Info", lines=2)

    # ---------------- PRESETS ----------------
    gr.Markdown("## Presets")

    with gr.Row():
        b1 = gr.Button("Basic")
        b2 = gr.Button("Voice")
        b3 = gr.Button("Emotion")
        b4 = gr.Button("Default")
        b5 = gr.Button("Expressive")
        b6 = gr.Button("Refs")

    # ---------------- IMPORTANT: api_name=False everywhere ----------------
    btn_gen.click(
        _generate,
        [text, voice, emotion, voice_ref, emo_ref, cfg, temp, expo, seed],
        [audio, meta],
        api_name=False
    )

    btn_refresh.click(
        _refresh,
        [],
        [voice, emotion, status],
        api_name=False
    )

    demo.load(
        _refresh,
        [],
        [voice, emotion, status],
        api_name=False
    )

    b1.click(p_basic, [], [text, voice, emotion, cfg, temp, expo, seed], api_name=False)
    b2.click(p_voice, [], [text, voice, emotion, cfg, temp, expo, seed], api_name=False)
    b3.click(p_emotion, [], [text, voice, emotion, cfg, temp, expo, seed], api_name=False)
    b4.click(p_default, [], [text, voice, emotion, cfg, temp, expo, seed], api_name=False)
    b5.click(p_expressive, [], [text, voice, emotion, cfg, temp, expo, seed], api_name=False)
    b6.click(p_refs, [], [text, voice, emotion, cfg, temp, expo, seed], api_name=False)


# IMPORTANT FIX: prevents schema crash again
demo.queue(
    default_concurrency_limit=1,
    api_open=False
).launch()