app.py

import os
import re
import time
import uuid
import queue
import threading
import subprocess
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import requests
import gradio as gr
from scipy.signal import resample_poly
from faster_whisper import WhisperModel

try:
    import webrtcvad
except ImportError as e:
    raise ImportError(
        "缺少 webrtcvad，请先安装：pip install webrtcvad-wheels\n"
        "（用 webrtcvad-wheels 这个 fork，预编译好的 wheel，装起来比原版 webrtcvad 省心）"
    ) from e

# =============================================================================
# 基础配置（与原来一致）
# =============================================================================
API_KEY = os.environ.get("SILICONFLOW_API_KEY", "")

ASR_URL = "https://api.siliconflow.cn/v1/audio/transcriptions"
CHAT_URL = "https://api.siliconflow.cn/v1/chat/completions"
TTS_URL = "https://api.siliconflow.cn/v1/audio/speech"

ASR_MODEL = "FunAudioLLM/SenseVoiceSmall"
MT_MODEL = "tencent/Hunyuan-MT-7B"
TTS_MODEL = "FunAudioLLM/CosyVoice2-0.5B"

AUDIO_DIR = "/tmp/chuanshengtong_audio"
os.makedirs(AUDIO_DIR, exist_ok=True)

FRENCH_ASR_MODEL_SIZE = "medium"
french_asr_model = WhisperModel(FRENCH_ASR_MODEL_SIZE, device="cpu", compute_type="int8")

# ---------------------------------------------------------------------------
# Piper 本地法语 TTS（原样保留）
# ---------------------------------------------------------------------------
PIPER_BIN = "piper"
PIPER_MODEL_DIR = "piper_models"
os.makedirs(PIPER_MODEL_DIR, exist_ok=True)

PIPER_MODEL_FR = os.path.join(PIPER_MODEL_DIR, "fr_FR-siwis-medium.onnx")
PIPER_MODEL_FR_JSON = os.path.join(PIPER_MODEL_DIR, "fr_FR-siwis-medium.onnx.json")

PIPER_MODEL_URL = "https://huggingface.co/rhasspy/piper-voices/resolve/main/fr/fr_FR/siwis/medium/fr_FR-siwis-medium.onnx"
PIPER_MODEL_JSON_URL = "https://huggingface.co/rhasspy/piper-voices/resolve/main/fr/fr_FR/siwis/medium/fr_FR-siwis-medium.onnx.json"


def ensure_piper_model():
    try:
        if not os.path.exists(PIPER_MODEL_FR):
            print("正在下载 Piper 法语模型（onnx）...")
            r = requests.get(PIPER_MODEL_URL, timeout=120)
            r.raise_for_status()
            with open(PIPER_MODEL_FR, "wb") as f:
                f.write(r.content)
        if not os.path.exists(PIPER_MODEL_FR_JSON):
            print("正在下载 Piper 法语模型配置（json）...")
            r = requests.get(PIPER_MODEL_JSON_URL, timeout=60)
            r.raise_for_status()
            with open(PIPER_MODEL_FR_JSON, "wb") as f:
                f.write(r.content)
        print("Piper 法语模型已就绪")
    except Exception as e:
        print(f"⚠️ Piper 法语模型下载失败：{e}（法语合成方向届时会报错，其他方向不受影响）")


ensure_piper_model()

# =============================================================================
# 翻译方向配置（原样保留，同传模式只会用到中↔英这两个方向）
# =============================================================================
LANG_MAP = {
    "英语 → 中文": {"src": "English", "tgt": "Chinese", "src_tag": "EN", "tgt_tag": "中",
                    "asr": "siliconflow", "tts": "siliconflow", "voice": "FunAudioLLM/CosyVoice2-0.5B:alex"},
    "汉语 → 英文": {"src": "Chinese", "tgt": "English", "src_tag": "中", "tgt_tag": "EN",
                    "asr": "siliconflow", "tts": "siliconflow", "voice": "FunAudioLLM/CosyVoice2-0.5B:david"},
    "汉语 → 法文": {"src": "Chinese", "tgt": "French", "src_tag": "中", "tgt_tag": "FR",
                    "asr": "siliconflow", "tts": "piper_fr", "voice": None},
    "法语 → 中文": {"src": "French", "tgt": "Chinese", "src_tag": "FR", "tgt_tag": "中",
                    "asr": "local_french", "tts": "siliconflow", "voice": "FunAudioLLM/CosyVoice2-0.5B:alex"},
    "英语 → 法文": {"src": "English", "tgt": "French", "src_tag": "EN", "tgt_tag": "FR",
                    "asr": "siliconflow", "tts": "piper_fr", "voice": None},
    "法语 → 英文": {"src": "French", "tgt": "English", "src_tag": "FR", "tgt_tag": "EN",
                    "asr": "local_french", "tts": "siliconflow", "voice": "FunAudioLLM/CosyVoice2-0.5B:david"},
}

