app.py

# ================================================================
# MTP Voice - app.py para Hugging Face Space (Gradio, CPU)
#
# Este Space reutiliza EXACTAMENTE el modelo MTP-2.5 (RoPE + SwiGLU +
# RMSNorm + KV-cache) del Space de texto, y le agrega:
#
#   - STT (voz -> texto): faster-whisper, modelo "base" cuantizado int8.
#     Entiende el audio que le mandes (mic o archivo) y lo transcribe
#     a español antes de pasarlo al modelo.
#   - TTS (texto -> voz): Piper, modelo ONNX es_ES-davefx-medium.
#     Sintetiza la respuesta de MTP en audio natural en español.
#
# OPTIMIZACIÓN CPU (todo el pipeline corre en CPU, sin GPU):
#   - faster-whisper con compute_type="int8" (cuantizado -> ~2-4x más
#     rápido que float32 en CPU, con pérdida de precisión mínima).
#   - beam_size=1 (greedy) + condition_on_previous_text=False +
#     language="es" fijo (se salta la detección automática de idioma,
#     que es un forward extra del encoder) + vad_filter=True (no
#     procesa silencios, que es tiempo de cómputo tirado a la basura).
#   - Piper es un modelo ONNX pequeño (~60MB) diseñado para tiempo real
#     en CPU (corre incluso en Raspberry Pi), no hace falta cuantizarlo
#     más.
#   - Respuestas de voz con techo de tokens más bajo que el chat de
#     texto: una respuesta hablada de 700 tokens tarda mucho en
#     generarse Y en sintetizarse: no tiene sentido para un asistente
#     de voz. Se limita a algo conversacional.
#   - Los 3 modelos (MTP, Whisper, Piper) se cargan UNA sola vez al
#     arrancar el Space y quedan en memoria; no se recargan por request.
# ================================================================
import os
import io
import math
import wave
import tempfile
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import sentencepiece as spm
from starlette.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
from fastapi import UploadFile, File
from fastapi.responses import Response
from pydantic import BaseModel
from typing import Optional
from huggingface_hub import hf_hub_download

# ---------------- Optimización para CPU ----------------
N_CPU = max(1, os.cpu_count() or 1)
torch.set_num_threads(N_CPU)
try:
    torch.set_num_interop_threads(1)
except RuntimeError:
    pass
torch.set_grad_enabled(False)

DEVICE = "cpu"
_HAS_SDPA = hasattr(F, "scaled_dot_product_attention")

REPO_ID = "TeszenAI/MTP-2"  # <-- mismo repo del checkpoint que en el Space de texto
FILENAME = "MTP2_5_MODEL.pt"

# Voz de Piper (español, calidad media, rápida en CPU). Podés cambiarla por
# cualquier otra voz de https://huggingface.co/rhasspy/piper-voices/tree/main/es
PIPER_REPO = "rhasspy/piper-voices"
PIPER_VOICE_DIR = "es/es_ES/davefx/medium"
PIPER_VOICE_NAME = "es_ES-davefx-medium"

# Modelo de Whisper para STT. "base" es el mejor punto medio velocidad/precisión
# en CPU sin GPU. Si el Space tiene más de 2 vCPU y sobra tiempo, se puede subir
# a "small" cambiando solo esta constante.
WHISPER_MODEL_SIZE = "base"


# ================================================================
# ARQUITECTURA MTP-2.x: RoPE + SwiGLU, con KV-cache (idéntica al Space
# de texto: si cambiás algo acá, tiene que coincidir con el checkpoint)
# ================================================================
def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rope(q, k, cos, sin):
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_rot = (q * cos) + (rotate_half(q) * sin)
    k_rot = (k * cos) + (rotate_half(k) * sin)
    return q_rot, k_rot


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_seq_len, base=10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._build_cache(max_seq_len)

    def _build_cache(self, seq_len):
        t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)
        self.max_seq_len_cached = seq_len

    def forward(self, seq_len, device, dtype, offset=0):
        if offset + seq_len > self.max_seq_len_cached:
            self._build_cache(offset + seq_len)
        cos = self.cos_cached[offset:offset + seq_len].to(device=device, dtype=dtype)
        sin = self.sin_cached[offset:offset + seq_len].to(device=device, dtype=dtype)
        return cos, sin


