path: root/mimo_tts.py

#!/usr/bin/env python3
"""MiMo V2.5 TTS — 独立语音合成脚本

用法:
    # 基本用法
    python3 mimo_tts.py "要合成的文字"

    # 指定音色和输出路径
    python3 mimo_tts.py "Hello world" --voice Mia --output hello.wav

    # 使用语音设计模型（通过文字描述生成音色）
    python3 mimo_tts.py "Once upon a time" --model voicedesign \
        --style "A warm storyteller voice with British accent"

    # 列出所有可用音色
    python3 mimo_tts.py --list-voices

    # 流式输出（PCM16 24kHz）
    python3 mimo_tts.py "你好世界" --stream --output output.wav

环境变量:
    XIAOMI_API_KEY    — 必需，小米 MiMo API Key
    XIAOMI_BASE_URL   — 可选，API 端点（默认 https://api.xiaomimimo.com/v1）

获取 API Key: https://mimo.mi.com
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import urllib.request
from pathlib import Path

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

API_BASE_DEFAULT = "https://api.xiaomimimo.com/v1"

VOICES = [
    {"id": "mimo_default", "display": "默认 (auto)", "lang": "auto", "gender": "-"},
    {"id": "冰糖",         "display": "冰糖 Bingtang",  "lang": "zh",  "gender": "♀"},
    {"id": "茉莉",         "display": "茉莉 Moli",      "lang": "zh",  "gender": "♀"},
    {"id": "苏打",         "display": "苏打 Suda",      "lang": "zh",  "gender": "♂"},
    {"id": "白桦",         "display": "白桦 Baihua",    "lang": "zh",  "gender": "♂"},
    {"id": "Mia",          "display": "Mia",            "lang": "en",  "gender": "♀"},
    {"id": "Chloe",        "display": "Chloe",          "lang": "en",  "gender": "♀"},
    {"id": "Milo",         "display": "Milo",           "lang": "en",  "gender": "♂"},
    {"id": "Dean",         "display": "Dean",           "lang": "en",  "gender": "♂"},
]

MODELS = {
    "tts":          "mimo-v2.5-tts",
    "voicedesign":  "mimo-v2.5-tts-voicedesign",
    "voiceclone":   "mimo-v2.5-tts-voiceclone",
}


# ---------------------------------------------------------------------------
# Core API
# ---------------------------------------------------------------------------

def get_api_base() -> str:
    return os.environ.get("XIAOMI_BASE_URL", API_BASE_DEFAULT).rstrip("/")


def synthesize(
    text: str,
    output_path: str,
    *,
    voice: str = "mimo_default",
    model: str = "mimo-v2.5-tts",
    style: str | None = None,
    api_key: str | None = None,
) -> str:
    """合成语音并写入文件，返回文件路径。"""
    api_key = api_key or os.environ.get("XIAOMI_API_KEY")
    if not api_key:
        raise RuntimeError("XIAOMI_API_KEY is not set. Get one at https://mimo.mi.com")

    messages: list = []
    if style:
        messages.append({"role": "user", "content": style})
    messages.append({"role": "assistant", "content": text})

    audio_config: dict = {"format": "wav"}
    if "voicedesign" not in model and "voiceclone" not in model:
        audio_config["voice"] = voice

    body = {
        "model": model,
        "messages": messages,
        "audio": audio_config,
    }

    data = json.dumps(body).encode()
    req = urllib.request.Request(
        f"{get_api_base()}/chat/completions",
        data=data,
        headers={"Content-Type": "application/json", "api-key": api_key},
        method="POST",
    )

    with urllib.request.urlopen(req, timeout=60) as resp:
        resp_data = json.loads(resp.read().decode())

    audio_b64 = resp_data["choices"][0]["message"]["audio"]["data"]
    audio_bytes = base64.b64decode(audio_b64)

    output_path = str(Path(output_path).expanduser())
    if not output_path.endswith(".wav"):
        output_path = os.path.splitext(output_path)[0] + ".wav"
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)

    with open(output_path, "wb") as f:
        f.write(audio_bytes)

    return output_path


def synthesize_stream(
    text: str,
    *,
    voice: str = "mimo_default",
    model: str = "mimo-v2.5-tts",
    style: str | None = None,
    api_key: str | None = None,
):
    """流式合成，yield PCM16 字节块 (24kHz mono)。"""
    api_key = api_key or os.environ.get("XIAOMI_API_KEY")
    if not api_key:
        raise RuntimeError("XIAOMI_API_KEY is not set.")

    messages: list = []
    if style:
        messages.append({"role": "user", "content": style})
    messages.append({"role": "assistant", "content": text})

    audio_config: dict = {"format": "pcm16"}
    if "voicedesign" not in model and "voiceclone" not in model:
        audio_config["voice"] = voice

    body = {"model": model, "messages": messages, "audio": audio_config, "stream": True}
    data = json.dumps(body).encode()
    req = urllib.request.Request(
        f"{get_api_base()}/chat/completions",
        data=data,
        headers={"Content-Type": "application/json", "api-key": api_key},
        method="POST",
    )

    with urllib.request.urlopen(req, timeout=120) as resp:
        for line in resp:
            line = line.strip()
            if not line or not line.startswith(b"data: "):
                continue
            payload = line[6:]
            if payload == b"[DONE]":
                break
            try:
                chunk = json.loads(payload)
                audio_data = chunk["choices"][0]["delta"].get("audio", {})
                if isinstance(audio_data, dict) and "data" in audio_data:
                    yield base64.b64decode(audio_data["data"])
            except (json.JSONDecodeError, KeyError, IndexError):
                continue


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def list_voices():
    print(f"{'ID':<15} {'Display':<20} {'Lang':<6} {'Gender'}")
    print("-" * 50)
    for v in VOICES:
        print(f"{v['id']:<15} {v['display']:<20} {v['lang']:<6} {v['gender']}")


def main():
    parser = argparse.ArgumentParser(
        description="MiMo V2.5 TTS — 语音合成",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("text", nargs="?", help="要合成的文字")
    parser.add_argument("-v", "--voice", default="mimo_default",
                        help="音色 ID (默认: mimo_default)")
    parser.add_argument("-m", "--model", default="tts",
                        choices=["tts", "voicedesign", "voiceclone"],
                        help="模型 (默认: tts)")
    parser.add_argument("-s", "--style", default=None,
                        help="风格/语气描述 (user message)")
    parser.add_argument("-o", "--output", default=None,
                        help="输出文件路径 (默认: ~/mimo_tts_output.wav)")
    parser.add_argument("--stream", action="store_true",
                        help="流式模式 (输出 PCM16 WAV)")
    parser.add_argument("--list-voices", action="store_true",
                        help="列出所有可用音色")
    parser.add_argument("--list-models", action="store_true",
                        help="列出所有可用模型")

    args = parser.parse_args()

    if args.list_voices:
        list_voices()
        return
    if args.list_models:
        for key, mid in MODELS.items():
            print(f"  {key:<15} → {mid}")
        return

    if not args.text:
        parser.error("请提供要合成的文字")

    model_id = MODELS[args.model]
    output = args.output or os.path.expanduser("~/mimo_tts_output.wav")

    if args.stream:
        print(f"流式合成中... voice={args.voice} model={model_id}")
        collected = bytearray()
        for chunk in synthesize_stream(args.text, voice=args.voice, model=model_id, style=args.style):
            collected.extend(chunk)
            print(f"  收到 {len(chunk)} 字节", end="\r")
        # Write as WAV (PCM16 24kHz mono)
        import struct, wave
        output = output if output.endswith(".wav") else output + ".wav"
        with wave.open(output, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(24000)
            wf.writeframes(bytes(collected))
        print(f"\n✅ 已保存: {output} ({len(collected)} 字节 PCM)")
    else:
        print(f"合成中... voice={args.voice} model={model_id}")
        result = synthesize(args.text, output, voice=args.voice, model=model_id, style=args.style)
        size_kb = os.path.getsize(result) / 1024
        print(f"✅ 已保存: {result} ({size_kb:.1f} KB)")


if __name__ == "__main__":
    main()