# 同传模式目前只开放这两个方向（对应"1A 2A"：CPU + 先做中↔英）
STREAMING_LANG_MAP = {
    "英语 → 中文": LANG_MAP["英语 → 中文"],
    "汉语 → 英文": LANG_MAP["汉语 → 英文"],
}

# =============================================================================
# 原来的批处理函数（点击录音模式用，原样保留）
# =============================================================================


def transcribe_siliconflow(audio_path):
    with open(audio_path, "rb") as f:
        files = {"file": (os.path.basename(audio_path), f, "audio/wav")}
        data = {"model": ASR_MODEL}
        headers = {"Authorization": "Bearer " + API_KEY}
        r = requests.post(ASR_URL, headers=headers, data=data, files=files, timeout=60)
    r.raise_for_status()
    text = r.json().get("text", "").strip()
    text = text.replace("<|", "").replace("|>", "")
    return text


def transcribe_french_local(audio_path):
    segments, _info = french_asr_model.transcribe(
        audio_path, language="fr", task="transcribe", vad_filter=True,
    )
    return "".join(seg.text for seg in segments).strip()


def transcribe(audio_path, engine):
    if engine == "local_french":
        return transcribe_french_local(audio_path)
    return transcribe_siliconflow(audio_path)


def translate(text, src, tgt):
    headers = {"Authorization": "Bearer " + API_KEY, "Content-Type": "application/json"}
    payload = {
        "model": MT_MODEL,
        "messages": [
            {
                "role": "system",
                "content": f"You are a professional translator. Translate the user text from {src} to {tgt}. "
                           f"Output ONLY the translated text, no explanations, no quotes, no pinyin.",
            },
            {"role": "user", "content": text},
        ],
        "temperature": 0.3,
        "max_tokens": 512,
    }
    r = requests.post(CHAT_URL, headers=headers, json=payload, timeout=60)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"].strip()


def speak_siliconflow(text, voice, out_path):
    headers = {"Authorization": "Bearer " + API_KEY, "Content-Type": "application/json"}
    payload = {"model": TTS_MODEL, "input": text, "voice": voice, "response_format": "mp3"}
    r = requests.post(TTS_URL, headers=headers, json=payload, timeout=60)
    r.raise_for_status()
    with open(out_path, "wb") as f:
        f.write(r.content)
    return out_path


def speak_piper_french(text, out_path):
    if not os.path.exists(PIPER_MODEL_FR):
        raise RuntimeError(f"未找到 Piper 法语模型文件：{PIPER_MODEL_FR}，请检查启动日志中的下载是否成功")
    cmd = [PIPER_BIN, "--model", PIPER_MODEL_FR, "--output_file", out_path]
    proc = subprocess.run(cmd, input=text.encode("utf-8"), capture_output=True)
    if proc.returncode != 0:
        err = proc.stderr.decode("utf-8", errors="ignore")
        raise RuntimeError(f"Piper 合成失败：{err}")
    return out_path


def speak(text, cfg):
    uid = uuid.uuid4().hex[:8]
    if cfg["tts"] == "piper_fr":
        out_path = os.path.join(AUDIO_DIR, f"{uid}.wav")
        return speak_piper_french(text, out_path)
    else:
        out_path = os.path.join(AUDIO_DIR, f"{uid}.mp3")
        return speak_siliconflow(text, cfg["voice"], out_path)


def process(audio_path, mode_label, history):
    history = history or []
    if not API_KEY:
        history.append({"role": "assistant", "content": "⚠️ 服务器未配置 SILICONFLOW_API_KEY，请检查 Space 的 Settings -> Repository secrets"})
        return history, None, None
    if audio_path is None:
        return history, None, None
    cfg = LANG_MAP[mode_label]
    try:
        text = transcribe(audio_path, cfg["asr"])
        if not text:
            return history, None, None
        trans = translate(text, cfg["src"], cfg["tgt"])
        tts_path = speak(trans, cfg)
    except Exception as e:
        history.append({"role": "assistant", "content": f"⚠️ 出错了：{e}"})
        return history, None, None
    history.append({"role": "user", "content": f"[{cfg['src_tag']}] {text}"})
    history.append({"role": "assistant", "content": f"[{cfg['tgt_tag']}] {trans}"})
    history.append({"role": "assistant", "content": {"path": tts_path}})
    return history, None, tts_path