class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd)
        self.proj = nn.Linear(n_embd, n_embd)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        self.register_buffer("mask", mask)

    def forward(self, x, cos, sin, past_kv=None, use_cache=False):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        q, k = apply_rope(q, k, cos, sin)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        present_kv = (k, v) if use_cache else None
        is_causal = (past_kv is None) and (T > 1)

        if _HAS_SDPA:
            out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=None, dropout_p=0.0, is_causal=is_causal,
            )
        else:
            att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
            if is_causal:
                Tk = k.size(-2)
                causal_mask = torch.tril(torch.ones(T, Tk, device=x.device, dtype=torch.bool))
                att = att.masked_fill(~causal_mask, float("-inf"))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            out = att @ v

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        out = self.resid_dropout(self.proj(out))
        return out, present_kv


class SwiGLU(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        hidden = int(2 * (4 * n_embd) / 3)
        hidden = ((hidden + 7) // 8) * 8
        self.w_gate = nn.Linear(n_embd, hidden, bias=False)
        self.w_up = nn.Linear(n_embd, hidden, bias=False)
        self.w_down = nn.Linear(hidden, n_embd, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.dropout(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return norm * self.weight


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = RMSNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = RMSNorm(n_embd)
        self.ff = SwiGLU(n_embd, dropout)

    def forward(self, x, cos, sin, past_kv=None, use_cache=False):
        attn_out, present_kv = self.attn(self.ln1(x), cos, sin, past_kv=past_kv, use_cache=use_cache)
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x, present_kv


class MTP(nn.Module):
    def __init__(self, vocab_size, block_size, n_layer, n_head, n_embd, dropout):
        super().__init__()
        self.block_size = block_size
        self.head_dim = n_embd // n_head
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.rope = RotaryEmbedding(self.head_dim, max_seq_len=block_size)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
        self.ln_f = RMSNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight

    def forward(self, idx, past_key_values=None, use_cache=False, pos_offset=0):
        B, T = idx.shape
        x = self.tok_emb(idx)
        x = self.drop(x)
        cos, sin = self.rope(T, idx.device, x.dtype, offset=pos_offset)

        new_past = [] if use_cache else None
        for i, block in enumerate(self.blocks):
            past_kv = past_key_values[i] if past_key_values is not None else None
            x, present_kv = block(x, cos, sin, past_kv=past_kv, use_cache=use_cache)
            if use_cache:
                new_past.append(present_kv)

        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits, new_past


# ---------------- Carga del checkpoint MTP (una sola vez) ----------------
print("[MTP] Descargando checkpoint desde el Hub...")
ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
checkpoint = torch.load(ckpt_path, map_location=DEVICE)

cfg = checkpoint["config"]
special = checkpoint["special_tokens"]
gen_defaults = checkpoint["generation_defaults"]

PAD_ID, BOS_ID, EOS_ID, UNK_ID = special["pad_id"], special["bos_id"], special["eos_id"], special["unk_id"]

sp = spm.SentencePieceProcessor()
sp.load_from_serialized_proto(checkpoint["spm_model_bytes"])

model = MTP(
    vocab_size=cfg["vocab_size"], block_size=cfg["block_size"],
    n_layer=cfg["n_layer"], n_head=cfg["n_head"],
    n_embd=cfg["n_embd"], dropout=cfg["dropout"],
).to(DEVICE)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

BLOCK_SIZE = cfg["block_size"]

print(f"[MTP] Cargado ({checkpoint['meta']['model_name']}, "
      f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos)"
      f" | SDPA={'sí' if _HAS_SDPA else 'no (fallback manual)'}")


import re as _re_indent


def protect_indentation(text):
    lines = text.split("\n")
    new_lines = []
    for line in lines:
        stripped = line.lstrip(" ")
        n_spaces = len(line) - len(stripped)
        n_levels = n_spaces // 4
        remainder = n_spaces % 4
        if n_levels > 0:
            prefix = " " + " ".join(["<tab>"] * n_levels) + " " + " " * remainder
        else:
            prefix = " " * remainder
        new_lines.append(prefix + stripped)
    text = "\n".join(new_lines)

    def _repl(m):
        n = len(m.group())
        return " " + " ".join(["<nl>"] * n) + " "

    text = _re_indent.sub(r"\n+", _repl, text)
    return text


def restore_indentation(text):
    text = _re_indent.sub(r"(<nl>\s*)+", lambda m: "\n" * m.group().count("<nl>"), text)
    text = _re_indent.sub(r"(<tab>\s*)+", lambda m: "    " * m.group().count("<tab>"), text)
    return text


def encode_text(s):
    return sp.encode(protect_indentation(s), out_type=int)


def decode_ids(ids):
    text = sp.decode([i for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID)])
    return restore_indentation(text)


def _block_repeated_ngrams(generated_ids, logits, ngram_size):
    if ngram_size <= 0 or len(generated_ids) < ngram_size:
        return logits
    prefix = tuple(generated_ids[-(ngram_size - 1):])
    banned = set()
    for i in range(len(generated_ids) - ngram_size + 1):
        if tuple(generated_ids[i:i + ngram_size - 1]) == prefix:
            banned.add(generated_ids[i + ngram_size - 1])
    if banned:
        logits[0, list(banned)] = float("-inf")
    return logits


@torch.inference_mode()
def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size=3):
    past_key_values = None
    cache_len = 0

    for _ in range(max_new_tokens):
        total_len = idx.shape[1]

        if total_len <= BLOCK_SIZE:
            if past_key_values is None:
                logits, past_key_values = model(idx, use_cache=True)
                cache_len = total_len
            else:
                last_token = idx[:, -1:]
                logits, past_key_values = model(
                    last_token,
                    past_key_values=past_key_values,
                    use_cache=True,
                    pos_offset=cache_len,
                )
                cache_len += 1
            logits = logits[:, -1, :]
        else:
            idx_cond = idx[:, -BLOCK_SIZE:]
            logits, past_key_values = model(idx_cond, use_cache=True)
            cache_len = BLOCK_SIZE
            logits = logits[:, -1, :]

        logits = logits / max(temperature, 1e-5)

        if repetition_penalty and repetition_penalty != 1.0:
            unique_ids = torch.unique(idx[0])
            logits[0, unique_ids] /= repetition_penalty

        logits = _block_repeated_ngrams(idx[0].tolist(), logits, no_repeat_ngram_size)

        if top_k is not None and top_k > 0:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float("-inf")

        probs = F.softmax(logits, dim=-1)

        if top_p is not None and 0 < top_p < 1:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cum_probs = torch.cumsum(sorted_probs, dim=-1)
            cutoff = cum_probs > top_p
            cutoff[:, 1:] = cutoff[:, :-1].clone()
            cutoff[:, 0] = False
            sorted_probs[cutoff] = 0.0
            sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
            next_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))
        else:
            next_id = torch.multinomial(probs, num_samples=1)

        idx = torch.cat([idx, next_id], dim=1)
        if next_id.item() == EOS_ID:
            break

    return idx