# =============================================================================
# 同声传译模式：新增部分开始
# =============================================================================

SAMPLE_RATE = 16000
FRAME_MS = 30                          # webrtcvad 只接受 10/20/30ms 帧
FRAME_SAMPLES = SAMPLE_RATE * FRAME_MS // 1000   # 480
SILENCE_END_MS = 600                   # 静音超过这个时长，判定"一句话说完了"——按你实测效果调
SILENCE_END_FRAMES = SILENCE_END_MS // FRAME_MS
PARTIAL_UPDATE_SEC = 1.0               # 隔多久刷新一次"正在识别"的临时字幕，太频繁会拖垮 CPU
VAD_AGGRESSIVENESS = 2                 # 0~3，越大越倾向于把不确定的帧判成"非语音"

# 中/英同传共用一个本地模型，别分开加载浪费内存
# CPU 环境：small + int8 是速度和准确率的折中；机器扛不住可以降到 base，效果会差一些
zh_en_asr_model = WhisperModel("small", device="cpu", compute_type="int8")


def to_16k_mono_float32(sr: int, audio: np.ndarray) -> np.ndarray:
    """把 Gradio 传来的任意采样率/声道音频，统一转成 16kHz 单声道 float32 [-1, 1]"""
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    audio = audio.astype(np.float32)
    if np.abs(audio).max() > 1.0:
        audio = audio / 32768.0  # 原始是 int16 range
    if sr != SAMPLE_RATE:
        audio = resample_poly(audio, SAMPLE_RATE, sr).astype(np.float32)
    return audio


def float32_to_pcm16_bytes(audio_f32: np.ndarray) -> bytes:
    clipped = np.clip(audio_f32, -1.0, 1.0)
    return (clipped * 32767).astype(np.int16).tobytes()


class VadSegmenter:
    """逐 30ms 帧跑 WebRTC VAD，判断当前是"说话中"还是"停顿到可以收尾了" """

    def __init__(self, aggressiveness=VAD_AGGRESSIVENESS):
        self.vad = webrtcvad.Vad(aggressiveness)
        self.silence_frames = 0
        self.in_speech = False
        self.last_is_speech = False

    def feed(self, frame_bytes: bytes) -> str:
        """返回 'speaking' / 'pause_end' / 'silence'"""
        is_speech = self.vad.is_speech(frame_bytes, SAMPLE_RATE)
        self.last_is_speech = is_speech
        if is_speech:
            self.silence_frames = 0
            self.in_speech = True
            return "speaking"
        if self.in_speech:
            self.silence_frames += 1
            if self.silence_frames >= SILENCE_END_FRAMES:
                self.in_speech = False
                self.silence_frames = 0
                return "pause_end"
            return "speaking"  # 短暂停顿，还不算说完一句
        return "silence"


class StreamingClauseASR:
    """
    维护"当前正在说的这一句"的音频缓冲区，随缓冲增长周期性重新识别。
    用 LocalAgreement 策略：前后两次识别结果里一致的稳定前缀才算"确认文本"
    展示出来，剩余部分标为"待确认"（下一轮可能改），这样能做到边说边出字，
    又不会因为 whisper 反复改主意导致画面疯狂抖动。
    """

    def __init__(self, model: WhisperModel, language: str):
        self.model = model
        self.language = language  # "zh" or "en"
        self.buffer = np.zeros(0, dtype=np.float32)
        self.confirmed_text = ""
        self.last_hyp = ""

    def add_audio(self, frame_f32: np.ndarray):
        self.buffer = np.concatenate([self.buffer, frame_f32])

    def partial_transcribe(self):
        if len(self.buffer) < SAMPLE_RATE * 0.3:
            return self.confirmed_text, ""
        segments, _ = self.model.transcribe(
            self.buffer, language=self.language, task="transcribe",
            beam_size=1, vad_filter=False, condition_on_previous_text=False,
        )
        hyp = "".join(s.text for s in segments).strip()
        stable = self._common_prefix(self.last_hyp, hyp)
        self.last_hyp = hyp
        if len(stable) > len(self.confirmed_text):
            self.confirmed_text = stable
        tentative = hyp[len(self.confirmed_text):]
        return self.confirmed_text, tentative

    @staticmethod
    def _common_prefix(a: str, b: str) -> str:
        n = min(len(a), len(b))
        i = 0
        while i < n and a[i] == b[i]:
            i += 1
        return a[:i]

    def finalize(self) -> str:
        """VAD 判定一句话说完了，做一次高质量的完整识别，拿到这句话的最终文本"""
        if len(self.buffer) < SAMPLE_RATE * 0.2:
            text = self.confirmed_text
        else:
            segments, _ = self.model.transcribe(
                self.buffer, language=self.language, task="transcribe",
                beam_size=5, vad_filter=True, condition_on_previous_text=False,
            )
            text = "".join(s.text for s in segments).strip()
        self.buffer = np.zeros(0, dtype=np.float32)
        self.confirmed_text = ""
        self.last_hyp = ""
        return text


# --- 后台流水线：翻译、合成各跑一条常驻线程，互不阻塞 -----------------------

translate_queue: "queue.Queue" = queue.Queue()
speak_queue: "queue.Queue" = queue.Queue()

_session_lock = threading.Lock()
_session_result_queues: dict[str, "queue.Queue"] = {}


def get_session_queue(session_id: str) -> "queue.Queue":
    with _session_lock:
        if session_id not in _session_result_queues:
            _session_result_queues[session_id] = queue.Queue()
        return _session_result_queues[session_id]


def translate_worker():
    while True:
        session_id, clause_id, text, cfg = translate_queue.get()
        try:
            trans = translate(text, cfg["src"], cfg["tgt"])
            speak_queue.put((session_id, clause_id, trans, cfg))
        except Exception as e:
            get_session_queue(session_id).put(("error", clause_id, f"翻译出错：{e}"))
        finally:
            translate_queue.task_done()


def speak_worker():
    while True:
        session_id, clause_id, trans, cfg = speak_queue.get()
        try:
            tts_path = speak(trans, cfg)
            tagged = f"[{cfg['tgt_tag']}] {trans}"
            get_session_queue(session_id).put(("done", clause_id, tagged, tts_path))
        except Exception as e:
            get_session_queue(session_id).put(("error", clause_id, f"合成出错：{e}"))
        finally:
            speak_queue.task_done()


threading.Thread(target=translate_worker, daemon=True).start()
threading.Thread(target=speak_worker, daemon=True).start()


# --- 每个浏览器会话独立的流水线状态 ------------------------------------------

@dataclass
class SessionState:
    session_id: str
    vad: VadSegmenter = field(default_factory=VadSegmenter)
    asr: Optional[StreamingClauseASR] = None
    lang_code: Optional[str] = None
    leftover: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=np.float32))
    last_partial_ts: float = 0.0
    history: list = field(default_factory=list)
    partial_index: Optional[int] = None
    clause_counter: int = 0
    pending_clause_index: dict = field(default_factory=dict)


def _update_partial_caption(session: SessionState, cfg, confirmed, tentative):
    line = f"[{cfg['src_tag']}] {confirmed}{tentative}"
    if session.partial_index is not None:
        session.history[session.partial_index]["content"] = line
    else:
        session.history.append({"role": "user", "content": line})
        session.partial_index = len(session.history) - 1


def _finalize_clause(session: SessionState, cfg):
    text = session.asr.finalize()
    if not text.strip():
        return  # VAD 误触发（比如咳嗽声、噪音），没识别出实际内容，直接跳过
    if session.partial_index is not None:
        session.history[session.partial_index]["content"] = f"[{cfg['src_tag']}] {text}"
        session.partial_index = None
    else:
        session.history.append({"role": "user", "content": f"[{cfg['src_tag']}] {text}"})

    session.clause_counter += 1
    clause_id = session.clause_counter
    session.history.append({"role": "assistant", "content": "⏳ 翻译中…"})
    session.pending_clause_index[clause_id] = len(session.history) - 1
    translate_queue.put((session.session_id, clause_id, text, cfg))


def _flush_session_results(session: SessionState):
    q = get_session_queue(session.session_id)
    latest_audio = None
    while True:
        try:
            item = q.get_nowait()
        except queue.Empty:
            break
        kind = item[0]
        if kind == "error":
            _, clause_id, msg = item
            idx = session.pending_clause_index.pop(clause_id, None)
            if idx is not None:
                session.history[idx]["content"] = f"⚠️ {msg}"
        elif kind == "done":
            _, clause_id, tagged_text, tts_path = item
            idx = session.pending_clause_index.pop(clause_id, None)
            if idx is not None:
                session.history[idx]["content"] = tagged_text
            latest_audio = tts_path
    return latest_audio