# Techo de tokens para VOZ: una respuesta hablada larga es lenta de generar Y
# de sintetizar. 220 tokens ya da respuestas de sobra para una conversación
# hablada natural, sin hacer esperar al usuario.
VOICE_MAX_TOKENS_HARD_LIMIT = 220


def run_inference(text, max_new_tokens=None, temperature=None, top_k=None, top_p=None,
                   repetition_penalty=None, no_repeat_ngram_size=None, hard_limit=700):
    max_new_tokens = int(max_new_tokens) if max_new_tokens else gen_defaults["max_new_tokens"]
    temperature = float(temperature) if temperature is not None else gen_defaults["temperature"]
    top_k = int(top_k) if top_k is not None else gen_defaults["top_k"]
    top_p = float(top_p) if top_p is not None else gen_defaults["top_p"]
    repetition_penalty = float(repetition_penalty) if repetition_penalty is not None else gen_defaults["repetition_penalty"]
    no_repeat_ngram_size = int(no_repeat_ngram_size) if no_repeat_ngram_size is not None else gen_defaults.get("no_repeat_ngram_size", 3)

    max_new_tokens = max(1, min(max_new_tokens, hard_limit))

    prefix = f"Usuario: {text}\nMTP: "
    ids = [BOS_ID] + encode_text(prefix)
    idx = torch.tensor([ids], dtype=torch.long, device=DEVICE)

    out = generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty, no_repeat_ngram_size)
    new_ids = out[0].tolist()[len(ids):]
    return decode_ids(new_ids).strip()


# ================================================================
# STT: faster-whisper (voz -> texto), optimizado para CPU
# ================================================================
from faster_whisper import WhisperModel

print(f"[STT] Cargando faster-whisper '{WHISPER_MODEL_SIZE}' (int8, CPU)...")
whisper_model = WhisperModel(
    WHISPER_MODEL_SIZE,
    device="cpu",
    compute_type="int8",     # cuantizado: 2-4x más rápido en CPU que float32
    cpu_threads=N_CPU,
    num_workers=1,
)
print("[STT] Whisper listo.")


def transcribe_audio(audio_path):
    """Transcribe un archivo de audio a texto en español. Devuelve '' si no
    detecta voz (silencio, ruido, etc)."""
    if not audio_path:
        return ""
    segments, _info = whisper_model.transcribe(
        audio_path,
        language="es",                # nos saltamos la detección de idioma
        task="transcribe",
        beam_size=1,                  # greedy: mucho más rápido que beam_size=5
        best_of=1,
        temperature=0.0,
        condition_on_previous_text=False,
        vad_filter=True,              # no procesa silencios
        without_timestamps=True,
    )
    return " ".join(seg.text.strip() for seg in segments).strip()


# ================================================================
# TTS: Piper (texto -> voz), optimizado para CPU
# ================================================================
from piper import PiperVoice
from piper.config import SynthesisConfig, PiperConfig
import json as _json
import onnxruntime as _ort

print(f"[TTS] Descargando voz Piper '{PIPER_VOICE_NAME}'...")
piper_model_path = hf_hub_download(repo_id=PIPER_REPO, filename=f"{PIPER_VOICE_DIR}/{PIPER_VOICE_NAME}.onnx")
piper_config_path = hf_hub_download(repo_id=PIPER_REPO, filename=f"{PIPER_VOICE_DIR}/{PIPER_VOICE_NAME}.onnx.json")

# --- Carga manual (en vez de PiperVoice.load) para controlar los threads de
# ONNX Runtime. Por qué: por defecto ORT puede repartir tanto intra-op como
# inter-op de forma que compite con los threads de PyTorch/Whisper y termina
# gastando más CPU total sin sintetizar más rápido (sobre-suscripción de
# hilos). Fijamos intra_op = todos los núcleos (para que ESTA inferencia use
# el CPU disponible) e inter_op = 1 (no lanza pools de threads extra), que es
# el patrón recomendado para servir un solo modelo pequeño en CPU.
with open(piper_config_path, "r", encoding="utf-8") as f:
    _piper_cfg_dict = _json.load(f)

_sess_opts = _ort.SessionOptions()
_sess_opts.intra_op_num_threads = N_CPU
_sess_opts.inter_op_num_threads = 1
_sess_opts.graph_optimization_level = _ort.GraphOptimizationLevel.ORT_ENABLE_ALL

piper_voice = PiperVoice(
    session=_ort.InferenceSession(piper_model_path, sess_options=_sess_opts, providers=["CPUExecutionProvider"]),
    config=PiperConfig.from_dict(_piper_cfg_dict),
)
print("[TTS] Piper listo.")

# length_scale < 1.0 = habla un poco más rápido -> genera menos muestras de
# audio por respuesta -> menos cómputo real por frase, sin tocar el modelo
# (misma voz, mismo timbre, solo un pelín más ágil). 0.92 es el punto donde
# todavía suena natural pero ya se nota más liviano en CPU.
PIPER_LENGTH_SCALE = 0.92


def synthesize_speech(text):
    """Sintetiza texto a un .wav en disco."""
    if not text or not text.strip():
        return None
    fd, path = tempfile.mkstemp(suffix=".wav")
    os.close(fd)
    syn_config = SynthesisConfig(length_scale=PIPER_LENGTH_SCALE)
    with wave.open(path, "wb") as wav_file:
        piper_voice.synthesize_wav(text, wav_file, syn_config=syn_config)
    return path


# ================================================================
# Lógica compartida: un turno completo de voz (audio -> texto -> audio)
# ================================================================
def voice_turn(user_text):
    """Dado un texto (ya transcrito o tipeado), genera la respuesta de MTP
    y la sintetiza en audio. Devuelve (respuesta_texto, ruta_audio)."""
    reply = run_inference(user_text, hard_limit=VOICE_MAX_TOKENS_HARD_LIMIT,
                           max_new_tokens=gen_defaults.get("max_new_tokens", 200))
    if not reply:
        reply = "No pude generar una respuesta."
    audio_path = synthesize_speech(reply)
    return reply, audio_path