def ingest_audio_chunk(session: Optional[SessionState], mode_label: str, new_chunk):
    """
    Gradio 流式麦克风的回调，大约每 stream_every 秒被调用一次，
    每次带来一小段新录到的音频。这里做：分帧喂 VAD -> 判断是否说完一句 ->
    刷新临时字幕 -> 顺便把后台线程做完的翻译/合成结果捞回来更新界面。
    """
    if new_chunk is None:
        return session, (session.history if session else []), None

    if session is None:
        session = SessionState(session_id=uuid.uuid4().hex)

    cfg = STREAMING_LANG_MAP[mode_label]
    lang_code = "zh" if cfg["src"] == "Chinese" else "en"
    if session.asr is None or session.lang_code != lang_code:
        # 切换了翻译方向，重建识别器和 VAD 状态，避免残留上一方向的缓冲
        session.asr = StreamingClauseASR(zh_en_asr_model, lang_code)
        session.lang_code = lang_code
        session.vad = VadSegmenter()
        session.leftover = np.zeros(0, dtype=np.float32)

    if not API_KEY:
        session.history.append({"role": "assistant", "content": "⚠️ 未配置 SILICONFLOW_API_KEY"})
        return session, session.history, None

    sr, chunk = new_chunk
    audio_f32 = to_16k_mono_float32(sr, chunk)
    session.leftover = np.concatenate([session.leftover, audio_f32])

    while len(session.leftover) >= FRAME_SAMPLES:
        frame = session.leftover[:FRAME_SAMPLES]
        session.leftover = session.leftover[FRAME_SAMPLES:]
        frame_bytes = float32_to_pcm16_bytes(frame)
        status = session.vad.feed(frame_bytes)
        if session.vad.last_is_speech or session.vad.in_speech:
            session.asr.add_audio(frame)
        if status == "pause_end":
            _finalize_clause(session, cfg)

    now = time.time()
    if session.vad.in_speech and now - session.last_partial_ts >= PARTIAL_UPDATE_SEC:
        session.last_partial_ts = now
        confirmed, tentative = session.asr.partial_transcribe()
        if confirmed or tentative:
            _update_partial_caption(session, cfg, confirmed, tentative)

    tts_path = _flush_session_results(session)

    return session, session.history, (tts_path if tts_path else gr.update())


# =============================================================================
# UI
# =============================================================================

with gr.Blocks(title="传声筒") as demo:
    gr.Markdown("## 🎙️ 传声筒")

    with gr.Tab("同声传译（中↔英，持续聆听）"):
        gr.Markdown(
            "点一下开始，持续说话即可，不用点停止。程序会自动按停顿分句，"
            "边听边识别、边翻译边播放译文，通常滞后几秒。\n"
            "⚠️ 实验性功能：识别用的是本地 CPU 模型，参数（停顿判定时长、模型大小）"
            "可能需要根据你的机器实际表现调整。"
        )
        mode_stream = gr.Radio(
            choices=list(STREAMING_LANG_MAP.keys()),
            value="英语 → 中文",
            label="翻译方向",
        )
        audio_in_stream = gr.Audio(
            sources=["microphone"],
            type="numpy",
            streaming=True,
            label="持续录音",
        )
        chatbot_stream = gr.Chatbot(label="同传记录", height=460)
        tts_out_stream = gr.Audio(label="译文语音", autoplay=True, interactive=False)

        session_state = gr.State(None)

        audio_in_stream.stream(
            fn=ingest_audio_chunk,
            inputs=[session_state, mode_stream, audio_in_stream],
            outputs=[session_state, chatbot_stream, tts_out_stream],
            stream_every=0.5,
        )

        mode_stream.change(
            fn=lambda: (None, [], None),
            inputs=None,
            outputs=[session_state, chatbot_stream, tts_out_stream],
        )

    with gr.Tab("点击录音（六个方向，原来的模式）"):
        gr.Markdown(
            "点麦克风开始录音，说一句完整的话，再点一次停止，自动识别+翻译+朗读译文。\n"
            "识别与翻译主要使用硅基流动，法语识别、法语合成使用本地方案补充。"
        )
        mode = gr.Radio(choices=list(LANG_MAP.keys()), value="英语 → 中文", label="翻译方向")
        audio_in = gr.Audio(sources=["microphone"], type="filepath", label="点击录音，再点一次停止")
        chatbot = gr.Chatbot(label="传译记录（含可重听音频）", height=460)
        tts_out = gr.Audio(label="最新译文语音", autoplay=True, interactive=False)

        audio_in.stop_recording(
            fn=process,
            inputs=[audio_in, mode, chatbot],
            outputs=[chatbot, audio_in, tts_out],
        )

demo.launch()