# ================================================================
# Interfaz Gradio "MTP Voice" (mismo look & feel dark/glass del chat MTP)
# ================================================================
CUSTOM_CSS = """
:root {
    --bg-color: #0f0f0f;
    --surface-input: #1a1a1a;
    --surface-hover: #2a2a2a;
    --user-bubble: #282828;
    --text-primary: #f1f1f1;
    --text-secondary: #8e8e8e;
    --text-muted: #555555;
    --accent-color: #4a9eff;
}
.gradio-container { background: var(--bg-color) !important; font-family: 'Inter', sans-serif !important; }
#mtp-header {
    display: flex; align-items: center; gap: 14px; padding: 6px 4px 18px 4px;
}
#mtp-header img {
    width: 44px; height: 44px; border-radius: 50%; object-fit: cover;
    box-shadow: 0 4px 16px rgba(0,0,0,0.5);
}
#mtp-header .title { font-size: 1.3rem; font-weight: 600; color: var(--text-primary); }
#mtp-header .subtitle { font-size: 0.82rem; color: var(--text-secondary); }

#mtp-chat-box {
    background: transparent; border: none; min-height: 320px; max-height: 480px;
    overflow-y: auto; padding: 4px 2px;
}
.msg-row { display: flex; gap: 12px; width: 100%; margin-bottom: 18px; }
.msg-row.user { justify-content: flex-end; }
.msg-row.bot { justify-content: flex-start; align-items: flex-start; }
.msg-content { line-height: 1.55; font-size: 0.98rem; word-wrap: break-word; max-width: 82%; white-space: pre-wrap; }
.user .msg-content { background-color: var(--user-bubble); padding: 11px 16px; border-radius: 20px; color: #fff; }
.bot .msg-content-wrapper { display: flex; flex-direction: column; gap: 4px; max-width: 88%; }
.bot .msg-text { color: var(--text-primary); white-space: pre-wrap; }
.bot-avatar {
    width: 30px; height: 30px; min-width: 30px; border-radius: 50%;
    background: linear-gradient(135deg, var(--accent-color), #2a6fd6);
    display: flex; align-items: center; justify-content: center; font-size: 0.75rem; color: #fff; font-weight: 700;
}
#mtp-empty {
    text-align: center; color: var(--text-secondary); font-size: 0.92rem; padding: 60px 10px;
}

#mtp-mic-row { display: flex; justify-content: center; margin: 10px 0 4px 0; }
#mtp-mic-row .wrap { border-radius: 24px !important; }

.gr-button-primary, button.primary {
    background: var(--accent-color) !important; border: none !important;
}
"""

WELCOME = """
<div id="mtp-empty">
  🎙️ Grabá un mensaje o escribí abajo.<br>MTP te va a responder también con voz.
</div>
"""


def render_history_html(history):
    if not history:
        return WELCOME
    rows = []
    for turn in history:
        if turn["role"] == "user":
            rows.append(
                f'<div class="msg-row user"><div class="msg-content">{_escape(turn["text"])}</div></div>'
            )
        else:
            rows.append(
                '<div class="msg-row bot">'
                '<div class="bot-avatar">M</div>'
                '<div class="msg-content-wrapper">'
                f'<div class="msg-text">{_escape(turn["text"])}</div>'
                '</div></div>'
            )
    return '<div>' + "".join(rows) + '</div>'


def _escape(text):
    return (
        (text or "")
        .replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace("\n", "<br>")
    )


def on_audio_turn(audio_path, history):
    transcript = transcribe_audio(audio_path)
    if not transcript:
        # No se detectó voz: no agregamos turno, solo avisamos.
        return render_history_html(history), None, history, gr.update(value=None)
    reply, wav_path = voice_turn(transcript)
    history = history + [{"role": "user", "text": transcript}, {"role": "bot", "text": reply}]
    return render_history_html(history), wav_path, history, gr.update(value=None)


def on_text_turn(user_text, history):
    user_text = (user_text or "").strip()
    if not user_text:
        return render_history_html(history), None, history, gr.update(value="")
    reply, wav_path = voice_turn(user_text)
    history = history + [{"role": "user", "text": user_text}, {"role": "bot", "text": reply}]
    return render_history_html(history), wav_path, history, gr.update(value="")


with gr.Blocks(title="MTP Voice") as demo:
    history_state = gr.State([])

    gr.HTML(
        '<div id="mtp-header">'
        '<img src="https://i.postimg.cc/wv3nLLGN/image.png">'
        '<div><div class="title">MTP Voice</div>'
        '<div class="subtitle">MTP-2.5 con voz · corriendo en CPU</div></div>'
        '</div>'
    )

    chat_html = gr.HTML(WELCOME, elem_id="mtp-chat-box")

    with gr.Row(elem_id="mtp-mic-row"):
        mic = gr.Audio(sources=["microphone"], type="filepath", label="Grabar mensaje", show_label=False)

    with gr.Row():
        text_input = gr.Textbox(placeholder="...o escribí acá", show_label=False, scale=5)
        send_btn = gr.Button("Enviar", variant="primary", scale=1)

    audio_out = gr.Audio(label="Respuesta de MTP", autoplay=True)

    mic.stop_recording(
        fn=on_audio_turn,
        inputs=[mic, history_state],
        outputs=[chat_html, audio_out, history_state, mic],
    )
    send_btn.click(
        fn=on_text_turn,
        inputs=[text_input, history_state],
        outputs=[chat_html, audio_out, history_state, text_input],
    )
    text_input.submit(
        fn=on_text_turn,
        inputs=[text_input, history_state],
        outputs=[chat_html, audio_out, history_state, text_input],
    )

demo.queue(max_size=16)


# ================================================================
# API REST (para integrarlo desde PHP, igual que el Space de texto):
#   POST /transcribe  (multipart, campo "audio")        -> {"text": "..."}
#   POST /generate     {"text": "..."}                    -> {"reply": "..."}
#   POST /tts          {"text": "..."}                    -> audio/wav (bytes)
# Chainear las 3 desde PHP te da voz->texto->respuesta->voz sin tocar Python.
# ================================================================
class GenerateRequest(BaseModel):
    text: str
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    repetition_penalty: Optional[float] = None
    no_repeat_ngram_size: Optional[int] = None


class TTSRequest(BaseModel):
    text: str


PORT = int(os.environ.get("PORT", 7860))
demo.launch(
    server_name="0.0.0.0",
    server_port=PORT,
    prevent_thread_lock=True,
    ssr_mode=False,
    css=CUSTOM_CSS,
    theme=gr.themes.Base(),
    app_kwargs={
        "middleware": [
            Middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]),
        ]
    },
)

app = demo.app


@app.post("/generate")
def generate_endpoint(req: GenerateRequest):
    if not req.text or not req.text.strip():
        return {"reply": "Escribe algo para que pueda responder."}
    try:
        reply = run_inference(
            req.text,
            max_new_tokens=req.max_tokens,
            temperature=req.temperature,
            top_k=req.top_k,
            top_p=req.top_p,
            repetition_penalty=req.repetition_penalty,
            no_repeat_ngram_size=req.no_repeat_ngram_size,
        )
        return {"reply": reply or "No pude generar una respuesta."}
    except Exception as e:
        return {"reply": f"Error del modelo: {e}"}


@app.post("/transcribe")
async def transcribe_endpoint(audio: UploadFile = File(...)):
    try:
        suffix = os.path.splitext(audio.filename or "audio.wav")[1] or ".wav"
        fd, tmp_path = tempfile.mkstemp(suffix=suffix)
        with os.fdopen(fd, "wb") as f:
            f.write(await audio.read())
        text = transcribe_audio(tmp_path)
        os.remove(tmp_path)
        return {"text": text}
    except Exception as e:
        return {"text": "", "error": str(e)}


@app.post("/tts")
def tts_endpoint(req: TTSRequest):
    if not req.text or not req.text.strip():
        return Response(status_code=400, content="texto vacío")
    wav_path = synthesize_speech(req.text)
    if not wav_path:
        return Response(status_code=500, content="no se pudo sintetizar")
    with open(wav_path, "rb") as f:
        data = f.read()
    os.remove(wav_path)
    return Response(content=data, media_type="audio/wav")


@app.get("/generate")
def generate_health():
    return {"status": "ok", "info": "MTP Voice: /generate, /transcribe, /tts"}


demo.block_thread()