first commit

This commit is contained in:
2026-08-15 14:43:56 +08:00
commit 63b8c28e1d
340 changed files with 59515 additions and 0 deletions
+28
View File
@@ -0,0 +1,28 @@
# Python
__pycache__/
*.py[cod]
.venv/
runtime/
config/.venv_python_path.txt
# Node / frontend build
frontend/admin/node_modules/
frontend/admin/.vite/
# Build output
build/
dist/
*.spec
# Runtime data
logs/
data/music_state.json
data/queue_state.json
data/song_requests.json
data/tts_state.json
data/bilibili_credentials.json
web/music_cover.jpg
# IDE
.idea/
.workbuddy/
+206
View File
@@ -0,0 +1,206 @@
# BetterGI 直播联动
监听 B 站直播间弹幕,管理观众排队、积分、扫码上号、BetterGI 配置组执行、点歌和 TTS 播报。
## 环境
- **Python**: conda 环境 `Live-streaming``E:\Programs\Anaconda3\envs\Live-streaming\python.exe`
- **Node.js**: 构建前端用
- **BetterGI**: 自动化脚本引擎,路径在 `config/config.json` 中配置
```powershell
# 安装 Python 依赖
E:\Programs\Anaconda3\envs\Live-streaming\python.exe -m pip install -r requirements.txt
# 下载 TTS 模型(约 1.2GB,仅需一次)
E:\Programs\Anaconda3\envs\Live-streaming\python.exe -c "from huggingface_hub import snapshot_download; snapshot_download('Qwen/Qwen3-TTS-12Hz-0.6B-Base')"
# 安装前端依赖并构建
cd frontend\admin
npm install
npm run build
```
## 启动
```powershell
.venv\Scripts\python.exe app\main.py --role all --host 0.0.0.0 --port 5191
```
| 参数 | 默认值 | 说明 |
|------|--------|------|
| `--role` | `all` | `queue` 仅排队/Web / `music` 仅音乐监听 / `tts` TTS 状态窗口 / `all` 同时启动全部(推荐,否则音乐页面不会随 SMTC 更新) |
| `--host` | `0.0.0.0` | 绑定地址 |
| `--port` | `8086` | 端口 |
启动后日志会打印本机和局域网访问地址,例如 `http://192.168.31.80:5191/admin`
## 配置文件
配置文件统一为 `config/config.json`,支持热加载,保存后无需重启。
### bilibili — 直播间连接
| 字段 | 说明 |
|------|------|
| `room_id` | 直播间号(短号也可) |
| `sessdata` | B 站 Cookie SESSDATA,用于监听弹幕和发送回复 |
| `bili_jct` | CSRF token,发弹幕必须。和 SESSDATA 配套获取 |
获取方式:浏览器登录 B 站 → F12 → Application → Cookies → bilibili.com → 复制 `SESSDATA``bili_jct`
> **注意**`sessdata` 和 `bili_jct` 属于敏感登录凭证,不要分享本文件。缺少 `bili_jct` 时只能收弹幕不能回复。
### queue — 排队积分
| 字段 | 默认值 | 说明 |
|------|--------|------|
| `initial_points` | 10 | 新用户初始积分 |
| `signin_points` | 10 | 每日签到获得积分 |
| `max_points` | 30 | 积分上限 |
| `points_per_minute` | 1 | 队首每分钟扣除积分 |
| `admin_window_seconds` | 180 | 队首上号窗口(超时过号) |
| `default_group` | 薄荷 | 队列空时默认运行的 BetterGI 配置组 |
### global — 全局
| 字段 | 说明 |
|------|------|
| `admin_uids` | 管理员 UID 列表(一级用户,可用重置等特权指令) |
| `log_level` | 日志级别:`DEBUG` / `INFO` / `WARNING` |
### broadcast — 弹幕播报与 TTS
| 字段 | 默认值 | 说明 |
|------|--------|------|
| `enable_danmu_reply` | true | 指令回复发到直播间 |
| `enable_system_danmu` | true | 系统通知发到直播间 |
| `enable_tts` | false | TTS 语音播报(需先下载模型) |
| `danmu_interval_sec` | 5 | 弹幕发送间隔 |
| `tts_provider` | none | `none` / `faster-qwen3-tts` / `dots-tts` |
### commands — 内置指令别名
11 个内置指令,每个可单独启用/禁用、自定义别名:
| key | 默认别名 | 功能 |
|-----|----------|------|
| `queue` | 排队 | 加入排队队列 |
| `signin` | 签到 | 每日签到 |
| `login` | 上号 | 队首触发扫码上号 |
| `confirm_yes` | 是 | 确认账号正确 |
| `confirm_no` | 不是 | 确认账号不正确,重新扫码 |
| `run` | 执行, 跑, 开始 | 执行配置组(需带参数) |
| `leave` | 退出 | 退出队列 |
| `reset` | 重置 | 一级用户重启原神和 BGI |
| `points` | 积分 | 查询积分 |
| `queue_list` | 队列 | 查看排队情况 |
| `help` | 帮助 | 显示帮助 |
### rules — 弹幕触发规则
```json
{
"keyword": "开始",
"match_type": "exact",
"groups": ["子探测单元"],
"cooldown": 60,
"admin_only": false,
"reply": "收到,开始执行子探测单元任务"
}
```
- `match_type``contains` / `exact` / `startswith` / `regex`
- 匹配后直接启动 `groups` 指定的 BetterGI 配置组
### music_monitor — 音乐监听与点歌
监听 Windows 系统媒体会话(SMTC),观众可通过弹幕点歌。
`request_player.commands`:点歌触发词,默认 `["点歌", "dg"]`
### system — 系统定时任务
| 字段 | 说明 |
|------|------|
| `enable_startup_shortcut` | 开机自启 |
| `auto_reboot_time` | 每日自动重启时间 |
| `launch_bilibili_live_time` | 定时启动直播姬 |
| `launch_genshin_time` | 定时启动原神 |
## 后台管理
```
http://127.0.0.1:5191/admin
```
| 页面 | 功能 |
|------|------|
| 总览 | 运行状态、排队人数、服务健康、快捷操作 |
| 配置 | B 站连接、BetterGI 路径、队列参数、前台视觉 |
| 规则 | 弹幕触发规则 + 内置指令别名管理 |
| 队列 | 当前排队列表,支持移出、加减分、清空 |
| 用户 | 所有用户积分管理,支持加减分、踢出、删除 |
| 媒体 | 弹幕播报开关、TTS 配置与测试、音乐监听 |
| 日志 | 系统日志 + BetterGI 日志 |
| JSON | 原始 JSON 编辑 |
后台每 3 秒自动刷新,数据实时同步。
## 开发
### 前端
```powershell
cd frontend\admin
npm run build # 构建到 web/admin/
```
源码:`frontend/admin/src/main.js`Vue3 单文件组件),`styles.css`
### 后端
```powershell
# 编译检查
E:\Programs\Anaconda3\envs\Live-streaming\python.exe -m py_compile app/danmu_queue.py
```
核心文件:
| 文件 | 功能 |
|------|------|
| `app/main.py` | 入口,启动 Web 服务和子模块 |
| `app/danmu_queue.py` | 弹幕监听、指令处理、队列管理、BGI 控制、TTS、Web API |
| `app/music_monitor.py` | SMTC 音乐监听与点歌调度 |
| `app/core/runtime_paths.py` | 运行时路径解析 |
## 打包
```powershell
# 轻量包(不含 TTS
.\build.bat
# 完整包(含 TTS/GPU 依赖)
.\build.bat -FullTts
```
输出:`dist/LiveStreaming/LiveStreaming.exe`
## 数据文件
| 路径 | 内容 |
|------|------|
| `data/users.json` | 用户账号与积分 |
| `data/queue_state.json` | 排队状态(启动时清空) |
| `data/ref_audio.wav` | TTS 参考音频 |
| `data/music_state.json` | 音乐播放状态 |
| `data/song_requests.json` | 点歌队列 |
| `logs/danmu_queue.log` | 主程序日志 |
## 弹幕链路排查
1. 确认 `config.json``room_id` 正确(浏览器打开直播间,URL 中数字)
2. 确认 `sessdata` 有效:看日志是否有 `SESSDATA有效, 账号=xxx`
3. 确认 WebSocket 连接:看日志是否有 `认证成功,开始监听弹幕`
4. 发弹幕测试:在直播间发 `帮助`,日志应出现 `[弹幕] 用户(uid): 帮助`
5. 如连接断线(`WinError 10013`),用管理员权限启动程序
6. DEBUG 统计:日志每 60 秒输出 `[弹幕统计] 收到N个raw | 解析N个包 | OP_MESSAGE=N | 提取弹幕=N`
+154
View File
@@ -0,0 +1,154 @@
from __future__ import annotations
import base64
import hashlib
import hmac
import json
import secrets
import time
from datetime import datetime
from pathlib import Path
from typing import Any
SESSION_COOKIE_NAME = "live_admin_session"
def _utc_now_iso() -> str:
return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
def _is_local_ip(ip: str) -> bool:
ip = (ip or "").strip()
return ip in {"127.0.0.1", "::1", "::ffff:127.0.0.1", "localhost"}
class AdminAuthManager:
def __init__(self, data_dir: str | Path, logger, *, session_ttl_sec: int = 12 * 60 * 60):
self.data_dir = Path(data_dir)
self.data_dir.mkdir(parents=True, exist_ok=True)
self.path = self.data_dir / "admin_auth.json"
self.audit_path = self.data_dir / "admin_audit.log"
self.logger = logger
self.session_ttl_sec = int(session_ttl_sec)
self._sessions: dict[str, dict[str, Any]] = {}
self._data = self._default_data()
self._load()
def _default_data(self) -> dict[str, Any]:
return {
"version": 1,
"created_at": "",
"password_hash": "",
"password_salt": "",
"password_updated_at": "",
}
def _load(self) -> None:
if not self.path.exists():
return
try:
data = json.loads(self.path.read_text(encoding="utf-8"))
if isinstance(data, dict):
self._data.update(data)
except Exception as exc:
self.logger.warning(f"[后台认证] 读取认证文件失败: {exc}")
def _save(self) -> None:
tmp = self.path.with_suffix(".tmp")
tmp.write_text(json.dumps(self._data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
tmp.replace(self.path)
def _hash_password(self, password: str, salt: bytes) -> str:
digest = hashlib.scrypt(
password.encode("utf-8"),
salt=salt,
n=2 ** 14,
r=8,
p=1,
dklen=64,
)
return base64.b64encode(digest).decode("ascii")
def is_bootstrapped(self) -> bool:
return bool(self._data.get("password_hash") and self._data.get("password_salt"))
def can_bootstrap_ip(self, ip: str) -> bool:
return _is_local_ip(ip)
def bootstrap(self, password: str) -> None:
password = str(password or "")
if len(password) < 8:
raise ValueError("后台密码至少 8 位")
salt = secrets.token_bytes(16)
now = _utc_now_iso()
self._data = {
"version": 1,
"created_at": self._data.get("created_at") or now,
"password_hash": self._hash_password(password, salt),
"password_salt": base64.b64encode(salt).decode("ascii"),
"password_updated_at": now,
}
self._sessions.clear()
self._save()
def verify_password(self, password: str) -> bool:
if not self.is_bootstrapped():
return False
try:
salt = base64.b64decode(self._data["password_salt"])
except Exception:
return False
current = self._hash_password(str(password or ""), salt)
return hmac.compare_digest(current, self._data.get("password_hash", ""))
def create_session(self, client_ip: str, user_agent: str = "") -> tuple[str, int]:
self.cleanup_sessions()
token = secrets.token_urlsafe(32)
now = int(time.time())
expires_at = now + self.session_ttl_sec
self._sessions[token] = {
"created_at": now,
"expires_at": expires_at,
"client_ip": client_ip or "",
"user_agent": (user_agent or "")[:240],
}
return token, expires_at
def cleanup_sessions(self) -> None:
now = int(time.time())
expired = [token for token, session in self._sessions.items() if int(session.get("expires_at", 0)) <= now]
for token in expired:
self._sessions.pop(token, None)
def get_session(self, token: str) -> dict[str, Any] | None:
self.cleanup_sessions()
if not token:
return None
session = self._sessions.get(token)
if not session:
return None
session["expires_at"] = int(time.time()) + self.session_ttl_sec
return session
def destroy_session(self, token: str) -> None:
if token:
self._sessions.pop(token, None)
def write_audit(self, *, action: str, target: str = "", client_ip: str = "", session_id: str = "", detail: str = "") -> None:
entry = {
"at": _utc_now_iso(),
"action": action,
"target": target,
"client_ip": client_ip,
"session_id": session_id[:12],
"detail": detail,
}
line = json.dumps(entry, ensure_ascii=False)
try:
with open(self.audit_path, "a", encoding="utf-8") as handle:
handle.write(line + "\n")
except Exception as exc:
self.logger.warning(f"[后台认证] 写审计日志失败: {exc}")
self.logger.info(f"[审计] action={action} target={target or '-'} ip={client_ip or '-'} detail={detail or '-'}")
+43
View File
@@ -0,0 +1,43 @@
from __future__ import annotations
import asyncio
import time
class AdminEventBus:
def __init__(self):
self._subscribers: set[asyncio.Queue] = set()
self._counter = 0
self._lock = asyncio.Lock()
async def subscribe(self) -> asyncio.Queue:
queue: asyncio.Queue = asyncio.Queue(maxsize=64)
async with self._lock:
self._subscribers.add(queue)
return queue
async def unsubscribe(self, queue: asyncio.Queue) -> None:
async with self._lock:
self._subscribers.discard(queue)
def subscriber_count(self) -> int:
return len(self._subscribers)
async def publish(self, event_type: str, payload):
self._counter += 1
event = {
"id": self._counter,
"type": event_type,
"ts": time.time(),
"payload": payload,
}
for queue in list(self._subscribers):
if queue.full():
try:
queue.get_nowait()
except asyncio.QueueEmpty:
pass
try:
queue.put_nowait(event)
except asyncio.QueueFull:
pass
+601
View File
@@ -0,0 +1,601 @@
from __future__ import annotations
import asyncio
import base64
import ctypes
import ctypes.wintypes
import hashlib
import http.cookiejar
import http.cookies
import io
import json
import logging
import os
import secrets
import time
import urllib.parse
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Callable
_PUBLIC_KEY_PEM = """-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDLgd2OAkcGVtoE3ThUREbio0Eg
Uc/prcajMKXvkCKFCWhJYJcLkcM2DKKcSeFpD/j6Boy538YXnR6VhcuUJOhH2x71
nzPjfdTcqMz7djHum0qSZA0AyCBDABUqCrfNgCiJ00Ra7GmRj+YCK1NJEuewlb40
JNrRuoEUXpabUzGB8QIDAQAB
-----END PUBLIC KEY-----"""
_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
_QR_GENERATE_URL = "https://passport.bilibili.com/x/passport-login/web/qrcode/generate"
_QR_POLL_URL = "https://passport.bilibili.com/x/passport-login/web/qrcode/poll"
_QR_HEADERS = {
"Referer": "https://www.bilibili.com/",
"Origin": "https://www.bilibili.com",
}
class _DataBlob(ctypes.Structure):
_fields_ = [("cbData", ctypes.wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]
def _blob(data: bytes) -> tuple[_DataBlob, Any]:
buffer = ctypes.create_string_buffer(data)
return _DataBlob(len(data), ctypes.cast(buffer, ctypes.POINTER(ctypes.c_byte))), buffer
def _dpapi_encrypt(value: str) -> str:
if os.name != "nt":
raise RuntimeError("B站刷新令牌安全存储仅支持 Windows DPAPI")
source, source_buffer = _blob(value.encode("utf-8"))
entropy, entropy_buffer = _blob(b"Live-streaming:bilibili-refresh-token:v1")
output = _DataBlob()
ok = ctypes.windll.crypt32.CryptProtectData(
ctypes.byref(source), None, ctypes.byref(entropy), None, None, 0,
ctypes.byref(output),
)
_ = source_buffer, entropy_buffer
if not ok:
raise ctypes.WinError()
try:
encrypted = ctypes.string_at(output.pbData, output.cbData)
return base64.b64encode(encrypted).decode("ascii")
finally:
ctypes.windll.kernel32.LocalFree(output.pbData)
def _dpapi_decrypt(value: str) -> str:
if os.name != "nt":
raise RuntimeError("B站刷新令牌安全存储仅支持 Windows DPAPI")
source, source_buffer = _blob(base64.b64decode(value))
entropy, entropy_buffer = _blob(b"Live-streaming:bilibili-refresh-token:v1")
output = _DataBlob()
ok = ctypes.windll.crypt32.CryptUnprotectData(
ctypes.byref(source), None, ctypes.byref(entropy), None, None, 0,
ctypes.byref(output),
)
_ = source_buffer, entropy_buffer
if not ok:
raise ctypes.WinError()
try:
return ctypes.string_at(output.pbData, output.cbData).decode("utf-8")
finally:
ctypes.windll.kernel32.LocalFree(output.pbData)
class BilibiliCredentialStore:
def __init__(self, path: str | Path):
self.path = Path(path)
def save_refresh_token(self, refresh_token: str) -> None:
token = str(refresh_token or "").strip()
if not token:
raise ValueError("refresh_token 不能为空")
payload = {
"version": 1,
"provider": "windows_dpapi_current_user",
"refresh_token_protected": _dpapi_encrypt(token),
}
self.path.parent.mkdir(parents=True, exist_ok=True)
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
with open(tmp, "w", encoding="utf-8", newline="\n") as handle:
json.dump(payload, handle, ensure_ascii=False, indent=2)
handle.write("\n")
tmp.replace(self.path)
def load_refresh_token(self) -> str:
if not self.path.exists():
return ""
with open(self.path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
protected = str(payload.get("refresh_token_protected") or "")
return _dpapi_decrypt(protected) if protected else ""
def is_configured(self) -> bool:
if not self.path.exists():
return False
try:
with open(self.path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
return bool(str(payload.get("refresh_token_protected") or ""))
except (OSError, ValueError, TypeError):
return False
def _read_der_length(data: bytes, offset: int) -> tuple[int, int]:
first = data[offset]
offset += 1
if first < 0x80:
return first, offset
count = first & 0x7F
return int.from_bytes(data[offset:offset + count], "big"), offset + count
def _read_der_tlv(data: bytes, offset: int, expected_tag: int | None = None) -> tuple[int, bytes, int]:
tag = data[offset]
if expected_tag is not None and tag != expected_tag:
raise ValueError(f"DER tag 不匹配: expected={expected_tag:#x}, actual={tag:#x}")
length, content_offset = _read_der_length(data, offset + 1)
end = content_offset + length
return tag, data[content_offset:end], end
def _public_numbers() -> tuple[int, int]:
body = "".join(line for line in _PUBLIC_KEY_PEM.splitlines() if not line.startswith("-----"))
der = base64.b64decode(body)
_, spki, _ = _read_der_tlv(der, 0, 0x30)
_, _, offset = _read_der_tlv(spki, 0, 0x30)
_, bit_string, _ = _read_der_tlv(spki, offset, 0x03)
_, rsa_key, _ = _read_der_tlv(bit_string[1:], 0, 0x30)
_, modulus_bytes, rsa_offset = _read_der_tlv(rsa_key, 0, 0x02)
_, exponent_bytes, _ = _read_der_tlv(rsa_key, rsa_offset, 0x02)
return int.from_bytes(modulus_bytes, "big"), int.from_bytes(exponent_bytes, "big")
def _mgf1(seed: bytes, length: int) -> bytes:
result = bytearray()
counter = 0
while len(result) < length:
result.extend(hashlib.sha256(seed + counter.to_bytes(4, "big")).digest())
counter += 1
return bytes(result[:length])
def _rsa_oaep_sha256_encrypt(message: bytes) -> str:
modulus, exponent = _public_numbers()
key_size = (modulus.bit_length() + 7) // 8
digest_size = hashlib.sha256().digest_size
if len(message) > key_size - 2 * digest_size - 2:
raise ValueError("待加密内容过长")
label_hash = hashlib.sha256(b"").digest()
padding = b"\x00" * (key_size - len(message) - 2 * digest_size - 2)
data_block = label_hash + padding + b"\x01" + message
seed = secrets.token_bytes(digest_size)
data_mask = _mgf1(seed, key_size - digest_size - 1)
masked_data = bytes(left ^ right for left, right in zip(data_block, data_mask))
seed_mask = _mgf1(masked_data, digest_size)
masked_seed = bytes(left ^ right for left, right in zip(seed, seed_mask))
encoded = b"\x00" + masked_seed + masked_data
encrypted = pow(int.from_bytes(encoded, "big"), exponent, modulus)
return encrypted.to_bytes(key_size, "big").hex()
def _parse_cookie(cookie_text: str) -> dict[str, str]:
parsed = http.cookies.SimpleCookie()
parsed.load(str(cookie_text or "").replace("; ", ";"))
return {name: morsel.value for name, morsel in parsed.items()}
def _cookie_header(values: dict[str, str]) -> str:
return "; ".join(f"{name}={value}" for name, value in values.items() if value)
def _request_json(url: str, *, cookie: str = "", data: dict[str, str] | None = None,
opener: urllib.request.OpenerDirector | None = None,
headers: dict[str, str] | None = None,
retries: int = 0) -> tuple[dict, Any]:
body = urllib.parse.urlencode(data).encode("utf-8") if data is not None else None
request = urllib.request.Request(url, data=body, method="POST" if body is not None else "GET")
request.add_header("User-Agent", _USER_AGENT)
for name, value in (headers or {}).items():
request.add_header(name, value)
if cookie:
request.add_header("Cookie", cookie)
if body is not None:
request.add_header("Content-Type", "application/x-www-form-urlencoded")
retry_count = max(0, int(retries))
for attempt in range(retry_count + 1):
try:
response = (opener or urllib.request.build_opener()).open(request, timeout=15)
return json.loads(response.read().decode("utf-8")), response
except urllib.error.HTTPError:
raise
except (urllib.error.URLError, ConnectionError, TimeoutError, OSError) as exc:
if attempt >= retry_count:
raise RuntimeError("连接B站登录服务失败,请稍后重试") from exc
time.sleep(0.4 * (attempt + 1))
raise AssertionError("unreachable")
def _seed_cookie_jar(jar: http.cookiejar.CookieJar, values: dict[str, str]) -> None:
for name, value in values.items():
jar.set_cookie(http.cookiejar.Cookie(
version=0, name=name, value=value, port=None, port_specified=False,
domain=".bilibili.com", domain_specified=True, domain_initial_dot=True,
path="/", path_specified=True, secure=False, expires=None, discard=True,
comment=None, comment_url=None, rest={}, rfc2109=False,
))
def _jar_values(jar: http.cookiejar.CookieJar) -> dict[str, str]:
return {cookie.name: cookie.value for cookie in jar}
def _login_url_cookie_values(url: str) -> dict[str, str]:
values: dict[str, str] = {}
query = urllib.parse.urlparse(str(url or "")).query
for item in query.split("&"):
raw_name, separator, raw_value = item.partition("=")
if not separator:
continue
name = urllib.parse.unquote_plus(raw_name)
if name in {"SESSDATA", "bili_jct", "DedeUserID", "DedeUserID__ckMd5", "sid", "buvid3"}:
values[name] = raw_value
return values
def _render_qr_png(content: str) -> bytes:
try:
import qrcode
from qrcode.constants import ERROR_CORRECT_M
except ImportError as exc:
raise RuntimeError("缺少 qrcode 依赖,请重新安装 requirements.txt") from exc
qr = qrcode.QRCode(
version=None,
error_correction=ERROR_CORRECT_M,
box_size=8,
border=3,
)
qr.add_data(content)
qr.make(fit=True)
image = qr.make_image(fill_color="black", back_color="white")
output = io.BytesIO()
image.save(output, format="PNG")
return output.getvalue()
class BilibiliQrLogin:
"""服务端持有二维码密钥和 CookieJar,前端只获取二维码图片与状态。"""
_STATUS_MESSAGES = {
"idle": "尚未开始扫码登录",
"awaiting_scan": "请使用哔哩哔哩客户端扫码",
"awaiting_confirm": "已扫码,请在手机上确认登录",
"completed": "登录成功,Cookie 与刷新凭据已更新",
"expired": "二维码已过期,请重新生成",
"failed": "扫码登录失败",
}
def __init__(
self,
*,
credential_store: BilibiliCredentialStore,
update_cookie: Callable[[dict[str, str]], None],
logger: logging.Logger,
on_logged_in: Callable[[], Any] | None = None,
ttl_seconds: int = 180,
):
self.credential_store = credential_store
self.update_cookie = update_cookie
self.logger = logger
self.on_logged_in = on_logged_in
self.ttl_seconds = max(60, int(ttl_seconds))
self._lock = asyncio.Lock()
self._session: dict[str, Any] | None = None
def _snapshot(self) -> dict[str, Any]:
session = self._session or {}
state = str(session.get("state") or "idle")
expires_at = float(session.get("expires_at") or 0)
expires_in = max(0, int(expires_at - time.time())) if expires_at else 0
return {
"success": True,
"state": state,
"message": str(session.get("message") or self._STATUS_MESSAGES.get(state, "")),
"expires_in": expires_in,
"has_qr_image": state in {"awaiting_scan", "awaiting_confirm"} and expires_in > 0,
"credential_configured": self.credential_store.is_configured(),
"account": session.get("account"),
}
def snapshot(self) -> dict[str, Any]:
return self._snapshot()
def _start_sync(self) -> dict[str, Any]:
jar = http.cookiejar.CookieJar()
opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
payload, _ = _request_json(
_QR_GENERATE_URL,
opener=opener,
headers=_QR_HEADERS,
retries=2,
)
if payload.get("code") != 0:
raise RuntimeError(f"B站二维码申请失败 code={payload.get('code')}")
data = payload.get("data") or {}
qr_url = str(data.get("url") or "").strip()
qr_key = str(data.get("qrcode_key") or "").strip()
if not qr_url or not qr_key:
raise RuntimeError("B站二维码响应缺少必要字段")
now = time.time()
return {
"state": "awaiting_scan",
"message": self._STATUS_MESSAGES["awaiting_scan"],
"created_at": now,
"expires_at": now + self.ttl_seconds,
"qr_url": qr_url,
"qr_key": qr_key,
"jar": jar,
"opener": opener,
"account": None,
}
async def start(self) -> dict[str, Any]:
async with self._lock:
try:
self._session = await asyncio.to_thread(self._start_sync)
except Exception as exc:
self._session = {
"state": "failed",
"message": str(exc) or type(exc).__name__,
"expires_at": 0,
}
self.logger.warning("[B站扫码登录] 二维码申请失败: %s", type(exc).__name__)
return self._snapshot()
self.logger.info("[B站扫码登录] 二维码已生成,等待扫码")
return self._snapshot()
def _poll_sync(self, session: dict[str, Any]) -> dict[str, Any]:
url = _QR_POLL_URL + "?" + urllib.parse.urlencode({"qrcode_key": session["qr_key"]})
payload, _ = _request_json(
url,
opener=session["opener"],
headers=_QR_HEADERS,
retries=2,
)
if payload.get("code") != 0:
return {"state": "failed", "message": f"B站扫码状态查询失败 code={payload.get('code')}"}
data = payload.get("data") or {}
status_code = int(data.get("code") or 0)
if status_code == 86101:
return {"state": "awaiting_scan", "message": self._STATUS_MESSAGES["awaiting_scan"]}
if status_code == 86090:
return {"state": "awaiting_confirm", "message": self._STATUS_MESSAGES["awaiting_confirm"]}
if status_code == 86038:
return {"state": "expired", "message": self._STATUS_MESSAGES["expired"]}
if status_code != 0:
return {"state": "failed", "message": str(data.get("message") or f"扫码失败 code={status_code}")}
refresh_token = str(data.get("refresh_token") or "").strip()
if not refresh_token:
return {"state": "failed", "message": "扫码成功响应缺少 refresh_token"}
cookie_values = _jar_values(session["jar"])
for name, value in _login_url_cookie_values(str(data.get("url") or "")).items():
cookie_values.setdefault(name, value)
if not cookie_values.get("SESSDATA") or not cookie_values.get("bili_jct"):
return {"state": "failed", "message": "扫码成功但登录 Cookie 不完整"}
cookie = _cookie_header(cookie_values)
nav, _ = _request_json(
"https://api.bilibili.com/x/web-interface/nav",
cookie=cookie,
headers={"Referer": "https://www.bilibili.com/"},
retries=2,
)
nav_data = nav.get("data") or {}
if nav.get("code") != 0 or not bool(nav_data.get("isLogin")):
return {"state": "failed", "message": "扫码 Cookie 登录验证失败"}
self.credential_store.save_refresh_token(refresh_token)
self.update_cookie(cookie_values)
return {
"state": "completed",
"message": self._STATUS_MESSAGES["completed"],
"account": {
"mid": str(nav_data.get("mid") or ""),
"uname": str(nav_data.get("uname") or "B站账号"),
},
}
async def poll(self) -> dict[str, Any]:
callback_needed = False
async with self._lock:
if not self._session:
return self._snapshot()
state = str(self._session.get("state") or "idle")
if state in {"completed", "expired", "failed"}:
return self._snapshot()
if time.time() >= float(self._session.get("expires_at") or 0):
self._session.update(state="expired", message=self._STATUS_MESSAGES["expired"])
return self._snapshot()
try:
result = await asyncio.to_thread(self._poll_sync, self._session)
except Exception as exc:
self.logger.warning("[B站扫码登录] 状态查询异常: %s", type(exc).__name__)
self._session.update(
state="failed",
message=str(exc) or f"扫码状态查询异常: {type(exc).__name__}",
)
return self._snapshot()
previous_state = state
self._session.update(result)
callback_needed = previous_state != "completed" and result.get("state") == "completed"
snapshot = self._snapshot()
if callback_needed:
self.logger.info("[B站扫码登录] 登录成功,Cookie 与刷新凭据已更新")
if self.on_logged_in:
callback_result = self.on_logged_in()
if asyncio.iscoroutine(callback_result):
await callback_result
return snapshot
async def qr_png(self) -> bytes:
async with self._lock:
if not self._session or self._session.get("state") not in {"awaiting_scan", "awaiting_confirm"}:
raise RuntimeError("当前没有可用的登录二维码")
if time.time() >= float(self._session.get("expires_at") or 0):
self._session.update(state="expired", message=self._STATUS_MESSAGES["expired"])
raise RuntimeError("登录二维码已过期")
content = str(self._session.get("qr_url") or "")
return await asyncio.to_thread(_render_qr_png, content)
class BilibiliCookieRefresher:
def __init__(
self,
*,
credential_store: BilibiliCredentialStore,
get_cookie: Callable[[], str],
update_cookie: Callable[[dict[str, str]], None],
logger: logging.Logger,
check_interval_seconds: int = 6 * 60 * 60,
on_refreshed: Callable[[], Any] | None = None,
is_enabled: Callable[[], bool] | None = None,
get_check_interval_seconds: Callable[[], int] | None = None,
):
self.credential_store = credential_store
self.get_cookie = get_cookie
self.update_cookie = update_cookie
self.logger = logger
self.check_interval_seconds = max(3600, int(check_interval_seconds))
self.on_refreshed = on_refreshed
self.is_enabled = is_enabled or (lambda: True)
self.get_check_interval_seconds = get_check_interval_seconds
self._stop = False
self._wake = asyncio.Event()
def stop(self) -> None:
self._stop = True
self._wake.set()
def wake(self) -> None:
self._wake.set()
def _current_interval(self) -> int:
if not self.get_check_interval_seconds:
return self.check_interval_seconds
try:
return max(3600, int(self.get_check_interval_seconds()))
except (TypeError, ValueError):
return self.check_interval_seconds
def _check_and_refresh_sync(self) -> dict[str, Any]:
refresh_token = self.credential_store.load_refresh_token()
if not refresh_token:
return {"status": "disabled", "message": "未配置刷新令牌"}
current_cookie = self.get_cookie()
current_values = _parse_cookie(current_cookie)
csrf = current_values.get("bili_jct", "")
if not current_values.get("SESSDATA") or not csrf:
return {"status": "failed", "message": "当前 Cookie 缺少 SESSDATA 或 bili_jct"}
info, _ = _request_json(
"https://passport.bilibili.com/x/passport-login/web/cookie/info?" +
urllib.parse.urlencode({"csrf": csrf}),
cookie=current_cookie,
retries=2,
)
if info.get("code") != 0:
return {"status": "failed", "message": f"登录状态检查失败 code={info.get('code')}"}
if not bool((info.get("data") or {}).get("refresh")):
return {"status": "valid", "message": "Cookie 当前无需刷新"}
timestamp = str((info.get("data") or {}).get("timestamp") or "")
correspond_path = _rsa_oaep_sha256_encrypt(f"refresh_{timestamp}".encode("utf-8"))
request = urllib.request.Request(
f"https://www.bilibili.com/correspond/1/{correspond_path}",
headers={"User-Agent": _USER_AGENT, "Cookie": current_cookie},
)
html = urllib.request.urlopen(request, timeout=15).read().decode("utf-8", errors="replace")
match = __import__("re").search(r'<div\s+id=["\']1-name["\']>([^<]+)</div>', html)
if not match:
return {"status": "failed", "message": "未获取到 refresh_csrf"}
refresh_csrf = match.group(1).strip()
old_refresh_token = refresh_token
jar = http.cookiejar.CookieJar()
_seed_cookie_jar(jar, current_values)
opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
refreshed, _ = _request_json(
"https://passport.bilibili.com/x/passport-login/web/cookie/refresh",
data={
"csrf": csrf,
"refresh_csrf": refresh_csrf,
"source": "main_web",
"refresh_token": old_refresh_token,
},
opener=opener,
)
if refreshed.get("code") != 0:
return {"status": "failed", "message": f"Cookie 刷新失败 code={refreshed.get('code')}"}
new_refresh_token = str((refreshed.get("data") or {}).get("refresh_token") or "")
if not new_refresh_token:
return {"status": "failed", "message": "刷新响应缺少新 refresh_token"}
new_values = dict(current_values)
new_values.update(_jar_values(jar))
new_csrf = new_values.get("bili_jct", "")
new_cookie = _cookie_header(new_values)
confirmed, _ = _request_json(
"https://passport.bilibili.com/x/passport-login/web/confirm/refresh",
cookie=new_cookie,
data={"csrf": new_csrf, "refresh_token": old_refresh_token},
)
if confirmed.get("code") != 0:
return {"status": "failed", "message": f"刷新确认失败 code={confirmed.get('code')}"}
nav, _ = _request_json(
"https://api.bilibili.com/x/web-interface/nav",
cookie=new_cookie,
retries=2,
)
if nav.get("code") != 0 or not bool((nav.get("data") or {}).get("isLogin")):
return {"status": "failed", "message": "新 Cookie 登录验证失败"}
self.credential_store.save_refresh_token(new_refresh_token)
self.update_cookie(new_values)
return {"status": "refreshed", "message": "Cookie 已刷新并验证"}
async def check_once(self) -> dict[str, Any]:
try:
result = await asyncio.to_thread(self._check_and_refresh_sync)
except Exception as exc:
self.logger.warning("[B站凭据] 自动检查异常: %s", type(exc).__name__)
return {"status": "failed", "message": type(exc).__name__}
status = result.get("status")
if status == "refreshed":
self.logger.info("[B站凭据] Cookie 已自动续期并完成登录验证")
if self.on_refreshed:
callback_result = self.on_refreshed()
if asyncio.iscoroutine(callback_result):
await callback_result
elif status == "valid":
self.logger.info("[B站凭据] Cookie 有效,当前无需续期")
elif status == "disabled":
self.logger.warning("[B站凭据] 自动续期未启用:未配置刷新令牌")
else:
self.logger.warning("[B站凭据] 自动续期失败:%s", result.get("message", "未知错误"))
return result
async def run(self) -> None:
while not self._stop:
if self.is_enabled():
await self.check_once()
try:
await asyncio.wait_for(self._wake.wait(), timeout=self._current_interval())
self._wake.clear()
except asyncio.TimeoutError:
continue
+1
View File
@@ -0,0 +1 @@
"""Core helpers for the live streaming app."""
+39
View File
@@ -0,0 +1,39 @@
"""Runtime path helpers.
All mutable runtime data is resolved from the executable directory when the
program is frozen, and from the project root while running from source.
"""
from __future__ import annotations
import os
import sys
from pathlib import Path
def app_root() -> Path:
if getattr(sys, "frozen", False):
return Path(sys.executable).resolve().parent
return Path(__file__).resolve().parents[2]
APP_ROOT = app_root()
CONFIG_DIR = APP_ROOT / "config"
DATA_DIR = APP_ROOT / "data"
WEB_DIR = APP_ROOT / "web"
LOG_DIR = APP_ROOT / "logs"
INTEGRATIONS_DIR = APP_ROOT / "integrations"
VENDOR_DIR = APP_ROOT / "vendor"
DOTS_TTS_SRC = VENDOR_DIR / "dots.tts-main" / "src"
def project_path(value: str | os.PathLike, *, base: Path | None = None) -> Path:
path = Path(value)
if path.is_absolute():
return path
return (base or APP_ROOT) / path
def ensure_runtime_dirs() -> None:
for path in (CONFIG_DIR, DATA_DIR, WEB_DIR, LOG_DIR):
path.mkdir(parents=True, exist_ok=True)
+654
View File
@@ -0,0 +1,654 @@
"""
BetterGI 弹幕联动主程序
========================
监听 B 站直播间弹幕 -> 关键词匹配 -> 调用 BetterGI.exe --startGroups 执行配置组
依赖: pip install websockets brotli
运行: python danmu_bettergi.py
"""
import asyncio
import json
import logging
import os
import struct
import subprocess
import sys
import time
import urllib.request
import zlib
from pathlib import Path
APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
sys.path.insert(0, str(APP_DIR))
from core.runtime_paths import APP_ROOT as PROJECT_ROOT, CONFIG_DIR, LOG_DIR, ensure_runtime_dirs, project_path
try:
import websockets
import brotli
except ImportError:
print("缺少依赖,请先运行: pip install websockets brotli")
sys.exit(1)
# ============== B站直播弹幕协议常量 ==============
HEADER_LEN = 16
OP_HEARTBEAT = 2 # 心跳请求
OP_HEARTBEAT_REPLY = 3 # 心跳响应(人气值)
OP_MESSAGE = 5 # 业务消息
OP_AUTH = 7 # 认证请求
OP_AUTH_REPLY = 8 # 认证响应
PROTO_JSON = 0 # 明文JSON
PROTO_ZLIB = 2 # zlib压缩
PROTO_BROTLI = 3 # brotli压缩
# ============== 弹幕服务器发现 ==============
def get_danmu_server(room_id: int) -> dict:
"""通过B站API获取弹幕服务器地址和token。优先使用host_server_list中的新服务器。"""
url = f"https://api.live.bilibili.com/room/v1/Danmu/getConf?room_id={room_id}"
req = urllib.request.Request(url)
req.add_header("User-Agent", "Mozilla/5.0")
resp = urllib.request.urlopen(req, timeout=10)
data = json.loads(resp.read())["data"]
token = data["token"]
host_list = data.get("host_server_list", [])
if host_list:
entry = host_list[0]
host = entry["host"]
port = entry.get("wss_port", 443)
else:
host = data["host"]
port = data.get("wss_port", 443)
return {
"ws_url": f"wss://{host}:{port}/sub",
"token": token,
"host": host,
"port": port,
}
def get_buvid3(sessdata: str = "") -> str:
"""通过B站finger/spi接口获取buvid3 (2024+协议认证必需)。"""
url = "https://api.bilibili.com/x/frontend/finger/spi"
req = urllib.request.Request(url)
req.add_header("User-Agent", "Mozilla/5.0")
if sessdata:
req.add_header("Cookie", f"SESSDATA={sessdata}")
try:
resp = urllib.request.urlopen(req, timeout=10)
data = json.loads(resp.read())
if data.get("code") == 0:
return data["data"]["b_3"]
except Exception:
pass
return ""
# ============== 配置管理 ==============
class Config:
def __init__(self, path: str):
self.path = Path(path)
self.data = {}
self.reload()
def reload(self):
with open(self.path, "r", encoding="utf-8") as f:
self.data = json.load(f)
@property
def room_id(self) -> int:
return int(self.data["bilibili"]["room_id"])
@property
def sessdata(self) -> str:
return self.data["bilibili"].get("sessdata", "")
@property
def bettergi_exe(self) -> str:
return self.data["bettergi"]["exe_path"]
@property
def bettergi_work_dir(self) -> str:
wd = self.data["bettergi"].get("work_dir", "")
return wd if wd else str(Path(self.bettergi_exe).parent)
@property
def default_cooldown(self) -> int:
return int(self.data["global"].get("default_cooldown", 30))
@property
def admin_uids(self) -> set:
return set(int(x) for x in self.data["global"].get("admin_uids", []))
@property
def rules(self) -> list:
return self.data.get("rules", [])
@property
def restart_mode(self) -> str:
"""重启模式: gentle(温和,跳过) / aggressive(激进,先杀再启)"""
return self.data["global"].get("restart_mode", "gentle")
# ============== 冷却管理 ==============
class CooldownManager:
"""按规则记录最后触发时间,防止同一指令被弹幕刷屏重复触发。"""
def __init__(self):
self._last_fire: dict = {} # keyword -> timestamp
def can_fire(self, keyword: str, cooldown: int) -> bool:
now = time.time()
last = self._last_fire.get(keyword, 0)
return (now - last) >= cooldown
def mark_fired(self, keyword: str):
self._last_fire[keyword] = time.time()
# ============== BetterGI 调用 ==============
class BetterGIRunner:
"""封装 BetterGI.exe --startGroups 调用。"""
def __init__(self, exe_path: str, work_dir: str, logger: logging.Logger,
restart_mode: str = "gentle"):
self.exe_path = exe_path
self.work_dir = work_dir
self.logger = logger
self.restart_mode = restart_mode # gentle / aggressive
self._busy = False # 标记是否正在执行任务
self._lock = asyncio.Lock()
# 缓存 cancelTaskHotkey (从 BetterGI Config.json 读)
self._cancel_hotkey = None
self._load_cancel_hotkey()
def _load_cancel_hotkey(self):
"""从 BetterGI 的 Config.json 读取取消任务热键。"""
try:
import os
cfg_path = os.path.join(self.work_dir, "User", "Config.json")
if os.path.exists(cfg_path):
with open(cfg_path, "r", encoding="utf-8") as f:
bgi_cfg = json.load(f)
hk = bgi_cfg.get("hotKeyConfig", {}).get("cancelTaskHotkey", "")
if hk:
self._cancel_hotkey = hk
self.logger.info(f"已读取 BetterGI 取消任务热键: {hk}")
except Exception as e:
self.logger.debug(f"读取 cancelTaskHotkey 失败(不影响使用): {e}")
async def _stop_bgi_gracefully(self):
"""优雅停止 BetterGI: 先按取消热键,再 taskkill 兜底。"""
# ① 模拟按取消热键 (让 BGI 内部收尾,保存进度)
if self._cancel_hotkey:
self.logger.info(f"[激进] 按取消热键: {self._cancel_hotkey}")
try:
await self._send_key(self._cancel_hotkey)
except Exception as e:
self.logger.warning(f"[激进] 按热键失败: {e}")
# 给 BetterGI 5 秒收尾时间
await asyncio.sleep(5)
# ② taskkill 强杀兜底
self.logger.info("[激进] taskkill /F /IM BetterGI.exe")
try:
proc = await asyncio.create_subprocess_exec(
"taskkill", "/F", "/IM", "BetterGI.exe",
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
# CREATE_NO_WINDOW = 0x08000000,避免弹黑窗
creationflags=0x08000000,
)
await proc.communicate()
except Exception as e:
self.logger.debug(f"[激进] taskkill 执行: {e}")
# 再等 1 秒让进程完全退出
await asyncio.sleep(1)
async def _send_key(self, key: str):
"""模拟按键。优先用 pywin32,其次 keyboard 库,都没有则跳过。"""
# 简单映射: BetterGI 的热键名通常是 "F9" "Ctrl+P" 这种
# 优先尝试 pywin32 (最稳定,不弹窗)
try:
import win32api
import win32con
# 简单支持单键 (F1-F24, 字母, 数字)
vk_map = {
"F1": win32con.VK_F1, "F2": win32con.VK_F2, "F3": win32con.VK_F3,
"F4": win32con.VK_F4, "F5": win32con.VK_F5, "F6": win32con.VK_F6,
"F7": win32con.VK_F7, "F8": win32con.VK_F8, "F9": win32con.VK_F9,
"F10": win32con.VK_F10, "F11": win32con.VK_F11, "F12": win32con.VK_F12,
"ESC": win32con.VK_ESCAPE, "ESCAPE": win32con.VK_ESCAPE,
}
# 处理 Ctrl+X Shift+X 这种组合键
parts = key.replace("+", " ").split()
main_key = parts[-1].upper()
ctrl = "CTRL" in [p.upper() for p in parts[:-1]]
shift = "SHIFT" in [p.upper() for p in parts[:-1]]
alt = "ALT" in [p.upper() for p in parts[:-1]]
vk = vk_map.get(main_key)
if vk is None and len(main_key) == 1:
vk = ord(main_key.upper()) # 字母键
if vk is None:
self.logger.warning(f"[激进] 不支持的热键: {key}, 跳过热键直接 taskkill")
return
if ctrl:
win32api.keybd_event(win32con.VK_CONTROL, 0, 0, 0)
if shift:
win32api.keybd_event(win32con.VK_SHIFT, 0, 0, 0)
if alt:
win32api.keybd_event(win32con.VK_MENU, 0, 0, 0)
win32api.keybd_event(vk, 0, 0, 0) # 按下
win32api.keybd_event(vk, 0, win32con.KEYEVENTF_KEYUP, 0) # 抬起
if ctrl:
win32api.keybd_event(win32con.VK_CONTROL, 0, win32con.KEYEVENTF_KEYUP, 0)
if shift:
win32api.keybd_event(win32con.VK_SHIFT, 0, win32con.KEYEVENTF_KEYUP, 0)
if alt:
win32api.keybd_event(win32con.VK_MENU, 0, win32con.KEYEVENTF_KEYUP, 0)
self.logger.info(f"[激进] 已模拟按键: {key}")
except ImportError:
self.logger.warning(
"[激进] 未安装 pywin32, 无法模拟热键。"
"可运行: pip install pywin32 (可选,不装则直接 taskkill)"
)
async def run_groups(self, groups: list) -> bool:
"""异步启动配置组。"""
async with self._lock:
if self._busy:
if self.restart_mode == "aggressive":
self.logger.warning(
f"[激进模式] BetterGI 正在执行任务,先杀再启: {groups}"
)
# 释放锁,执行停止 (停止可能耗时,不持有锁)
self._busy = False
else:
self.logger.warning(
f"[温和模式] BetterGI 正在执行任务,跳过本次触发: {groups}"
)
return False
else:
self._busy = True
# 激进模式: 先停止旧任务
if self.restart_mode == "aggressive":
await self._stop_bgi_gracefully()
async with self._lock:
self._busy = True
try:
cmd = [self.exe_path, "--startGroups"] + groups
self.logger.info(f"调用 BetterGI: {' '.join(cmd)}")
proc = await asyncio.create_subprocess_exec(
*cmd,
cwd=self.work_dir,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
stdout, stderr = await proc.communicate()
if proc.returncode in (0, 553):
self.logger.info(f"BetterGI 启动成功: {groups} (rc={proc.returncode})")
return True
else:
self.logger.error(
f"BetterGI 启动失败 rc={proc.returncode}: "
f"{stderr.decode('gbk', errors='replace')}"
)
return False
except Exception as e:
self.logger.exception(f"调用 BetterGI 异常: {e}")
return False
finally:
self._busy = False
# ============== 弹幕匹配引擎 ==============
class DanmuMatcher:
"""把弹幕文本匹配到对应规则。"""
def __init__(self, config: Config, cooldown: CooldownManager,
runner: BetterGIRunner, logger: logging.Logger):
self.config = config
self.cooldown = cooldown
self.runner = runner
self.logger = logger
async def handle(self, text: str, uid: int, uname: str):
text = (text or "").strip()
if not text:
return
admins = self.config.admin_uids
is_admin = uid in admins
for rule in self.config.rules:
if not self._match(text, rule):
continue
keyword = rule["keyword"]
groups = rule["groups"]
cooldown_sec = rule.get("cooldown", self.config.default_cooldown)
admin_only = rule.get("admin_only", False)
reply = rule.get("reply", "")
# 权限检查
if admin_only and not is_admin:
self.logger.info(
f"[权限不足] {uname}({uid}) 弹幕'{text}' 命中'{keyword}' 但需管理员"
)
return # 注意: 这里return了,继续检查下一条规则(如果有)
# 冷却检查
if not self.cooldown.can_fire(keyword, cooldown_sec):
remain = int(cooldown_sec - (time.time() - self.cooldown._last_fire.get(keyword, 0)))
self.logger.info(
f"[冷却中] '{keyword}' 还需 {remain}s, 来自 {uname}({uid})"
)
return # 冷却不触发,但日志已打印
# 触发
self.cooldown.mark_fired(keyword)
self.logger.info(
f"[触发] {uname}({uid}) 弹幕'{text}' -> 配置组 {groups}"
)
if reply:
self.logger.info(f"[回复提示] {reply}")
await self.runner.run_groups(groups)
return # 一条弹幕只触发第一个命中的规则
@staticmethod
def _match(text: str, rule: dict) -> bool:
keyword = rule["keyword"]
mtype = rule.get("match_type", "contains")
if mtype == "exact":
return text == keyword
elif mtype == "contains":
return keyword in text
elif mtype == "startswith":
return text.startswith(keyword)
elif mtype == "regex":
import re
return re.search(keyword, text) is not None
return False
# ============== B站弹幕协议 ==============
def make_packet(op: int, body: bytes = b"") -> bytes:
"""组装协议包。header 16字节 + body。"""
if isinstance(body, str):
body = body.encode("utf-8")
total = HEADER_LEN + len(body)
header = struct.pack(">IHHII", total, HEADER_LEN, 1, op, 1)
return header + body
def parse_packets(data: bytes):
"""一个 WebSocket 帧可能含多个协议包,循环切分。"""
offset = 0
packets = []
while offset < len(data):
if offset + HEADER_LEN > len(data):
break
total, header_len, proto_ver, op, seq = struct.unpack(
">IHHII", data[offset:offset + HEADER_LEN]
)
body = data[offset + header_len:offset + total]
packets.append((proto_ver, op, body))
offset += total
return packets
def decode_body(proto_ver: int, body: bytes) -> bytes:
"""按协议版本解压 body。"""
if proto_ver == PROTO_JSON:
return body
if proto_ver == PROTO_ZLIB:
return zlib.decompress(body)
if proto_ver == PROTO_BROTLI:
return brotli.decompress(body)
return body
def _clean_uid(value) -> int:
"""只接受纯十进制正整数 UID,匿名或含星号的值返回 0。"""
if isinstance(value, bool):
return 0
if isinstance(value, int):
return value if value > 0 else 0
if isinstance(value, str):
value = value.strip()
return int(value) if value.isdecimal() and int(value) > 0 else 0
return 0
def extract_danmu(body: bytes):
"""从消息体中提取弹幕(DANMU_MSG)。返回 [(text, uid, uname), ...]。"""
results = []
try:
decoded = decode_body(0, body) # body 已是解压后的,proto_ver 此处忽略
except Exception:
decoded = body
try:
msg = json.loads(decoded.decode("utf-8", errors="replace"))
except Exception:
return results
cmd = msg.get("cmd", "")
if cmd.startswith("DANMU_MSG"):
info = msg.get("info", [])
if not isinstance(info, list) or len(info) <= 2:
return results
text = str(info[1] if len(info) > 1 else "")
member = info[2]
if isinstance(member, (list, tuple)):
uid = _clean_uid(member[0] if len(member) > 0 else 0)
uname = str(member[1] if len(member) > 1 and member[1] is not None else "")
elif isinstance(member, dict):
uid = _clean_uid(member.get("uid_str") or member.get("uid") or member.get("mid"))
uname = str(member.get("uname") or member.get("name") or "")
else:
uid, uname = 0, ""
results.append((text, uid, uname))
return results
# ============== 弹幕客户端 ==============
class BliveClient:
"""B站直播弹幕 WebSocket 客户端,带自动重连。"""
def __init__(self, room_id: int, sessdata: str,
matcher: DanmuMatcher, logger: logging.Logger):
self.room_id = room_id
self.sessdata = sessdata
self.matcher = matcher
self.logger = logger
self._stop = False
self._buvid3 = ""
async def run(self):
"""主循环:断线自动重连,间隔递增。"""
# 启动时获取buvid3 (2024+协议认证必需)
self._buvid3 = get_buvid3(self.sessdata)
if self._buvid3:
self.logger.info(f"获取buvid3成功: {self._buvid3[:20]}...")
else:
self.logger.warning("获取buvid3失败,弹幕可能收不到! 将尝试匿名连接")
retry = 0
while not self._stop:
try:
# 每次连接前重新获取服务器地址(避免IP变化)
server_info = get_danmu_server(self.room_id)
ws_url = server_info["ws_url"]
token = server_info["token"]
self.logger.info(f"弹幕服务器: {server_info['host']}:{server_info['port']}")
await self._connect_once(ws_url, token)
retry = 0 # 连上后重置
except asyncio.CancelledError:
break
except Exception as e:
self.logger.warning(f"连接断开: {e}")
if self._stop:
break
retry += 1
wait = min(2 ** retry, 60) # 指数退避,最多60秒
self.logger.info(f"{wait}秒后重连(第{retry}次)...")
await asyncio.sleep(wait)
def stop(self):
self._stop = True
async def _connect_once(self, ws_url: str, token: str):
self.logger.info(f"连接直播间 room_id={self.room_id} ...")
# 禁用 websockets 自带 ping/pong (B站有自己的心跳协议)
async with websockets.connect(
ws_url,
max_size=None,
ping_interval=None,
ping_timeout=None,
close_timeout=5,
) as ws:
# 1. 发送认证包(使用API返回的token + buvid3)
auth = {
"uid": 0,
"roomid": self.room_id,
"protover": PROTO_BROTLI,
"platform": "web",
"type": 2,
"key": token or (self.sessdata or ""),
}
if self._buvid3:
auth["buvid"] = self._buvid3
await ws.send(make_packet(OP_AUTH, json.dumps(auth)))
self.logger.info("已发送认证包,等待响应...")
# 立即发送一个心跳,激活弹幕推送
await ws.send(make_packet(OP_HEARTBEAT))
# 2. 心跳任务
async def heartbeat():
while True:
await asyncio.sleep(15)
try:
await ws.send(make_packet(OP_HEARTBEAT))
except Exception:
break
hb_task = asyncio.create_task(heartbeat())
# 3. 接收循环
msg_counter = 0
async for raw in ws:
if isinstance(raw, str):
continue
for proto_ver, op, body in parse_packets(raw):
if op == OP_AUTH_REPLY:
# 检查认证是否真的成功
try:
auth_resp = json.loads(body.decode("utf-8", errors="replace"))
code = auth_resp.get("code", -1)
if code == 0:
self.logger.info(f"认证成功,开始监听弹幕 (响应: {auth_resp})")
else:
self.logger.error(f"认证失败! 响应: {auth_resp}")
except Exception:
self.logger.info(f"认证响应(原始): {body}")
elif op == OP_HEARTBEAT_REPLY:
# body 是 4 字节人气值(大端序int32)
if len(body) >= 4:
popularity = struct.unpack(">I", body[:4])[0]
self.logger.debug(f"人气值: {popularity}")
elif op == OP_MESSAGE:
# body 可能被压缩,按 proto_ver 解压后再切包
try:
decoded = decode_body(proto_ver, body)
except Exception as e:
self.logger.error(f"解压失败 proto={proto_ver} len={len(body)}: {e}")
continue
# 解压后可能内含多个子包
for sub_proto, sub_op, sub_body in parse_packets(decoded):
if sub_op == OP_MESSAGE:
msg_counter += 1
# 调试:打印每条消息的cmd类型
try:
msg_json = json.loads(sub_body.decode("utf-8", errors="replace"))
cmd = msg_json.get("cmd", "?")
self.logger.debug(f"#{msg_counter} cmd={cmd}")
except Exception:
self.logger.debug(f"#{msg_counter} 解析JSON失败")
continue
for text, uid, uname in extract_danmu(sub_body):
self.logger.info(f"[弹幕] {uname}({uid}): {text}")
await self.matcher.handle(text, uid, uname)
hb_task.cancel()
# ============== 日志 ==============
def setup_logger(config: Config) -> logging.Logger:
logger = logging.getLogger("danmu_bettergi")
logger.setLevel(getattr(logging, config.data["global"].get("log_level", "INFO")))
fmt = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s",
datefmt="%H:%M:%S")
sh = logging.StreamHandler(sys.stdout)
sh.setFormatter(fmt)
logger.addHandler(sh)
log_file = config.data["global"].get("log_file")
if log_file:
LOG_DIR.mkdir(parents=True, exist_ok=True)
fh = logging.FileHandler(project_path(log_file, base=LOG_DIR), encoding="utf-8")
fh.setFormatter(fmt)
logger.addHandler(fh)
return logger
# ============== 入口 ==============
async def main():
ensure_runtime_dirs()
config_path = CONFIG_DIR / "config.json"
if not config_path.exists():
legacy_config_path = PROJECT_ROOT / "config.json"
if legacy_config_path.exists():
config_path = legacy_config_path
if not config_path.exists():
print("找不到 config.json,请先配置!")
sys.exit(1)
config = Config(str(config_path))
logger = setup_logger(config)
logger.info("=" * 50)
logger.info("BetterGI 弹幕联动启动")
logger.info(f"直播间: {config.room_id}")
logger.info(f"BetterGI: {config.bettergi_exe}")
logger.info(f"规则数: {len(config.rules)}")
logger.info(f"管理员UID: {config.admin_uids or ''}")
logger.info(f"默认冷却: {config.default_cooldown}s")
logger.info(f"重启模式: {config.restart_mode}")
logger.info("=" * 50)
cooldown = CooldownManager()
runner = BetterGIRunner(
config.bettergi_exe, config.bettergi_work_dir, logger,
restart_mode=config.restart_mode,
)
matcher = DanmuMatcher(config, cooldown, runner, logger)
client = BliveClient(config.room_id, config.sessdata, matcher, logger)
try:
await client.run()
except KeyboardInterrupt:
logger.info("收到退出信号,正在停止...")
client.stop()
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
pass
+10205
View File
File diff suppressed because it is too large Load Diff
+462
View File
@@ -0,0 +1,462 @@
from __future__ import annotations
import io
import multiprocessing
import os
import threading
import time
import traceback
import uuid
from multiprocessing.connection import Connection
from typing import Any, Callable
MAX_NEW_TOKENS = 384
DEFAULT_SYNTHESIS_TIMEOUT_SECONDS = 120.0
DEFAULT_STARTUP_TIMEOUT_SECONDS = 480.0
# 启动失败后再次拉起 worker 的最小间隔:避免"启动超时→立即重启→再超时"的死循环
# 在系统高负载时持续加载 torch/CUDA,进一步加剧卡顿。
STARTUP_FAILURE_BACKOFF_SECONDS = 60.0
DEFAULT_CPU_THREADS = 4
DEFAULT_CPU_AFFINITY_COUNT = 8
DEFAULT_PROCESS_PRIORITY = "below_normal"
class FasterQwenWorkerError(RuntimeError):
pass
class FasterQwenWorkerTimeout(FasterQwenWorkerError):
pass
def _generation_kwargs(settings: dict[str, Any], text: str) -> dict[str, Any]:
return {
"text": text,
"language": str(settings.get("language") or "Chinese"),
"non_streaming_mode": bool(settings.get("non_streaming_mode", True)),
"max_new_tokens": MAX_NEW_TOKENS,
}
def _bounded_int(value: Any, default: int, *, minimum: int, maximum: int) -> int:
try:
parsed = int(value)
except (TypeError, ValueError):
parsed = default
return max(minimum, min(maximum, parsed))
def _configure_worker_environment(settings: dict[str, Any]) -> int:
cpu_threads = _bounded_int(
settings.get("cpu_threads"),
DEFAULT_CPU_THREADS,
minimum=1,
maximum=8,
)
thread_value = str(cpu_threads)
for name in (
"OMP_NUM_THREADS",
"MKL_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"NUMEXPR_NUM_THREADS",
"VECLIB_MAXIMUM_THREADS",
"BLIS_NUM_THREADS",
):
os.environ[name] = thread_value
os.environ["TOKENIZERS_PARALLELISM"] = "false"
return cpu_threads
def _configure_torch_threads(torch_module: Any, cpu_threads: int) -> None:
torch_module.set_num_threads(cpu_threads)
try:
torch_module.set_num_interop_threads(1)
except RuntimeError:
# PyTorch only allows setting interop threads before parallel work starts.
pass
def _apply_worker_process_limits(settings: dict[str, Any]) -> dict[str, Any]:
cpu_count = max(1, int(os.cpu_count() or 1))
affinity_count = _bounded_int(
settings.get("cpu_affinity_count"),
DEFAULT_CPU_AFFINITY_COUNT,
minimum=0,
maximum=min(cpu_count, 63),
)
priority = str(settings.get("process_priority") or DEFAULT_PROCESS_PRIORITY).strip().lower()
applied_affinity = 0
applied_priority = "default"
if os.name == "nt":
try:
import ctypes
from ctypes import wintypes
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
kernel32.GetCurrentProcess.restype = wintypes.HANDLE
kernel32.SetPriorityClass.argtypes = [wintypes.HANDLE, wintypes.DWORD]
kernel32.SetPriorityClass.restype = wintypes.BOOL
kernel32.SetProcessAffinityMask.argtypes = [wintypes.HANDLE, ctypes.c_size_t]
kernel32.SetProcessAffinityMask.restype = wintypes.BOOL
process_handle = kernel32.GetCurrentProcess()
priority_classes = {
"idle": 0x00000040,
"below_normal": 0x00004000,
"normal": 0x00000020,
}
priority_class = priority_classes.get(priority, priority_classes[DEFAULT_PROCESS_PRIORITY])
if kernel32.SetPriorityClass(process_handle, priority_class):
applied_priority = priority if priority in priority_classes else DEFAULT_PROCESS_PRIORITY
if affinity_count > 0:
affinity_mask = (1 << affinity_count) - 1
if kernel32.SetProcessAffinityMask(process_handle, affinity_mask):
applied_affinity = affinity_count
except Exception:
pass
return {
"cpu_affinity_count": applied_affinity,
"process_priority": applied_priority,
}
def _load_runtime(settings: dict[str, Any]) -> dict[str, Any]:
device = str(settings.get("device") or "cuda")
if device == "cpu" and "CUDA_VISIBLE_DEVICES" not in os.environ:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
cpu_threads = _configure_worker_environment(settings)
import torch
_configure_torch_threads(torch, cpu_threads)
from faster_qwen3_tts import FasterQwen3TTS
if device == "cpu":
torch.cuda.is_available = lambda: False
load_kwargs: dict[str, Any] = {}
if device == "cpu":
load_kwargs["device"] = "cpu"
model = FasterQwen3TTS.from_pretrained(
str(settings.get("model_name_or_path") or "Qwen/Qwen3-TTS-12Hz-0.6B-Base"),
**load_kwargs,
)
# voice_clone_prompt 不再预计算:当前 faster_qwen3_tts 的 FasterQwen3TTS 没有
# create_voice_clone_prompt 方法,预计算只会失败。改为每次合成时在
# _synthesize_wav 里直接传 ref_audio/ref_text/xvec_only 参数。
voice_clone_prompt = None
return {
"model": model,
"torch": torch,
"voice_clone_prompt": voice_clone_prompt,
"settings": settings,
"cpu_threads": cpu_threads,
}
def _synthesize_wav(runtime: dict[str, Any], text: str) -> bytes:
import soundfile as sf
model = runtime["model"]
torch = runtime["torch"]
settings = runtime["settings"]
safe_text = str(text or "").strip()[:80] or "欢迎来到直播间。"
kwargs = _generation_kwargs(settings, safe_text)
voice_clone_prompt = runtime.get("voice_clone_prompt")
if voice_clone_prompt is not None:
kwargs["voice_clone_prompt"] = voice_clone_prompt
else:
ref_audio = str(settings.get("ref_audio") or "")
if not ref_audio:
raise RuntimeError("Faster-Qwen3-TTS requires ref_audio")
kwargs.update({
"ref_audio": ref_audio,
"ref_text": str(settings.get("ref_text") or "") or None,
"xvec_only": bool(settings.get("xvec_only", True)),
"append_silence": bool(settings.get("append_silence", True)),
})
if not torch.cuda.is_available():
torch.backends.cudnn.enabled = False
with torch.inference_mode():
wavs, sample_rate = model.generate_voice_clone(**kwargs)
output = io.BytesIO()
audio = wavs[0]
if isinstance(audio, torch.Tensor):
audio = audio.cpu().numpy()
sf.write(output, audio, sample_rate, format="WAV")
return output.getvalue()
def faster_qwen_worker_main(connection: Connection, settings: dict[str, Any]) -> None:
try:
cpu_threads = _configure_worker_environment(settings)
process_limits = _apply_worker_process_limits(settings)
load_started = time.monotonic()
runtime = _load_runtime(settings)
load_ms = int((time.monotonic() - load_started) * 1000)
warmup_started = time.monotonic()
_synthesize_wav(runtime, "系统启动")
warmup_ms = int((time.monotonic() - warmup_started) * 1000)
connection.send({
"type": "ready",
"pid": os.getpid(),
"load_ms": load_ms,
"warmup_ms": warmup_ms,
"max_new_tokens": MAX_NEW_TOKENS,
"cpu_threads": cpu_threads,
**process_limits,
})
except BaseException as exc:
try:
connection.send({
"type": "startup_error",
"error_type": type(exc).__name__,
"error": str(exc),
"traceback": traceback.format_exc(),
})
except Exception:
pass
connection.close()
return
while True:
try:
message = connection.recv()
except (EOFError, OSError):
break
if not isinstance(message, dict):
continue
command = str(message.get("command") or "")
if command == "stop":
break
if command != "synthesize":
continue
request_id = str(message.get("request_id") or "")
started = time.monotonic()
try:
audio = _synthesize_wav(runtime, str(message.get("text") or ""))
connection.send({
"type": "result",
"request_id": request_id,
"audio": audio,
"duration_ms": int((time.monotonic() - started) * 1000),
"bytes": len(audio),
"max_new_tokens": MAX_NEW_TOKENS,
})
except BaseException as exc:
try:
connection.send({
"type": "error",
"request_id": request_id,
"error_type": type(exc).__name__,
"error": str(exc),
"traceback": traceback.format_exc(),
})
except Exception:
break
connection.close()
class FasterQwenWorkerClient:
def __init__(
self,
settings: dict[str, Any],
logger=None,
*,
synthesis_timeout_seconds: float = DEFAULT_SYNTHESIS_TIMEOUT_SECONDS,
startup_timeout_seconds: float = DEFAULT_STARTUP_TIMEOUT_SECONDS,
context=None,
process_target: Callable[..., None] | None = None,
):
self.settings = dict(settings)
self.logger = logger
self.synthesis_timeout_seconds = max(1.0, float(synthesis_timeout_seconds))
self.startup_timeout_seconds = max(10.0, float(startup_timeout_seconds))
self._context = context or multiprocessing.get_context("spawn")
self._process_target = process_target or faster_qwen_worker_main
self._lock = threading.RLock()
self._process = None
self._connection = None
self._worker_pid = 0
self._next_start_after = 0.0
@property
def worker_pid(self) -> int:
return int(self._worker_pid or 0)
def _log(self, level: str, message: str, *args) -> None:
if self.logger:
getattr(self.logger, level)(message, *args)
def _is_alive_locked(self) -> bool:
return bool(self._process is not None and self._process.is_alive())
def ensure_ready(self) -> dict[str, Any]:
with self._lock:
if self._is_alive_locked() and self._connection is not None:
return {"pid": self.worker_pid, "reused": True}
return self._start_worker_locked()
def _start_worker_locked(self) -> dict[str, Any]:
now = time.monotonic()
if now < self._next_start_after:
wait = int(self._next_start_after - now)
raise FasterQwenWorkerError(
f"TTS worker 启动退避中,距上次启动失败不足 {int(STARTUP_FAILURE_BACKOFF_SECONDS)} 秒,"
f"{wait} 秒后可重试"
)
self._terminate_worker_locked("replace_stale_worker", graceful=False)
parent_connection, child_connection = self._context.Pipe(duplex=True)
process = self._context.Process(
target=self._process_target,
args=(child_connection, self.settings),
name="FasterQwen3TTSWorker",
daemon=True,
)
process.start()
try:
child_connection.close()
except Exception:
pass
self._process = process
self._connection = parent_connection
self._worker_pid = int(getattr(process, "pid", 0) or 0)
self._log("info", "[FasterQwenTTS] worker 已启动, pid=%s,正在加载和预热", self.worker_pid)
if not parent_connection.poll(self.startup_timeout_seconds):
self._next_start_after = time.monotonic() + STARTUP_FAILURE_BACKOFF_SECONDS
self._terminate_worker_locked("startup_timeout", graceful=False)
raise FasterQwenWorkerError(
f"Faster-Qwen3-TTS worker startup exceeded {self.startup_timeout_seconds:.0f}s"
)
try:
message = parent_connection.recv()
except (EOFError, OSError) as exc:
self._next_start_after = time.monotonic() + STARTUP_FAILURE_BACKOFF_SECONDS
self._terminate_worker_locked("startup_connection_closed", graceful=False)
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker exited during startup") from exc
if not isinstance(message, dict) or message.get("type") != "ready":
if isinstance(message, dict):
error = str(message.get("error") or message.get("error_type") or "unknown startup error")
else:
error = "invalid startup response"
self._next_start_after = time.monotonic() + STARTUP_FAILURE_BACKOFF_SECONDS
self._terminate_worker_locked("startup_error", graceful=False)
raise FasterQwenWorkerError(error)
self._next_start_after = 0.0
self._log(
"info",
"[FasterQwenTTS] worker 预热完成, pid=%s load=%sms warmup=%sms max_new_tokens=%s",
self.worker_pid,
message.get("load_ms"),
message.get("warmup_ms"),
message.get("max_new_tokens"),
)
self._log(
"info",
"[FasterQwenTTS] worker 资源限制: cpu_threads=%s affinity=%s priority=%s",
message.get("cpu_threads"),
message.get("cpu_affinity_count"),
message.get("process_priority"),
)
return message
def synthesize(self, text: str) -> tuple[bytes, dict[str, Any]]:
with self._lock:
self.ensure_ready()
request_id = uuid.uuid4().hex
connection = self._connection
try:
connection.send({
"command": "synthesize",
"request_id": request_id,
"text": str(text or ""),
})
except (BrokenPipeError, EOFError, OSError) as exc:
self._restart_after_failure_locked("send_failed")
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker connection failed") from exc
if not connection.poll(self.synthesis_timeout_seconds):
self._log(
"error",
"[FasterQwenTTS] 单次合成超过 %.0f 秒,强制终止 worker pid=%s",
self.synthesis_timeout_seconds,
self.worker_pid,
)
restart_error = self._restart_after_failure_locked("synthesis_timeout")
suffix = f"; restart failed: {restart_error}" if restart_error else ""
raise FasterQwenWorkerTimeout(
f"Faster-Qwen3-TTS synthesis exceeded {self.synthesis_timeout_seconds:.0f}s{suffix}"
)
try:
message = connection.recv()
except (EOFError, OSError) as exc:
self._restart_after_failure_locked("worker_exited")
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker exited during synthesis") from exc
if not isinstance(message, dict) or message.get("request_id") != request_id:
self._restart_after_failure_locked("invalid_response")
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker returned an invalid response")
if message.get("type") == "error":
error = str(message.get("error") or message.get("error_type") or "synthesis failed")
lowered = error.lower()
if "cuda" in lowered or "out of memory" in lowered or "device-side" in lowered:
self._restart_after_failure_locked("cuda_error")
raise FasterQwenWorkerError(error)
if message.get("type") != "result" or not isinstance(message.get("audio"), bytes):
self._restart_after_failure_locked("invalid_result")
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker returned no audio")
return message["audio"], message
def _restart_after_failure_locked(self, reason: str) -> str:
self._terminate_worker_locked(reason, graceful=False)
try:
self._start_worker_locked()
return ""
except Exception as exc:
self._log("error", "[FasterQwenTTS] worker 自动重启失败: %s", exc)
return str(exc)
def _terminate_worker_locked(self, reason: str, *, graceful: bool) -> None:
process = self._process
connection = self._connection
self._process = None
self._connection = None
self._worker_pid = 0
if process is None:
if connection is not None:
try:
connection.close()
except Exception:
pass
return
if graceful and process.is_alive() and connection is not None:
try:
connection.send({"command": "stop"})
process.join(timeout=3.0)
except Exception:
pass
if process.is_alive():
self._log("warning", "[FasterQwenTTS] 终止 worker, reason=%s pid=%s", reason, process.pid)
process.terminate()
process.join(timeout=10.0)
if process.is_alive():
process.kill()
process.join(timeout=5.0)
if connection is not None:
try:
connection.close()
except Exception:
pass
def close(self) -> None:
with self._lock:
self._terminate_worker_locked("shutdown", graceful=True)
+569
View File
@@ -0,0 +1,569 @@
"""安全、幂等地将旧版 JSON/日志快照补录到统计数据库。"""
from __future__ import annotations
import argparse
import asyncio
import hashlib
import json
import re
from collections import Counter
from collections.abc import Mapping
from datetime import UTC, datetime, timedelta, timezone
from pathlib import Path
from typing import Any
try:
from .stats_store import DEFAULT_DATABASE_PATH, StatsStore, business_date
except ImportError: # 允许直接执行文件
from stats_store import DEFAULT_DATABASE_PATH, StatsStore, business_date
BEIJING_TZ = timezone(timedelta(hours=8), name="Asia/Shanghai")
_SOURCE_FILES = (
"users.json",
"song_requests.json",
"admin_audit.log",
"queue_state.json",
"music_state.json",
"tts_state.json",
)
_USER_FIELDS = {
"uname",
"points",
"last_signin_date",
"created_at",
"blocked_all",
"blocked_queue",
"blocked_song_request",
"note",
}
_SONG_FIELDS = {
"id",
"name",
"artist",
"duration_ms",
"duration_sec",
"keyword",
"uid",
"uname",
"requested_at",
"source",
"remove_after_play",
"started_at",
"finished_at",
"status",
}
_AUDIT_FIELDS = {"at", "action", "target", "client_ip", "session_id", "detail"}
TABLE_COLUMNS = {
"users": {
"platform", "platform_user_id", "display_name", "avatar_url", "user_level",
"is_admin", "first_seen_at_utc", "last_seen_at_utc", "snapshot_json",
},
"song_requests": {
"request_id", "platform", "platform_user_id", "requested_at_utc", "business_date",
"song_id", "song_name", "artist", "source", "status", "points_cost",
"queue_position", "payload_json",
},
"playback_sessions": {
"playback_id", "request_id", "song_id", "song_name", "started_at_utc",
"ended_at_utc", "business_date", "status", "duration_ms", "played_ms",
"stop_reason", "payload_json",
},
"admin_audit_events": {
"event_id", "occurred_at_utc", "business_date", "actor", "action",
"target_type", "target_id", "success", "remote_address_hash", "payload_json",
},
"point_transactions": {
"transaction_id", "platform", "platform_user_id", "occurred_at_utc",
"business_date", "amount", "balance_after", "reason", "reference_type",
"reference_id", "payload_json",
},
"events": {
"event_id", "event_type", "category", "occurred_at_utc", "business_date",
"payload_json",
},
}
_POINT_DETAIL_RE = re.compile(r"^delta=(-?\d+)\s+now=(-?\d+)$")
_COUNT_DETAIL_RE = re.compile(r"^count=(\d+)$")
_STATUS_MAP = {
"played": "completed",
"complete": "completed",
"finished": "completed",
"success": "completed",
"skipped": "skipped",
"skip": "skipped",
"interrupted": "interrupted",
"cancelled": "cancelled",
"canceled": "cancelled",
"play_error": "failed",
"error": "failed",
"failed": "failed",
"queued": "queued",
"pending": "queued",
"playing": "playing",
"active": "playing",
}
def _hash(*parts: Any) -> str:
encoded = json.dumps(parts, ensure_ascii=False, separators=(",", ":"), default=str).encode("utf-8", "replace")
return hashlib.sha256(encoded).hexdigest()
def _text(value: Any, limit: int) -> str | None:
if value is None:
return None
result = str(value).strip()
return result[:limit] if result else None
def _int(value: Any, *, minimum: int | None = None) -> int | None:
try:
result = int(value)
except (TypeError, ValueError, OverflowError):
return None
if minimum is not None and result < minimum:
return None
return result
def _bool(value: Any) -> bool:
return value is True or value == 1
def _utc_iso(
value: Any,
*,
naive_is_beijing: bool = True,
fallback: datetime | None = None,
timespec: str = "milliseconds",
) -> str:
parsed: datetime
if isinstance(value, (int, float)) and not isinstance(value, bool):
parsed = datetime.fromtimestamp(float(value), UTC)
elif isinstance(value, str) and value.strip():
raw = value.strip()
try:
parsed = datetime.fromisoformat(raw.replace("Z", "+00:00"))
except ValueError:
parsed = fallback or datetime.now(UTC)
else:
parsed = fallback or datetime.now(UTC)
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=BEIJING_TZ if naive_is_beijing else UTC)
return parsed.astimezone(UTC).isoformat(timespec=timespec).replace("+00:00", "Z")
def _file_time(path: Path) -> str:
return _utc_iso(datetime.fromtimestamp(path.stat().st_mtime, UTC), naive_is_beijing=False)
def _read_bytes(path: Path) -> bytes:
return path.read_bytes()
def _read_json(raw: bytes) -> Any:
return json.loads(raw.decode("utf-8-sig"))
def _fingerprint(raw: bytes) -> str:
return hashlib.sha256(raw).hexdigest()
def _normalize_status(value: Any, default: str) -> str:
key = str(value or "").strip().casefold()
return _STATUS_MAP.get(key, default)
def _json_payload(value: Mapping[str, Any]) -> dict[str, Any]:
return {key: item for key, item in value.items() if item is not None}
def _user_records(path: Path, raw: bytes) -> list[tuple[str, dict[str, Any]]]:
data = _read_json(raw)
if not isinstance(data, Mapping):
return []
snapshot_at = _file_time(path)
records: list[tuple[str, dict[str, Any]]] = []
for uid, original in data.items():
if not isinstance(original, Mapping):
continue
item = {key: original[key] for key in _USER_FIELDS if key in original}
platform_user_id = _text(uid, 64)
if not platform_user_id:
continue
created_at = _utc_iso(item.get("created_at"), fallback=datetime.fromtimestamp(path.stat().st_mtime, UTC))
snapshot = _json_payload({
"points": _int(item.get("points")),
"last_signin_date": _text(item.get("last_signin_date"), 10),
"blocked_all": _bool(item.get("blocked_all")),
"blocked_queue": _bool(item.get("blocked_queue")),
"blocked_song_request": _bool(item.get("blocked_song_request")),
"note_length": len(str(item.get("note") or "")),
"legacy_snapshot": True,
})
records.append(("users", {
"platform": "bilibili",
"platform_user_id": platform_user_id,
"display_name": _text(item.get("uname"), 128),
"avatar_url": None,
"user_level": None,
"is_admin": False,
"first_seen_at_utc": created_at,
"last_seen_at_utc": snapshot_at,
"snapshot_json": snapshot,
}))
return records
def _song_records(path: Path, raw: bytes) -> list[tuple[str, dict[str, Any]]]:
data = _read_json(raw)
if not isinstance(data, Mapping):
return []
snapshot_at = _file_time(path)
fallback = datetime.fromtimestamp(path.stat().st_mtime, UTC)
entries: list[dict[str, Any]] = []
sections = (("queue", data.get("queue")), ("active", [data.get("active")]), ("history", data.get("history")))
for section, values in sections:
if not isinstance(values, list):
continue
section_entries: dict[tuple[str | None, str, str | None, str], list[dict[str, Any]]] = {}
for queue_position, original in enumerate(values, 1):
if not isinstance(original, Mapping):
continue
item = {key: original[key] for key in _SONG_FIELDS if key in original}
requested = _utc_iso(item.get("requested_at"), fallback=fallback, timespec="microseconds")
source = _text(item.get("source"), 64) or "legacy"
business_key = (
_text(item.get("uid"), 64),
requested,
_text(item.get("id"), 128),
source,
)
entry = {
"section": section,
"item": item,
"requested": requested,
"started": _utc_iso(item["started_at"], timespec="microseconds") if item.get("started_at") is not None else None,
"ended": _utc_iso(item["finished_at"], timespec="microseconds") if item.get("finished_at") is not None else None,
"source": source,
"business_key": business_key,
"queue_position": queue_position if section == "queue" else None,
}
section_entries.setdefault(business_key, []).append(entry)
for business_key, duplicates in section_entries.items():
duplicates.sort(key=lambda entry: (
entry["started"] or "",
entry["ended"] or "",
str(entry["item"].get("status") or ""),
str(entry["item"].get("name") or ""),
str(entry["item"].get("artist") or ""),
str(entry["item"].get("duration_ms") or ""),
))
for duplicate_ordinal, entry in enumerate(duplicates):
entry["request_id"] = _hash("legacy-song-request", *business_key, duplicate_ordinal)
entries.append(entry)
section_priority = {"queue": 0, "active": 1, "history": 2}
requests: dict[str, dict[str, Any]] = {}
playbacks: dict[str, dict[str, Any]] = {}
for entry in sorted(entries, key=lambda value: section_priority[value["section"]]):
section = entry["section"]
item = entry["item"]
request_id = entry["request_id"]
confidence = "snapshot" if section in {"queue", "active"} else "history"
default_status = "queued" if section == "queue" else "playing" if section == "active" else "unknown"
status = _normalize_status(item.get("status"), default_status)
requests[request_id] = {
"request_id": request_id,
"platform": "bilibili",
"platform_user_id": _text(item.get("uid"), 64),
"requested_at_utc": entry["requested"],
"business_date": business_date(entry["requested"]),
"song_id": _text(item.get("id"), 128),
"song_name": _text(item.get("name"), 256),
"artist": _text(item.get("artist"), 256),
"source": entry["source"],
"status": status,
"points_cost": 0,
"queue_position": entry["queue_position"],
"payload_json": _json_payload({
"legacy_section": section,
"confidence": confidence,
"latest_80_only": section == "history",
"keyword_length": len(str(item.get("keyword") or "")),
"remove_after_play": _bool(item.get("remove_after_play")),
"snapshot_at_utc": snapshot_at if confidence == "snapshot" else None,
}),
}
started = entry["started"]
if not started:
continue
ended = entry["ended"]
played_ms = None
if ended:
start_dt = datetime.fromisoformat(started.replace("Z", "+00:00"))
end_dt = datetime.fromisoformat(ended.replace("Z", "+00:00"))
played_ms = max(0, round((end_dt - start_dt).total_seconds() * 1000))
playback_id = _hash("legacy-playback", request_id, started)
playbacks[playback_id] = {
"playback_id": playback_id,
"request_id": request_id,
"song_id": _text(item.get("id"), 128),
"song_name": _text(item.get("name"), 256),
"started_at_utc": started,
"ended_at_utc": ended,
"business_date": business_date(started),
"status": status,
"duration_ms": _int(item.get("duration_ms"), minimum=0),
"played_ms": played_ms,
"stop_reason": "snapshot" if confidence == "snapshot" else status,
"payload_json": {
"legacy_section": section,
"confidence": confidence,
"latest_80_only": section == "history",
},
}
return [
*(("song_requests", record) for record in requests.values()),
*(("playback_sessions", record) for record in playbacks.values()),
]
def _safe_detail(detail: Any) -> dict[str, Any]:
text = str(detail or "")
result: dict[str, Any] = {
"detail_length": len(text),
"detail_sha256": _hash("admin-detail", text),
}
count_match = _COUNT_DETAIL_RE.fullmatch(text)
if count_match:
result.update({"detail_kind": "count", "count": int(count_match.group(1))})
return result
try:
parsed = json.loads(text)
except (TypeError, ValueError):
parsed = None
if isinstance(parsed, Mapping):
allowed_flags = {key: _bool(parsed[key]) for key in ("blocked_all", "blocked_queue", "blocked_song_request") if key in parsed}
if allowed_flags:
result.update({"detail_kind": "flags", "flags": allowed_flags})
return result
result["detail_kind"] = "opaque"
return result
def _audit_records(raw: bytes) -> list[tuple[str, dict[str, Any]]]:
records: list[tuple[str, dict[str, Any]]] = []
duplicate_counts: Counter[str] = Counter()
for raw_line in raw.decode("utf-8-sig", "replace").splitlines():
line = raw_line.strip()
if not line:
continue
try:
original = json.loads(line)
except json.JSONDecodeError:
continue
if not isinstance(original, Mapping):
continue
line_hash = _hash("legacy-admin-audit-line", line)
duplicate_ordinal = duplicate_counts[line_hash]
duplicate_counts[line_hash] += 1
item = {key: original[key] for key in _AUDIT_FIELDS if key in original}
occurred = _utc_iso(item.get("at"), naive_is_beijing=True)
action = _text(item.get("action"), 96) or "unknown"
target = _text(item.get("target"), 256)
event_id = _hash("legacy-admin-audit", line_hash, duplicate_ordinal)
payload = _safe_detail(item.get("detail"))
session = _text(item.get("session_id"), 512)
if session:
payload["session_hash"] = _hash("admin-session", session)
records.append(("admin_audit_events", {
"event_id": event_id,
"occurred_at_utc": occurred,
"business_date": business_date(occurred),
"actor": "admin",
"action": action,
"target_type": "legacy_target",
"target_id": target,
"success": True,
"remote_address_hash": _hash("admin-ip", item.get("client_ip")) if item.get("client_ip") else None,
"payload_json": payload,
}))
point_match = _POINT_DETAIL_RE.fullmatch(str(item.get("detail") or "")) if action == "user_add_points" else None
if point_match and target:
records.append(("point_transactions", {
"transaction_id": _hash("legacy-admin-points", event_id),
"platform": "bilibili",
"platform_user_id": target,
"occurred_at_utc": occurred,
"business_date": business_date(occurred),
"amount": int(point_match.group(1)),
"balance_after": int(point_match.group(2)),
"reason": "admin_adjustment",
"reference_type": "admin_audit",
"reference_id": event_id,
"payload_json": {"legacy_import": True},
}))
return records
def _snapshot_event(path: Path, raw: bytes) -> list[tuple[str, dict[str, Any]]]:
data = _read_json(raw)
if not isinstance(data, Mapping):
return []
monitor_updated_at = (data.get("monitor") or {}).get("updated_at") if isinstance(data.get("monitor"), Mapping) else None
occurred = _utc_iso(
data.get("updated_at") or monitor_updated_at,
fallback=datetime.fromtimestamp(path.stat().st_mtime, UTC),
)
if path.name == "queue_state.json":
queue = data.get("queue") if isinstance(data.get("queue"), list) else []
payload = {
"queue_size": len(queue),
"has_active_user": bool(data.get("current_admin_uid")),
"has_group": bool(data.get("current_group")),
"default_running": _bool(data.get("default_running")),
"login_status": _text(data.get("login_status"), 64),
"has_user_finished_once": _bool(data.get("has_user_finished_once")),
"confidence": "snapshot",
}
event_type = "legacy.queue_snapshot"
category = "queue"
elif path.name == "music_state.json":
current = data.get("current") if isinstance(data.get("current"), Mapping) else {}
monitor = data.get("monitor") if isinstance(data.get("monitor"), Mapping) else {}
payload = {
"playing": _bool(data.get("playing")),
"current_title": _text(current.get("title"), 256),
"current_artist": _text(current.get("artist"), 256),
"duration": _int(current.get("duration"), minimum=0),
"progress": _int(current.get("progress"), minimum=0),
"playlist_size": len(data.get("playlist")) if isinstance(data.get("playlist"), list) else 0,
"request_size": len(data.get("requests")) if isinstance(data.get("requests"), list) else 0,
"monitor_online": _bool(monitor.get("online")),
"platform": _text(monitor.get("platform"), 32),
"confidence": "snapshot",
}
event_type = "legacy.music_snapshot"
category = "music"
else:
recent = data.get("recent_events") if isinstance(data.get("recent_events"), list) else []
payload = {
"enabled": _bool(data.get("enabled")),
"provider": _text(data.get("provider"), 64),
"model_loaded": _bool(data.get("model_loaded")),
"last_duration_ms": _int(data.get("last_duration_ms"), minimum=0),
"has_last_error": bool(data.get("last_error")),
"total_synthesized": _int(data.get("total_synthesized"), minimum=0),
"total_errors": _int(data.get("total_errors"), minimum=0),
"recent_event_count": len(recent),
"last_text_length": len(str(data.get("last_text") or "")),
"confidence": "snapshot",
}
event_type = "legacy.tts_snapshot"
category = "tts"
payload = _json_payload(payload)
return [("events", {
"event_id": _hash("legacy-snapshot", path.name, _fingerprint(raw)),
"event_type": event_type,
"category": category,
"occurred_at_utc": occurred,
"business_date": business_date(occurred),
"payload_json": payload,
})]
def _build_records(path: Path, raw: bytes) -> list[tuple[str, dict[str, Any]]]:
if path.name == "users.json":
return _user_records(path, raw)
if path.name == "song_requests.json":
return _song_records(path, raw)
if path.name == "admin_audit.log":
return _audit_records(raw)
return _snapshot_event(path, raw)
def validate_records(records: list[tuple[str, dict[str, Any]]]) -> None:
"""拒绝未知表、未知字段及非映射记录,避免导入边界被意外扩大。"""
for record_number, record in enumerate(records, 1):
if not isinstance(record, tuple) or len(record) != 2:
raise ValueError(f"{record_number} 条导入记录格式无效")
table, fields = record
allowed_columns = TABLE_COLUMNS.get(table)
if allowed_columns is None:
raise ValueError(f"{record_number} 条记录使用未知表: {table}")
if not isinstance(fields, Mapping):
raise ValueError(f"{record_number} 条记录字段不是映射")
unknown_columns = set(fields) - allowed_columns
if unknown_columns:
raise ValueError(f"{table} 包含未知字段: {sorted(unknown_columns)}")
if not fields:
raise ValueError(f"{table} 导入记录不能为空")
async def backfill_legacy_statistics(store: StatsStore, data_dir: str | Path, dry_run: bool = False) -> dict[str, Any]:
"""扫描旧数据并通过 ``StatsStore.import_once`` 原子、幂等地补录。"""
root = Path(data_dir)
result: dict[str, Any] = {"dry_run": bool(dry_run), "sources": {}, "record_counts": {}}
totals: Counter[str] = Counter()
for filename in _SOURCE_FILES:
path = root / filename
if not path.is_file():
result["sources"][filename] = {"status": "missing", "records": 0}
continue
raw = _read_bytes(path)
records = _build_records(path, raw)
validate_records(records)
counts = Counter(table for table, _ in records)
totals.update(counts)
imported = True
if not dry_run:
imported = await store.import_once(
f"legacy-history:{filename}",
"legacy_history",
records,
fingerprint=_fingerprint(raw),
checkpoint_key="sha256",
checkpoint_value=_fingerprint(raw),
metadata={"filename": filename, "record_count": len(records), "schema": 1},
)
result["sources"][filename] = {
"status": "dry_run" if dry_run else "completed" if imported else "failed",
"records": len(records),
"tables": dict(sorted(counts.items())),
}
result["record_counts"] = dict(sorted(totals.items()))
result["total_records"] = sum(totals.values())
return result
async def _main_async(args: argparse.Namespace) -> int:
data_dir = Path(args.data_dir)
if args.dry_run:
result = await backfill_legacy_statistics(StatsStore(args.database), data_dir, dry_run=True)
else:
store = StatsStore(args.database)
if not await store.start():
print(json.dumps({"status": "failed", "reason": "database_start_failed"}, ensure_ascii=False))
return 1
try:
result = await backfill_legacy_statistics(store, data_dir)
await store.flush()
finally:
await store.close()
print(json.dumps(result, ensure_ascii=False, sort_keys=True))
return 0 if all(source["status"] != "failed" for source in result["sources"].values()) else 1
def main() -> int:
parser = argparse.ArgumentParser(description="安全、幂等地补录旧版直播统计")
parser.add_argument("--data-dir", default=str(Path(__file__).resolve().parents[1] / "data"))
parser.add_argument("--database", default=str(DEFAULT_DATABASE_PATH))
parser.add_argument("--dry-run", action="store_true")
return asyncio.run(_main_async(parser.parse_args()))
if __name__ == "__main__":
raise SystemExit(main())
+260
View File
@@ -0,0 +1,260 @@
"""Unified launcher for source and frozen builds."""
from __future__ import annotations
import argparse
import asyncio
import ctypes
import os
import socket
import subprocess
import sys
from pathlib import Path
APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
sys.path.insert(0, str(APP_DIR))
from core.runtime_paths import APP_ROOT, DATA_DIR, ensure_runtime_dirs
ERROR_ALREADY_EXISTS = 183
MAIN_INSTANCE_MUTEX = "Local\\BetterGI_LiveStreaming_Main_5191"
class SingleInstanceLock:
"""Windows named mutex used by the queue-producing main process only."""
def __init__(self, name: str = MAIN_INSTANCE_MUTEX):
self.name = name
self._handle = None
def acquire(self) -> bool:
if os.name != "nt":
return True
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
handle = kernel32.CreateMutexW(None, False, self.name)
last_error = ctypes.get_last_error()
if not handle:
raise ctypes.WinError(last_error)
self._handle = handle
if last_error == ERROR_ALREADY_EXISTS:
self.close()
return False
return True
def close(self):
if self._handle is None or os.name != "nt":
return
ctypes.WinDLL("kernel32", use_last_error=True).CloseHandle(self._handle)
self._handle = None
class _JOBOBJECT_IO_COUNTERS(ctypes.Structure):
_fields_ = [
("ReadOperationCount", ctypes.c_ulonglong),
("WriteOperationCount", ctypes.c_ulonglong),
("OtherOperationCount", ctypes.c_ulonglong),
("ReadTransferCount", ctypes.c_ulonglong),
("WriteTransferCount", ctypes.c_ulonglong),
("OtherTransferCount", ctypes.c_ulonglong),
]
class _JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure):
_fields_ = [
("PerProcessUserTimeLimit", ctypes.c_longlong),
("PerJobUserTimeLimit", ctypes.c_longlong),
("LimitFlags", ctypes.c_ulong),
("MinimumWorkingSetSize", ctypes.c_size_t),
("MaximumWorkingSetSize", ctypes.c_size_t),
("ActiveProcessLimit", ctypes.c_ulong),
("Affinity", ctypes.c_size_t),
("PriorityClass", ctypes.c_ulong),
("SchedulingClass", ctypes.c_ulong),
]
class _JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure):
_fields_ = [
("BasicLimitInformation", _JOBOBJECT_BASIC_LIMIT_INFORMATION),
("IoInfo", _JOBOBJECT_IO_COUNTERS),
("ProcessMemoryLimit", ctypes.c_size_t),
("JobMemoryLimit", ctypes.c_size_t),
("PeakProcessMemoryUsed", ctypes.c_size_t),
("PeakJobMemoryUsed", ctypes.c_size_t),
]
class WindowsJob:
"""Kill spawned music/TTS processes automatically when the launcher exits."""
JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000
JOB_OBJECT_EXTENDED_LIMIT_INFORMATION = 9
def __init__(self):
self._handle = None
if os.name != "nt":
return
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
handle = kernel32.CreateJobObjectW(None, None)
if not handle:
raise ctypes.WinError(ctypes.get_last_error())
info = _JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
info.BasicLimitInformation.LimitFlags = self.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
ok = kernel32.SetInformationJobObject(
handle,
self.JOB_OBJECT_EXTENDED_LIMIT_INFORMATION,
ctypes.byref(info),
ctypes.sizeof(info),
)
if not ok:
error = ctypes.get_last_error()
kernel32.CloseHandle(handle)
raise ctypes.WinError(error)
self._handle = handle
def assign(self, process: subprocess.Popen):
if self._handle is None or os.name != "nt":
return
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
if not kernel32.AssignProcessToJobObject(self._handle, process._handle):
raise ctypes.WinError(ctypes.get_last_error())
def close(self):
if self._handle is None or os.name != "nt":
return
ctypes.WinDLL("kernel32", use_last_error=True).CloseHandle(self._handle)
self._handle = None
def _is_frozen() -> bool:
return bool(getattr(sys, "frozen", False))
def _role_command(role: str, port: int, host: str) -> list[str]:
if _is_frozen():
return [sys.executable, "--role", role, "--port", str(port), "--host", host]
return [sys.executable, str(Path(__file__).resolve()), "--role", role, "--port", str(port), "--host", host]
def _assert_port_available(host: str, port: int):
probe_host = "0.0.0.0" if host in {"", "::"} else host
family = socket.AF_INET6 if ":" in probe_host else socket.AF_INET
with socket.socket(family, socket.SOCK_STREAM) as sock:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0)
sock.bind((probe_host, port))
def _spawn_role(role: str, port: int, host: str, *, visible: bool = False) -> subprocess.Popen:
creationflags = 0
if os.name == "nt":
creationflags = 0x00000010 if visible else 0x08000000 # CREATE_NEW_CONSOLE / CREATE_NO_WINDOW
return subprocess.Popen(
_role_command(role, port, host),
cwd=APP_ROOT,
stdout=None if visible else subprocess.DEVNULL,
stderr=None if visible else subprocess.DEVNULL,
creationflags=creationflags,
)
async def _run_queue(host: str, port: int):
import danmu_queue
await danmu_queue.main(host=host, port=port)
async def _run_music(port: int):
import music_monitor
await music_monitor.run_monitor(port)
def _run_tts_monitor():
import tts_monitor
sys.argv = [
sys.argv[0],
"--state-file",
str(DATA_DIR / "tts_state.json"),
]
tts_monitor.main()
def _stop_children(children: list[subprocess.Popen], timeout: float = 5.0):
for child in children:
if child.poll() is None:
child.terminate()
for child in children:
if child.poll() is not None:
continue
try:
child.wait(timeout=timeout)
except subprocess.TimeoutExpired:
child.kill()
child.wait(timeout=timeout)
async def _run_all(host: str, port: int):
children: list[subprocess.Popen] = []
job = WindowsJob()
queue_task: asyncio.Task | None = None
try:
# Refuse stale/conflicting listeners before creating any helper process.
_assert_port_available(host, port)
queue_task = asyncio.create_task(_run_queue(host, port), name="queue-main")
children.append(_spawn_role("music", port, host, visible=False))
children.append(_spawn_role("tts", port, host, visible=True))
for child in children:
job.assign(child)
await queue_task
except Exception:
if queue_task is not None and not queue_task.done():
queue_task.cancel()
await asyncio.gather(queue_task, return_exceptions=True)
raise
finally:
_stop_children(children)
job.close()
def main():
ensure_runtime_dirs()
parser = argparse.ArgumentParser(description="BetterGI 直播联动统一入口")
parser.add_argument("--role", choices=["all", "queue", "music", "tts"], default="all")
parser.add_argument("--port", type=int, default=8086)
parser.add_argument("--host", default="0.0.0.0", help="Web service bind address")
args = parser.parse_args()
instance_lock = None
if args.role in {"all", "queue"}:
instance_lock = SingleInstanceLock()
if not instance_lock.acquire():
print("直播系统已经在运行,本次重复启动已拒绝。")
return 2
try:
if args.role == "tts":
_run_tts_monitor()
return 0
if args.role == "music":
asyncio.run(_run_music(args.port))
return 0
if args.role == "queue":
asyncio.run(_run_queue(args.host, args.port))
return 0
asyncio.run(_run_all(args.host, args.port))
return 0
except OSError as exc:
if getattr(exc, "winerror", None) == 10048 or getattr(exc, "errno", None) in {48, 98, 10048}:
print(f"直播端口 {args.port} 已被占用,服务未启动,也未创建辅助进程。")
return 3
raise
finally:
if instance_lock is not None:
instance_lock.close()
if __name__ == "__main__":
raise SystemExit(main())
+466
View File
@@ -0,0 +1,466 @@
"""Independent mpv audio player controlled through Windows JSON IPC."""
from __future__ import annotations
import asyncio
import json
import logging
import os
import subprocess
import threading
import time
from pathlib import Path
from typing import Any
class MpvPlayer:
def __init__(
self,
exe_path: str | Path,
logger: logging.Logger,
*,
pipe_name: str = "",
log_path: str | Path | None = None,
):
self.exe_path = Path(exe_path)
self.logger = logger
self.log_path = Path(log_path) if log_path else None
pipe_name = pipe_name or f"live_streaming_mpv_{os.getpid()}"
self.pipe_path = rf"\\.\pipe\{pipe_name}"
self.process: subprocess.Popen | None = None
self.current_url = ""
self.current_metadata: dict[str, Any] = {}
self.desired_state = "stopped"
self.generation = 0
self.started_at = 0.0
self.last_progress = 0.0
self.last_progress_at = 0.0
self.last_snapshot_at = 0.0
self.recovery_count = 0
self._request_id = 0
self._ipc_lock = asyncio.Lock()
self._pipe_state_lock = threading.Lock()
self._pipe = None
def available(self) -> bool:
return self.exe_path.is_file()
def running(self) -> bool:
return self.process is not None and self.process.poll() is None
def update_exe_path(self, exe_path: str | Path) -> None:
next_path = Path(exe_path)
if next_path == self.exe_path:
return
if self.running():
self.logger.warning(f"[mpv] 播放器路径已修改,将在进程下次重启后生效: {next_path}")
self.exe_path = next_path
async def ensure_started(self) -> bool:
if self.running():
return True
await asyncio.to_thread(self._reset_pipe_sync)
if not self.available():
self.logger.error(f"[mpv] 播放器不存在: {self.exe_path}")
return False
creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0) if os.name == "nt" else 0
args = [
str(self.exe_path),
"--idle=yes",
"--no-video",
"--force-window=no",
"--no-terminal",
"--msg-level=all=warn",
f"--input-ipc-server={self.pipe_path}",
"--keep-open=no",
"--audio-buffer=5",
"--cache=yes",
"--cache-secs=20",
"--demuxer-max-bytes=50MiB",
"--network-timeout=10",
]
if self.log_path:
self.log_path.parent.mkdir(parents=True, exist_ok=True)
args.append(f"--log-file={self.log_path}")
try:
self.process = subprocess.Popen(
args,
cwd=str(self.exe_path.parent),
stdin=subprocess.DEVNULL,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
creationflags=creationflags,
)
for _ in range(30):
if not self.running():
break
if await self._command(["get_property", "idle-active"], retry=False) is not None:
self.logger.info(f"[mpv] 播放服务已启动: {self.exe_path}")
return True
await asyncio.sleep(0.1)
except Exception as exc:
self.logger.error(f"[mpv] 启动失败: {exc}")
return False
def _pipe_request_sync(self, command: list[Any], request_id: int) -> Any:
request = json.dumps(
{"command": command, "request_id": request_id},
ensure_ascii=False,
).encode("utf-8") + b"\n"
pipe = None
try:
with self._pipe_state_lock:
pipe = self._pipe
if pipe is None or pipe.closed:
pipe = open(self.pipe_path, "r+b", buffering=0)
self._pipe = pipe
pipe.write(request)
deadline = time.time() + 2.5
while time.time() < deadline:
response = pipe.readline()
if not response:
continue
payload = json.loads(response.decode("utf-8", errors="replace"))
if payload.get("request_id") != request_id:
continue
if payload.get("error") != "success":
return None
if "data" not in payload:
return True
data = payload["data"]
if data is None and command and command[0] != "get_property":
return True
return data
except Exception:
self._reset_pipe_sync(pipe)
raise
return None
def _reset_pipe_sync(self, expected_pipe=None) -> None:
with self._pipe_state_lock:
pipe = self._pipe
if expected_pipe is not None and pipe is not expected_pipe:
return
self._pipe = None
if pipe is not None:
try:
pipe.close()
except Exception:
pass
async def _command(self, command: list[Any], *, retry: bool = True) -> Any:
async with self._ipc_lock:
attempts = 2 if retry else 1
for attempt in range(attempts):
self._request_id += 1
request_id = self._request_id
try:
return await asyncio.wait_for(
asyncio.to_thread(self._pipe_request_sync, command, request_id),
timeout=3.0,
)
except Exception:
self._reset_pipe_sync()
if attempt + 1 < attempts:
await asyncio.sleep(0.15)
return None
def _clear_current(self) -> None:
self.current_url = ""
self.current_metadata = {}
self.started_at = 0.0
self.last_progress = 0.0
self.last_progress_at = 0.0
self.recovery_count = 0
@staticmethod
def _path_matches(expected: str, actual: str) -> bool:
expected = str(expected or "").strip()
actual = str(actual or "").strip()
if not expected or not actual:
return False
if expected == actual:
return True
if expected.lower().startswith(("http://", "https://")):
return False
try:
actual_path = actual[8:] if actual.lower().startswith("file:///") else actual
return Path(expected).resolve() == Path(actual_path).resolve()
except Exception:
return False
async def _wait_until_loaded(
self,
*,
expected_generation: int,
start_at: float,
timeout: float = 12.0,
) -> bool:
deadline = time.time() + max(2.0, timeout)
expected_url = self.current_url
while time.time() < deadline:
if expected_generation != self.generation or self.desired_state != "playing":
return False
if not self.running():
return False
path = await self._command(["get_property", "path"], retry=False)
idle = await self._command(["get_property", "idle-active"], retry=False)
if self._path_matches(expected_url, str(path or "")) and idle is False:
if start_at > 0:
seek_result = None
seek_deadline = min(deadline, time.time() + 4.0)
while time.time() < seek_deadline:
duration = await self._command(["get_property", "duration"], retry=False)
try:
duration_value = max(0.0, float(duration or 0.0))
except (TypeError, ValueError):
duration_value = 0.0
if duration_value > 0:
seek_result = await self._command(["seek", float(start_at), "absolute+exact"])
if seek_result is not None:
break
await asyncio.sleep(0.2)
if seek_result is None:
self.logger.warning(f"[mpv] 续播定位失败: {start_at:.1f}")
return False
if await self._command(["set_property", "pause", False]) is None:
return False
progress_deadline = min(deadline, time.time() + 4.0)
baseline = max(0.0, float(start_at or 0.0))
while time.time() < progress_deadline:
progress = await self._command(["get_property", "time-pos"], retry=False)
duration = await self._command(["get_property", "duration"], retry=False)
try:
progress_value = max(0.0, float(progress or 0.0))
except (TypeError, ValueError):
progress_value = 0.0
try:
duration_value = max(0.0, float(duration or 0.0))
except (TypeError, ValueError):
duration_value = 0.0
idle_now = await self._command(["get_property", "idle-active"], retry=False)
if idle_now is False and (progress_value > 0.05 or duration_value > 0 or baseline > 0):
now = time.time()
self.started_at = now
self.last_progress = max(baseline, progress_value)
self.last_progress_at = now
return True
await asyncio.sleep(0.15)
await asyncio.sleep(0.15)
self.logger.warning(
f"[mpv] 音频加载超时,未进入可播放状态: "
f"{self.current_metadata.get('name') or self.current_url}"
)
return False
async def _load_current(self, *, start_at: float, expected_generation: int) -> bool:
if expected_generation != self.generation or self.desired_state != "playing" or not self.current_url:
return False
if not await self.ensure_started():
return False
if expected_generation != self.generation or self.desired_state != "playing":
return False
result = await self._command(["loadfile", self.current_url, "replace"])
if result is None or expected_generation != self.generation or self.desired_state != "playing":
return False
return await self._wait_until_loaded(
expected_generation=expected_generation,
start_at=start_at,
)
async def play(self, url: str, metadata: dict[str, Any], *, start_at: float = 0.0) -> bool:
self.generation += 1
generation = self.generation
self.desired_state = "playing"
self.current_url = str(url)
self.current_metadata = dict(metadata)
self.recovery_count = 0
ok = await self._load_current(start_at=start_at, expected_generation=generation)
if not ok and generation == self.generation:
self.desired_state = "stopped"
self._clear_current()
return ok
async def pause(self) -> bool:
if not self.current_url:
return False
self.desired_state = "paused"
if not self.running():
return True
return await self._command(["set_property", "pause", True]) is not None
async def resume(self) -> bool:
if not self.current_url:
return False
self.desired_state = "playing"
self.last_progress_at = time.time()
if not self.running():
return await self._load_current(start_at=self.last_progress, expected_generation=self.generation)
return await self._command(["set_property", "pause", False]) is not None
async def stop(self) -> bool:
self.generation += 1
self.desired_state = "stopped"
self._clear_current()
if not self.running():
return True
return await self._command(["stop"]) is not None
async def close(self) -> None:
self.generation += 1
self.desired_state = "stopped"
self._clear_current()
if self.running():
await self._command(["quit"])
await asyncio.sleep(0.15)
if self.running():
self.process.terminate()
try:
await asyncio.to_thread(self.process.wait, 2)
except Exception:
if self.running():
self.process.kill()
self.process = None
await asyncio.to_thread(self._reset_pipe_sync)
async def snapshot(self) -> dict[str, Any]:
now = time.time()
if not self.running():
return self._snapshot_payload(False, 0.0, 0.0, True, False, False, "", now)
progress, duration, paused, idle, eof, path = await asyncio.gather(
self._command(["get_property", "time-pos"]),
self._command(["get_property", "duration"]),
self._command(["get_property", "pause"]),
self._command(["get_property", "idle-active"]),
self._command(["get_property", "eof-reached"]),
self._command(["get_property", "path"]),
)
try:
progress_value = max(0.0, float(progress or 0.0))
except (TypeError, ValueError):
progress_value = 0.0
try:
duration_value = max(0.0, float(duration or self.current_metadata.get("duration_sec", 0) or 0))
except (TypeError, ValueError):
duration_value = 0.0
idle_value = bool(idle) if idle is not None else not bool(self.current_url)
paused_value = bool(paused)
eof_value = bool(eof)
playing = (
self.desired_state == "playing"
and bool(self.current_url)
and not paused_value
and not idle_value
and not eof_value
)
if progress_value > self.last_progress + 0.2:
self.last_progress = progress_value
self.last_progress_at = now
self.recovery_count = 0
self.last_snapshot_at = now
return self._snapshot_payload(
playing,
progress_value,
duration_value,
idle_value,
paused_value,
eof_value,
str(path or ""),
now,
)
def _snapshot_payload(
self,
playing: bool,
progress: float,
duration: float,
idle: bool,
paused: bool,
eof: bool,
path: str,
now: float,
) -> dict[str, Any]:
metadata = self.current_metadata
return {
"playing": playing,
"paused": paused,
"idle": idle,
"eof": eof,
"desired_state": self.desired_state,
"generation": self.generation,
"path": path,
"current": {
"id": str(metadata.get("id", "")),
"title": metadata.get("name") or ("暂无歌曲" if idle else "正在加载"),
"artist": metadata.get("artist") or "mpv",
"cover": metadata.get("cover", ""),
"cover_hash": metadata.get("cover_hash", ""),
"duration": duration,
"progress": progress,
"source": "mpv",
},
"playlist": [],
"requests": [],
"monitor": {
"online": self.running(),
"source": "mpv.ipc",
"platform": "mpv",
"updated_at": now,
"targets": [],
"allow_all": False,
},
}
def mark_ended(self) -> None:
self.generation += 1
self.desired_state = "stopped"
self._clear_current()
async def maintain(self, *, stall_seconds: float = 12.0) -> dict[str, Any]:
"""Recover only unexpected failures while the desired state is playing."""
if self.desired_state != "playing" or not self.current_url:
return {"action": "none"}
generation = self.generation
now = time.time()
if not self.running():
progress = self.last_progress
self.process = None
if await self._load_current(start_at=progress, expected_generation=generation):
self.logger.warning(f"[mpv] 进程退出后已从 {progress:.1f} 秒恢复")
return {"action": "process_restarted", "progress": progress}
return {"action": "failed", "reason": "process_restart_failed"}
state = await self.snapshot()
if generation != self.generation or self.desired_state != "playing":
return {"action": "superseded", "snapshot": state}
progress = float((state.get("current") or {}).get("progress", 0) or 0)
duration = float((state.get("current") or {}).get("duration", 0) or 0)
effective_progress = max(progress, self.last_progress)
near_end = duration > 0 and effective_progress >= max(0.0, duration - 2.0)
unloaded_after_progress = (
state.get("idle")
and not str(state.get("path") or "")
and effective_progress >= 0.5
)
if state.get("eof") or (state.get("idle") and near_end) or unloaded_after_progress:
self.mark_ended()
return {"action": "ended", "progress": effective_progress, "duration": duration, "snapshot": state}
loading_grace = now - self.started_at < 2.5
if state.get("paused") and not loading_grace:
if await self._command(["set_property", "pause", False]) is not None:
self.last_progress_at = now
self.logger.warning("[mpv] 检测到非预期暂停,已自动继续播放")
return {"action": "resumed", "progress": progress, "snapshot": state}
if state.get("idle") and not loading_grace:
return {"action": "reload_required", "reason": "unexpected_idle", "progress": progress, "snapshot": state}
if self.last_progress_at and not loading_grace and now - self.last_progress_at >= max(3.0, stall_seconds):
self.recovery_count += 1
if self.recovery_count == 1:
await self._command(["set_property", "pause", False])
self.last_progress_at = now
return {"action": "unstalled", "progress": progress, "snapshot": state}
return {"action": "reload_required", "reason": "stalled", "progress": progress, "snapshot": state}
return {"action": "none", "snapshot": state}
+965
View File
@@ -0,0 +1,965 @@
"""
音乐播放器 -> 直播间 UI 音乐状态同步
通过 Windows 10/11 的 SMTCSystem Media Transport Controls)读取当前媒体会话,
并推送到 BGI 直播间 Web 服务。默认偏向网易云音乐,但可在 config/config.json
的 music_monitor 节点或命令行中配置目标播放器。
播放状态判断(参考 now-playing-service/NeteaseMusicService.cs):
- 优先用 Windows 音频会话的峰值音量(volume>0 = Playing),最贴近"是否真在出声"
- volume=0 但 UIA 进度最近 1.5s 内变化 → 仍视为 Playing(静音播放/拖进度条场景)
- 否则 Paused
封面防闪:
- 只在 SMTC 缩略图 hash 变化时才重写 music_cover.jpg(原子 tmp→replace
- 同一 hash 不重写文件,前端不会因文件变化触发 reload
- 标题/歌手来自 cloudmusic 窗口标题(已稳定,不改)
"""
import argparse
import asyncio
import hashlib
import ctypes
import json
import os
import re
import subprocess
import sys
import time
import traceback
from pathlib import Path
from typing import Any
APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
sys.path.insert(0, str(APP_DIR))
import aiohttp
import winrt.windows.media.control as wmc
import winrt.windows.storage.streams as streams
try:
import uiautomation as uia
except Exception:
uia = None
# pycaw: 读取 Windows 音频会话峰值音量(参考 now-playing-service 的 CSCore.AudioMeterInformation
try:
from pycaw.pycaw import AudioUtilities # type: ignore
_PYCAW_OK = True
except Exception:
_PYCAW_OK = False
from core.runtime_paths import CONFIG_DIR, DATA_DIR, WEB_DIR, ensure_runtime_dirs
ensure_runtime_dirs()
CONFIG_FILE = CONFIG_DIR / "config.json"
MUSIC_FILE = DATA_DIR / "music_state.json"
COVER_FILE = WEB_DIR / "music_cover.jpg"
DEFAULT_TARGETS = [
"网易云音乐",
"Netease",
"CloudMusic",
"cloudmusic",
"YesPlayMusic",
"Listen1",
"QQMusic",
"qqmusic",
"spotify",
]
DEFAULT_MONITOR_CONFIG = {
"platform": "netease",
"targets": ["网易云音乐", "Netease", "CloudMusic", "cloudmusic"],
"allow_all": False,
"interval_sec": 1.0,
"holdover_ms": 1500,
"prefer_playing": True,
"keep_last_when_none": True,
"cover_enabled": True,
"auto_resume_enabled": False,
"auto_resume_interval_sec": 3,
"auto_resume_stall_sec": 10,
"extra_filter": "",
}
def load_monitor_config() -> dict[str, Any]:
cfg = dict(DEFAULT_MONITOR_CONFIG)
if CONFIG_FILE.exists():
try:
data = json.loads(CONFIG_FILE.read_text(encoding="utf-8"))
if isinstance(data.get("music_monitor"), dict):
cfg.update(data["music_monitor"])
except Exception:
pass
targets = cfg.get("targets") or []
if isinstance(targets, str):
targets = [x.strip() for x in targets.replace("", ",").split(",") if x.strip()]
extra = str(cfg.get("extra_filter", "")).strip()
if extra:
targets.append(extra)
cfg["targets"] = [str(x).strip() for x in targets if str(x).strip()]
return cfg
def session_identity(session) -> str:
display = getattr(session, "source_app_display_name", "") or ""
aumid = getattr(session, "source_app_user_model_id", "") or ""
return f"{display} {aumid}".strip()
def is_target_app(session, targets: list[str], allow_all: bool) -> bool:
if allow_all or not targets:
return True
combined = session_identity(session).lower()
return any(target.lower() in combined for target in targets)
def playback_status_int(session) -> int:
try:
return int(session.get_playback_info().playback_status)
except Exception:
return 0
def is_playing(session) -> bool:
return playback_status_int(session) == 4
async def choose_session(mgr, cfg: dict[str, Any]):
"""从全部媒体会话中选择最合适的目标会话。"""
targets = cfg.get("targets", [])
allow_all = bool(cfg.get("allow_all", False))
prefer_playing = bool(cfg.get("prefer_playing", True))
sessions = list(mgr.get_sessions())
candidates = [s for s in sessions if is_target_app(s, targets, allow_all)]
if not candidates:
current = mgr.get_current_session()
if current and is_target_app(current, targets, allow_all):
candidates = [current]
if not candidates:
return None
# 当前系统媒体会话通常最能代表用户正在操作的播放器;优先选择它,
# 避免网易云残留多个 SMTC 会话时反复读到已停止更新的旧会话。
current = mgr.get_current_session()
if current and current in candidates and (not prefer_playing or is_playing(current)):
return current
if prefer_playing:
playing = [s for s in candidates if is_playing(s)]
if playing:
return playing[0]
if current and current in candidates:
return current
return candidates[0]
async def read_thumbnail(thumb_ref):
if thumb_ref is None:
return None
try:
stream = await thumb_ref.open_read_async()
size = stream.size
if size <= 0 or size > 5 * 1024 * 1024:
return None
buffer = streams.Buffer(size)
await stream.read_async(buffer, size, streams.InputStreamOptions.READ_AHEAD)
data = bytes(buffer)
return data if len(data) > 100 else None
except Exception:
return None
def write_cover_atomic(data: bytes, previous_hash: str = "") -> tuple[str, str]:
cover_hash = hashlib.md5(data).hexdigest()
if cover_hash == previous_hash and COVER_FILE.exists():
return "/music_cover.jpg", cover_hash
COVER_FILE.parent.mkdir(parents=True, exist_ok=True)
tmp = COVER_FILE.with_suffix(".tmp")
tmp.write_bytes(data)
tmp.replace(COVER_FILE)
return "/music_cover.jpg", cover_hash
def parse_time_to_seconds(text: str) -> int | None:
parts = str(text).strip().split(":")
if len(parts) not in (2, 3):
return None
try:
nums = [int(x) for x in parts]
except ValueError:
return None
if any(x < 0 for x in nums) or nums[-1] >= 60:
return None
if len(nums) == 2:
return nums[0] * 60 + nums[1]
if nums[1] >= 60:
return None
return nums[0] * 3600 + nums[1] * 60 + nums[2]
def parse_progress_text(text: str) -> tuple[int, int] | None:
cleaned = str(text).replace(" ", "")
m = re.search(r"(\d{1,2}:\d{2}(?::\d{2})?)\s*[/|/|]\s*(\d{1,2}:\d{2}(?::\d{2})?)", cleaned)
if not m:
return None
current = parse_time_to_seconds(m.group(1))
total = parse_time_to_seconds(m.group(2))
if current is None or total is None or total <= 0 or current > total + 2:
return None
return max(0, current), max(0, total)
def _cloudmusic_pids() -> list[int]:
if not hasattr(ctypes, "windll"):
return []
try:
output = subprocess.check_output(
["tasklist", "/FI", "IMAGENAME eq cloudmusic.exe", "/FO", "CSV", "/NH"],
text=True, encoding="gbk", errors="ignore", creationflags=0x08000000,
)
pids = []
for line in output.splitlines():
parts = [p.strip().strip('"') for p in line.split(",")]
if len(parts) >= 2 and parts[0].lower() == "cloudmusic.exe":
try:
pids.append(int(parts[1]))
except ValueError:
pass
return pids
except Exception:
return []
def _window_titles_by_pids(pids: list[int]) -> list[str]:
if not pids or not hasattr(ctypes, "windll"):
return []
titles = []
user32 = ctypes.windll.user32
EnumWindowsProc = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_void_p)
def enum_proc(hwnd, lparam):
try:
if not user32.IsWindowVisible(hwnd):
return True
pid = ctypes.c_ulong()
user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid))
if pid.value not in pids:
return True
length = user32.GetWindowTextLengthW(hwnd)
if length <= 0:
return True
buf = ctypes.create_unicode_buffer(length + 1)
user32.GetWindowTextW(hwnd, buf, length + 1)
title = (buf.value or "").strip()
if title:
titles.append(title)
except Exception:
pass
return True
try:
user32.EnumWindows(EnumWindowsProc(enum_proc), 0)
except Exception:
pass
return titles
def read_netease_window_title() -> str:
"""参考 now-playing-service:从 cloudmusic 进程窗口标题取“歌名 - 歌手”。"""
for title in _window_titles_by_pids(_cloudmusic_pids()):
if " - " in title and "MediaPlayer" not in title:
return title.replace("/", " / ").strip()
return ""
def read_netease_progress_uia() -> tuple[int, int] | None:
"""参考 now-playing-service:只在 cloudmusic 播放窗口子树中解析 MM:SS / MM:SS。"""
if uia is None:
return None
try:
pids = set(_cloudmusic_pids())
root = uia.GetRootControl()
for win in root.GetChildren():
try:
pid = int(getattr(win, "ProcessId", 0) or 0)
name = (getattr(win, "Name", "") or "")
if pids and pid not in pids:
continue
if " - " not in name and "cloudmusic" not in (getattr(win, "ClassName", "") or "").lower():
continue
stack = list(win.GetChildren())
deadline = time.time() + 0.3
while stack and time.time() < deadline:
ctrl = stack.pop(0)
text_value = (getattr(ctrl, "Name", "") or "").strip()
parsed = parse_progress_text(text_value)
if parsed:
return parsed
try:
stack.extend(ctrl.GetChildren())
except Exception:
pass
except Exception:
continue
except Exception:
return None
return None
_LAST_NETEASE_TITLE = ""
_LAST_NETEASE_AT = 0.0
_LAST_PROGRESS_SECONDS = -1
_LAST_PROGRESS_CHANGE_AT = 0.0
_LAST_VOLUME_PEAK = 0.0
_LAST_PLAYING = False
_LAST_COVER_HASH = ""
_LAST_COVER_SOURCE = ""
_COVER_CACHE: dict[str, tuple[str, str]] = {}
_COVER_TASKS: dict[str, asyncio.Task] = {}
_SMTC_DISABLED_UNTIL = 0.0
_LAST_AUTO_RESUME_AT = 0.0
_AUTO_RESUME_TRACK_KEY = ""
_AUTO_RESUME_LAST_PROGRESS = -1.0
_AUTO_RESUME_LAST_PROGRESS_AT = 0.0
_AUTO_RESUME_PAUSED_SINCE = 0.0
_AUTO_RESUME_FAILURES = 0
async def get_target_smtc_session(cfg: dict[str, Any]):
mgr = await wmc.GlobalSystemMediaTransportControlsSessionManager.request_async()
return await choose_session(mgr, cfg)
async def try_resume_playback(cfg: dict[str, Any], force: bool = False) -> bool:
"""通过 SMTC 尝试恢复播放;暂停/停滞时允许重复发送播放命令。"""
global _SMTC_DISABLED_UNTIL
if time.time() < _SMTC_DISABLED_UNTIL:
return False
async def _inner() -> bool:
session = await get_target_smtc_session(cfg)
if session is None:
return False
if is_playing(session) and not force:
return False
result = await session.try_play_async()
return bool(result)
def _run():
return asyncio.run(_inner())
try:
return await asyncio.wait_for(asyncio.to_thread(_run), timeout=2.5)
except Exception:
_SMTC_DISABLED_UNTIL = time.time() + 5
return False
async def maybe_auto_resume(info: dict[str, Any] | None, cfg: dict[str, Any]) -> str:
"""定期恢复明确暂停,也检测“状态为播放但进度长时间不动”的假播放。"""
global _LAST_AUTO_RESUME_AT, _AUTO_RESUME_TRACK_KEY
global _AUTO_RESUME_LAST_PROGRESS, _AUTO_RESUME_LAST_PROGRESS_AT
global _AUTO_RESUME_PAUSED_SINCE, _AUTO_RESUME_FAILURES
if not bool(cfg.get("auto_resume_enabled", False)):
return ""
if info is None or not info.get("title"):
_AUTO_RESUME_TRACK_KEY = ""
_AUTO_RESUME_LAST_PROGRESS = -1.0
_AUTO_RESUME_LAST_PROGRESS_AT = 0.0
_AUTO_RESUME_PAUSED_SINCE = 0.0
_AUTO_RESUME_FAILURES = 0
return ""
now = time.time()
title = str(info.get("title") or "").strip()
artist = str(info.get("artist") or "").strip()
track_key = f"{title}\n{artist}"
progress = max(0.0, float(info.get("progress") or 0))
duration = max(0.0, float(info.get("duration") or 0))
playing = bool(info.get("playing"))
if track_key != _AUTO_RESUME_TRACK_KEY:
_AUTO_RESUME_TRACK_KEY = track_key
_AUTO_RESUME_LAST_PROGRESS = progress
_AUTO_RESUME_LAST_PROGRESS_AT = now
_AUTO_RESUME_PAUSED_SINCE = 0.0 if playing else now
_AUTO_RESUME_FAILURES = 0
return ""
if playing:
_AUTO_RESUME_PAUSED_SINCE = 0.0
elif not _AUTO_RESUME_PAUSED_SINCE:
_AUTO_RESUME_PAUSED_SINCE = now
if abs(progress - _AUTO_RESUME_LAST_PROGRESS) >= 0.5:
_AUTO_RESUME_LAST_PROGRESS = progress
_AUTO_RESUME_LAST_PROGRESS_AT = now
_AUTO_RESUME_FAILURES = 0
return ""
interval_sec = max(1.0, float(cfg.get("auto_resume_interval_sec", 3) or 3))
stall_sec = max(interval_sec, float(cfg.get("auto_resume_stall_sec", 10) or 10))
stalled = (
playing
and progress > 0
and (duration <= 0 or progress < max(0, duration - 2))
and now - _AUTO_RESUME_LAST_PROGRESS_AT >= stall_sec
)
if playing and not stalled:
return ""
if now - _LAST_AUTO_RESUME_AT < interval_sec:
return ""
_LAST_AUTO_RESUME_AT = now
# 明确暂停也强制发送 try_play,避免 SMTC 状态缓存或会话切换导致第一次命令被吞掉。
resumed = await try_resume_playback(cfg, force=True)
if resumed:
_AUTO_RESUME_LAST_PROGRESS_AT = now
_AUTO_RESUME_FAILURES = 0
return "stalled" if stalled else "paused"
_AUTO_RESUME_FAILURES += 1
# 连续恢复失败时缩短 SMTC 禁用窗口,下一轮重新获取会话并继续尝试。
if _AUTO_RESUME_FAILURES >= 2:
global _SMTC_DISABLED_UNTIL
_SMTC_DISABLED_UNTIL = min(_SMTC_DISABLED_UNTIL, now + 1)
return ""
async def read_smtc_snapshot(cfg: dict[str, Any]) -> dict[str, Any] | None:
"""在隔离线程里读取 SMTC,避免 WinRT 偶发卡死拖住 MusicMonitor 主循环。"""
global _SMTC_DISABLED_UNTIL
if time.time() < _SMTC_DISABLED_UNTIL:
return None
async def _inner():
session = await get_target_smtc_session(cfg)
if session is None:
return None
props = await session.try_get_media_properties_async()
playback = session.get_playback_info()
timeline = session.get_timeline_properties()
thumb_data = None
if bool(cfg.get("cover_enabled", True)):
thumb_data = await read_thumbnail(props.thumbnail)
return {
"title": props.title or "",
"artist": props.artist or "",
"duration": max(0, timeline.end_time.total_seconds()),
"progress": max(0, timeline.position.total_seconds()),
"playing": int(playback.playback_status) == 4,
"source": session_identity(session) or "smtc",
"thumb_data": thumb_data,
}
def _run():
return asyncio.run(_inner())
try:
return await asyncio.wait_for(asyncio.to_thread(_run), timeout=2.5)
except Exception:
_SMTC_DISABLED_UNTIL = time.time() + 5
return None
def _get_cloudmusic_audio_peak() -> float:
"""参考 now-playing-service: 累加 cloudmusic 所有音频会话的峰值音量。
volume>0 = 真在出声 = Playing。pycaw 不可用时返回 -1 表示未知。"""
if not _PYCAW_OK:
return -1.0
try:
from pycaw.pycaw import IAudioMeterInformation # type: ignore
total = 0.0
sessions = AudioUtilities.GetAllSessions()
for sess in sessions:
try:
proc = getattr(sess, "Process", None)
if proc is None:
continue
# pycaw 的 Process.name 是方法不是属性, 要调用
name_attr = getattr(proc, "name", None)
if callable(name_attr):
pname = name_attr()
else:
pname = str(name_attr or "")
pname = (pname or "").lower()
if "cloudmusic" in pname:
meter = sess._ctl.QueryInterface(IAudioMeterInformation)
total += meter.GetPeakValue()
except Exception:
continue
return total
except Exception:
return -1.0
def _decide_playing(volume_peak: float, progress_changed_recently: bool, holdover_sec: float) -> bool:
"""参考 now-playing-service: volume>0 → Playing; volume=0 但进度最近变化 → Playing; 否则保持/暂停。"""
global _LAST_PLAYING
if volume_peak > 0.00001:
_LAST_PLAYING = True
return True
if volume_peak < 0:
# pycaw 不可用,回退到进度判断
if progress_changed_recently:
_LAST_PLAYING = True
return True
_LAST_PLAYING = False
return False
# volume=0
if progress_changed_recently:
# 进度在动但没声音 → 静音播放,仍算 Playing
_LAST_PLAYING = True
return True
_LAST_PLAYING = False
return False
def _split_title_artist(window_title: str) -> tuple[str, str]:
title = (window_title or "").strip()
if " - " in title:
song, artist = title.split(" - ", 1)
return song.strip() or title, artist.strip()
return title, ""
import websockets as _ws_mod
def fiber_store_extract_js() -> str:
return r'''
function _ensureStore() {
try {
if (window._reduxStore) return true;
const rootEl = document.querySelector('#root');
const root = window._fiberRoot || (rootEl && rootEl._reactRootContainer && rootEl._reactRootContainer._internalRoot);
if (!root) return false;
let queue = [root.current || root];
let visited = 0;
while (queue.length > 0) {
let node = queue.shift();
if (!node) continue;
visited++;
if (visited > 20000) break;
if (node.memoizedProps && node.memoizedProps.store) { window._reduxStore = node.memoizedProps.store; return true; }
if (node.stateNode && node.stateNode.store) { window._reduxStore = node.stateNode.store; return true; }
let child = node.child;
while (child) { queue.push(child); child = child.sibling; }
}
return false;
} catch(err) { return false; }
}
'''
async def get_cdp_ws_url(port: int = 9222) -> str:
try:
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=1.2)) as session:
async with session.get(f"http://127.0.0.1:{port}/json") as resp:
if resp.status != 200:
return ""
targets = await resp.json(content_type=None)
for target in targets:
text = (str(target.get("url", "")) + " " + str(target.get("title", ""))).lower()
if target.get("type") == "page" and target.get("webSocketDebuggerUrl") and ("orpheus" in text or "music.163.com" in text):
return target.get("webSocketDebuggerUrl", "")
for target in targets:
if target.get("type") == "page" and target.get("webSocketDebuggerUrl"):
return target.get("webSocketDebuggerUrl", "")
except Exception:
return ""
return ""
async def read_netease_cdp_state(cfg: dict[str, Any], previous_hash: str = "") -> dict[str, Any] | None:
"""通过 CDP 读网易云 Redux 状态。只取 title/artist/progress/duration/picUrl
不取 playingplaying 交给音量峰值判断)。"""
port = int(((cfg.get("request_player") or {}).get("cdp_port", 9222)) or 9222)
ws_url = await get_cdp_ws_url(port)
if not ws_url:
return None
script = fiber_store_extract_js() + r'''
(function(){
if(!_ensureStore()) return null;
const state = window._reduxStore.getState();
const playing = state.playing || {};
const list = (state.playingList && state.playingList.curPlayingList) || [];
const id = playing.resourceTrackId || playing.onlineResourceId || playing.resourceId || playing.trackId;
function normId(x){ return x == null ? '' : String(x); }
let item = null;
if(id) item = list.find(x => normId(x.id || x.trackId || x.resourceId) === normId(id));
if(!item && list.length === 1) item = list[0];
const track = (item && (item.track || item.resource || item)) || {};
const artists = track.artists || track.ar || item?.artists || item?.ar || [];
let artist = '';
if(Array.isArray(artists)) artist = artists.map(a => a && a.name ? a.name : '').filter(Boolean).join('/');
else if(typeof artists === 'string') artist = artists;
const album = track.album || track.al || item?.album || item?.al || {};
const picUrl = album.picUrl || album.blurPicUrl || track.picUrl || item?.picUrl || '';
const durationMs = Number(track.duration || track.dt || item?.duration || item?.dt || 0);
const positionMs = Number(playing.position || playing.currentTime || playing.progress || 0);
return {
id: normId(id || item?.id || track.id),
title: track.name || item?.name || '',
artist: artist || '',
picUrl: picUrl || '',
duration: durationMs > 10000 ? durationMs / 1000 : durationMs,
progress: positionMs > 10000 ? positionMs / 1000 : positionMs,
};
})()
'''
try:
async with _ws_mod.connect(ws_url, open_timeout=1.5, close_timeout=0.5) as ws:
await ws.send(json.dumps({
"id": 1,
"method": "Runtime.evaluate",
"params": {"expression": script, "returnByValue": True, "awaitPromise": True},
}))
deadline = time.time() + 2
while time.time() < deadline:
raw = await asyncio.wait_for(ws.recv(), timeout=max(0.1, deadline - time.time()))
msg = json.loads(raw)
if msg.get("id") != 1:
continue
value = (((msg.get("result") or {}).get("result") or {}).get("value"))
if isinstance(value, dict) and (value.get("title") or value.get("id")):
return value
return None
except Exception:
return None
return None
async def download_cover_url(pic_url: str, previous_hash: str = "") -> tuple[str, str]:
global _LAST_COVER_SOURCE
if not pic_url:
return "", previous_hash
if pic_url.startswith("http://"):
pic_url = "https://" + pic_url[7:]
if pic_url == _LAST_COVER_SOURCE and previous_hash and COVER_FILE.exists():
return "/music_cover.jpg", previous_hash
try:
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=3)) as session:
async with session.get(pic_url) as resp:
if resp.status != 200:
return "", previous_hash
data = await resp.read()
if data:
cover_url, cover_hash = write_cover_atomic(data, previous_hash)
_LAST_COVER_SOURCE = pic_url
return cover_url, cover_hash
except Exception:
return "", previous_hash
return "", previous_hash
def cover_cache_key(title: str, artist: str) -> str:
return re.sub(r"\s+", " ", f"{title} - {artist}".strip().lower())
async def fetch_cover_task(key: str, title: str, artist: str):
"""后台补封面:失败也不能影响 MusicMonitor 主循环。"""
try:
keyword = " ".join(x for x in [title, artist] if x).strip()
if not keyword:
return
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
"Referer": "https://music.163.com/",
"Content-Type": "application/x-www-form-urlencoded",
"Cookie": "os=pc; appver=2.9.8;",
"X-Real-IP": "118.88.88.88",
"X-Forwarded-For": "118.88.88.88",
}
async with aiohttp.ClientSession(headers=headers, timeout=aiohttp.ClientTimeout(total=4)) as session:
async with session.post(
"https://music.163.com/api/search/get/web",
data={"s": keyword, "type": "1", "limit": "1", "offset": "0"},
) as resp:
result = await resp.json(content_type=None)
songs = ((result.get("result") or {}).get("songs") or [])
if not songs:
return
album = songs[0].get("album") or {}
pic_url = album.get("picUrl") or album.get("blurPicUrl") or ""
if not pic_url:
return
if pic_url.startswith("http://"):
pic_url = "https://" + pic_url[7:]
async with session.get(pic_url) as img_resp:
if img_resp.status != 200:
return
data = await img_resp.read()
if data:
_COVER_CACHE[key] = write_cover_atomic(data)
except Exception:
pass
finally:
_COVER_TASKS.pop(key, None)
def apply_cached_or_schedule_cover(info: dict[str, Any], cfg: dict[str, Any]):
if not bool(cfg.get("cover_enabled", True)):
return
key = cover_cache_key(info.get("title", ""), info.get("artist", ""))
if not key:
return
cached = _COVER_CACHE.get(key)
if cached:
info["cover"], info["cover_hash"] = cached
return
task = _COVER_TASKS.get(key)
if task is None or task.done():
try:
_COVER_TASKS[key] = asyncio.create_task(fetch_cover_task(key, info.get("title", ""), info.get("artist", "")))
except RuntimeError:
pass
def apply_netease_window_fallback(info: dict[str, Any], cfg: dict[str, Any]) -> dict[str, Any]:
"""用窗口标题修正歌曲名/艺人,用音量峰值+进度变化判断播放状态。
参考 now-playing-service/NeteaseMusicService.cs。"""
global _LAST_NETEASE_TITLE, _LAST_NETEASE_AT, _LAST_PROGRESS_SECONDS, _LAST_PROGRESS_CHANGE_AT
holdover = max(0, int(cfg.get("holdover_ms", 1500))) / 1000
win_title = read_netease_window_title()
now = time.time()
if win_title:
_LAST_NETEASE_TITLE = win_title
_LAST_NETEASE_AT = now
elif _LAST_NETEASE_TITLE and now - _LAST_NETEASE_AT <= holdover:
win_title = _LAST_NETEASE_TITLE
if win_title:
title, artist = _split_title_artist(win_title)
if title:
info["title"] = title
if artist:
info["artist"] = artist
info["source"] = "cloudmusic.window"
# 进度
parsed = read_netease_progress_uia()
if parsed:
progress, duration = parsed
info["progress"] = progress
info["duration"] = duration
if progress != _LAST_PROGRESS_SECONDS:
_LAST_PROGRESS_SECONDS = progress
_LAST_PROGRESS_CHANGE_AT = now
progress_changed_recently = (now - _LAST_PROGRESS_CHANGE_AT) <= holdover
# 播放状态:音量峰值优先,进度兜底
volume_peak = _get_cloudmusic_audio_peak()
info["playing"] = _decide_playing(volume_peak, progress_changed_recently, holdover)
return info
async def get_media_info(cfg: dict[str, Any], previous_cover_hash: str = ""):
"""网易云已开启 SMTC 后:优先信任 SMTC 的标题、封面、播放状态;CDP/窗口标题只做兜底。"""
is_netease = any("cloudmusic" in str(t).lower() for t in cfg.get("targets", []))
smtc = await read_smtc_snapshot(cfg)
info = {
"title": "",
"artist": "",
"duration": 0,
"progress": 0,
"playing": False,
"cover": "",
"cover_hash": previous_cover_hash,
"source": "smtc",
"updated_at": time.time(),
}
cover_url = ""
cover_hash = previous_cover_hash
if smtc is not None:
info["title"] = smtc.get("title", "")
info["artist"] = smtc.get("artist", "")
info["duration"] = float(smtc.get("duration") or 0)
info["progress"] = float(smtc.get("progress") or 0)
info["source"] = smtc.get("source") or "smtc"
info["playing"] = bool(smtc.get("playing"))
thumb_data = smtc.get("thumb_data")
if thumb_data:
cover_url, cover_hash = write_cover_atomic(thumb_data, previous_cover_hash)
if is_netease:
# SMTC 现在是主数据源;如果 SMTC 某些字段缺失,再用 CDP/窗口标题补齐。
# 现在网易云已开启 SMTC,音乐显示主链路不再碰 CDP/WebSocket,避免卡主循环。
# 如果 SMTC 暂时没给标题,再用窗口标题兜底;封面则保留上一轮,不再网络下载。
if not info.get("title"):
info = apply_netease_window_fallback(info, cfg)
if not info.get("title"):
return None
elif smtc is None:
return None
# 封面绝不因为某一轮没读到就清空,避免前台闪烁。
if cover_url:
info["cover"] = cover_url
elif previous_cover_hash and COVER_FILE.exists():
info["cover"] = "/music_cover.jpg"
else:
# SMTC/窗口标题没有封面时,按“歌名 + 歌手”异步补封面;成功后下一轮自动显示。
apply_cached_or_schedule_cover(info, cfg)
if not info.get("cover"):
info["cover"] = ""
info["cover_hash"] = info.get("cover_hash") or cover_hash
return info
def default_state() -> dict[str, Any]:
return {
"playing": False,
"current": {
"title": "暂无歌曲",
"artist": "未接入音乐源",
"cover": "",
"duration": 0,
"progress": 0,
"source": "",
},
"playlist": [],
"requests": [],
"monitor": {"online": False, "source": ""},
}
def load_state() -> dict[str, Any]:
if MUSIC_FILE.exists():
try:
state = json.loads(MUSIC_FILE.read_text(encoding="utf-8"))
if isinstance(state, dict):
return state
except Exception:
pass
return default_state()
async def save_state(state: dict[str, Any]):
MUSIC_FILE.parent.mkdir(parents=True, exist_ok=True)
tmp = MUSIC_FILE.with_suffix(".tmp")
tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
tmp.replace(MUSIC_FILE)
async def post_state(state: dict[str, Any], api_url: str) -> bool:
try:
async with aiohttp.ClientSession() as session:
async with session.post(api_url, json=state) as resp:
return resp.status == 200
except Exception:
return False
async def run_monitor(port: int, interval: float | None = None, cli_targets=None, allow_all: bool | None = None):
api_url = f"http://localhost:{port}/api/music"
state = load_state()
last_info_time = 0.0
print(f"[MusicMonitor] start api={api_url}", flush=True)
while True:
cfg = load_monitor_config()
if interval is not None:
cfg["interval_sec"] = interval
if cli_targets:
cfg["targets"] = cli_targets
if allow_all is not None:
cfg["allow_all"] = allow_all
sleep_sec = max(0.3, float(cfg.get("interval_sec", 1.0)))
holdover_sec = max(0, int(cfg.get("holdover_ms", 1500))) / 1000
keep_last = bool(cfg.get("keep_last_when_none", True))
try:
previous_cover_hash = ((state.get("current") or {}).get("cover_hash") or "")
info = await get_media_info(cfg, previous_cover_hash)
now = time.time()
resume_reason = await maybe_auto_resume(info, cfg) if info is not None else ""
if resume_reason:
info["playing"] = True
reason_text = "进度停滞" if resume_reason == "stalled" else "检测到暂停"
print(f"[MusicMonitor] auto resume playback ({reason_text})", flush=True)
if info is None:
monitor = state.setdefault("monitor", {})
monitor["online"] = False
monitor["updated_at"] = now
monitor["source"] = ""
monitor["platform"] = cfg.get("platform", "")
monitor["targets"] = cfg.get("targets", [])
monitor["allow_all"] = bool(cfg.get("allow_all", False))
within_holdover = bool(last_info_time) and now - last_info_time <= holdover_sec
if keep_last and within_holdover:
# SMTC 偶发丢一帧时短暂保留,避免页面闪烁。
pass
else:
# 超过保留时间后不能继续把旧歌曲伪装成“已暂停”。
# 仅清理播放器状态,保留点歌队列等独立数据。
empty = default_state()
state["playing"] = False
state["current"] = empty["current"]
if not keep_last:
state["playlist"] = []
else:
last_info_time = now
state["playing"] = info["playing"]
prev_cover = (state.get("current") or {}).get("cover", "")
prev_cover_hash = (state.get("current") or {}).get("cover_hash", "")
# 封面: 新值优先, 空值保留旧值(防闪)
new_cover = info.get("cover", "") or prev_cover
new_cover_hash = info.get("cover_hash", "") or prev_cover_hash
# 如果异步封面任务刚完成,本轮 info 可能尚未带 cover;这里再查一次缓存。
if not new_cover:
cached = _COVER_CACHE.get(cover_cache_key(info.get("title", ""), info.get("artist", "")))
if cached:
new_cover, new_cover_hash = cached
state["current"] = {
"title": info["title"],
"artist": info["artist"],
"cover": new_cover,
"cover_hash": new_cover_hash,
"duration": info["duration"],
"progress": info["progress"],
"source": info["source"],
}
state["monitor"] = {
"online": True,
"source": info["source"],
"platform": cfg.get("platform", ""),
"updated_at": info["updated_at"],
"targets": cfg.get("targets", []),
"allow_all": bool(cfg.get("allow_all", False)),
}
ok = await post_state(state, api_url)
if not ok:
await save_state(state)
cur = state.get("current") or {}
print(f"[MusicMonitor] {cur.get('title', '')} - {cur.get('artist', '')} "
f"({cur.get('progress', 0):.0f}s/{cur.get('duration', 0):.0f}s) "
f"playing={state.get('playing')} source={cur.get('source', '')}", flush=True)
except Exception as e:
print(f"[MusicMonitor] error: {e}", flush=True)
traceback.print_exc()
await asyncio.sleep(sleep_sec)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="BGI 直播间音乐状态同步")
parser.add_argument("--port", type=int, default=8086, help="本地 Web 服务端口号")
parser.add_argument("--interval", type=float, default=None, help="轮询间隔(秒),默认读取 config/config.json")
parser.add_argument("--allow-all", action="store_true", help="允许同步任意媒体会话(适合网页版播放器)")
parser.add_argument("--filter", default="", help="额外过滤关键字,匹配 source_app_display_name/aumid")
args = parser.parse_args()
targets = None
if args.filter:
cfg = load_monitor_config()
targets = list(cfg.get("targets", [])) + [args.filter]
asyncio.run(run_monitor(args.port, args.interval, targets, True if args.allow_all else None))
+320
View File
@@ -0,0 +1,320 @@
"""网易云音乐二维码登录会话。"""
from __future__ import annotations
import asyncio
import base64
import http.cookiejar
import http.cookies
import io
import json
import logging
import secrets
import string
import time
import urllib.parse
import urllib.request
from typing import Any, Callable
from Crypto.Cipher import AES
try:
from .netease_resolver import NeteaseResolver
except ImportError:
from netease_resolver import NeteaseResolver
_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
_WEAPI_NONCE = b"0CoJUm6Qyw8W8jud"
_WEAPI_IV = b"0102030405060708"
_WEAPI_PUBLIC_EXPONENT = 0x10001
_WEAPI_MODULUS = int(
"00e0b509f6259df8642dbc35662901477df22677ec152b5ff68ace615bb7b725"
"152b3ab17a876aea8a5aa76d2e417629ec4ee341f56135fccf695280104e0312"
"ecbda92557c93870114af6c9d05c4f7f0c3685b7a46bee255932575cce10b424"
"d813cfe4875d3e82047b97ddef52741d546b8e289dc6935b3ece0462db0a22b8"
"e7",
16,
)
_SECRET_ALPHABET = string.ascii_letters + string.digits
def _aes_encrypt_base64(content: bytes, key: bytes) -> bytes:
padding = AES.block_size - (len(content) % AES.block_size)
padded = content + bytes([padding]) * padding
encrypted = AES.new(key, AES.MODE_CBC, _WEAPI_IV).encrypt(padded)
return base64.b64encode(encrypted)
def _weapi_form(data: dict[str, Any], *, secret_key: str | None = None) -> dict[str, str]:
secret = secret_key or "".join(secrets.choice(_SECRET_ALPHABET) for _ in range(16))
if len(secret.encode("ascii")) != 16:
raise ValueError("weapi secret key must be 16 ASCII bytes")
serialized = json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
first_pass = _aes_encrypt_base64(serialized, _WEAPI_NONCE)
params = _aes_encrypt_base64(first_pass, secret.encode("ascii")).decode("ascii")
reversed_secret = int.from_bytes(secret[::-1].encode("ascii"), "big")
enc_sec_key = format(
pow(reversed_secret, _WEAPI_PUBLIC_EXPONENT, _WEAPI_MODULUS),
"x",
).zfill(256)
return {"params": params, "encSecKey": enc_sec_key}
def _render_qr_png(content: str) -> bytes:
try:
import qrcode
from qrcode.constants import ERROR_CORRECT_M
except ImportError as exc:
raise RuntimeError("缺少 qrcode 依赖,请重新安装 requirements.txt") from exc
qr = qrcode.QRCode(
version=None,
error_correction=ERROR_CORRECT_M,
box_size=8,
border=3,
)
qr.add_data(content)
qr.make(fit=True)
image = qr.make_image(fill_color="black", back_color="white")
output = io.BytesIO()
image.save(output, format="PNG")
return output.getvalue()
def _request_json(
url: str,
*,
opener: urllib.request.OpenerDirector | None = None,
data: dict[str, Any] | None = None,
) -> tuple[dict[str, Any], Any]:
encoded_data = None
if data is not None:
encoded_data = urllib.parse.urlencode(data).encode("utf-8")
request = urllib.request.Request(
url,
data=encoded_data,
headers={
"User-Agent": _USER_AGENT,
"Referer": "https://music.163.com/",
"Content-Type": "application/x-www-form-urlencoded",
},
)
response = (opener or urllib.request.build_opener()).open(request, timeout=15)
payload = json.loads(response.read().decode("utf-8", errors="replace"))
if not isinstance(payload, dict):
raise RuntimeError("网易云登录接口返回格式无效")
return payload, response
def _request_weapi_json(
url: str,
data: dict[str, Any],
*,
opener: urllib.request.OpenerDirector,
secret_key: str | None = None,
) -> tuple[dict[str, Any], Any]:
return _request_json(
url,
opener=opener,
data=_weapi_form(data, secret_key=secret_key),
)
def _cookie_values(jar: http.cookiejar.CookieJar, response: Any, payload: dict[str, Any]) -> dict[str, str]:
values = {cookie.name: cookie.value for cookie in jar}
headers = getattr(response, "headers", None)
raw_headers = headers.get_all("Set-Cookie", []) if headers and hasattr(headers, "get_all") else []
for raw_header in raw_headers:
parsed = http.cookies.SimpleCookie()
try:
parsed.load(raw_header)
except http.cookies.CookieError:
continue
values.update({name: morsel.value for name, morsel in parsed.items()})
for raw_cookie in (
payload.get("cookie"),
(payload.get("data") or {}).get("cookie") if isinstance(payload.get("data"), dict) else None,
):
if not raw_cookie:
continue
parsed = http.cookies.SimpleCookie()
try:
parsed.load(str(raw_cookie))
except http.cookies.CookieError:
continue
values.update({name: morsel.value for name, morsel in parsed.items()})
return values
class NeteaseQrLogin:
"""服务端保存二维码 key,登录成功后只把 MUSIC_U 交给保存回调。"""
_STATUS_MESSAGES = {
"idle": "尚未开始扫码登录",
"awaiting_scan": "请使用网易云音乐客户端扫码",
"awaiting_confirm": "已扫码,请在手机上确认登录",
"completed": "登录成功,MUSIC_U 已自动保存",
"expired": "二维码已过期,请重新生成",
"failed": "网易云扫码登录失败",
}
def __init__(
self,
*,
api_base: Callable[[], str],
get_saved_music_u: Callable[[], str],
save_music_u: Callable[[str, dict[str, Any]], Any],
logger: logging.Logger,
account_checker: Callable[[str, str], Any] | None = None,
ttl_seconds: int = 180,
) -> None:
self.api_base = api_base
self.get_saved_music_u = get_saved_music_u
self.save_music_u = save_music_u
self.logger = logger
self.account_checker = account_checker
self.ttl_seconds = max(60, int(ttl_seconds))
self._lock = asyncio.Lock()
self._session: dict[str, Any] | None = None
def _base_url(self) -> str:
return str(self.api_base() or "https://music.163.com").rstrip("/")
def snapshot(self) -> dict[str, Any]:
session = self._session or {}
state = str(session.get("state") or "idle")
expires_at = float(session.get("expires_at") or 0)
expires_in = max(0, int(expires_at - time.time())) if expires_at else 0
return {
"success": True,
"state": state,
"message": str(session.get("message") or self._STATUS_MESSAGES.get(state, "")),
"expires_in": expires_in,
"has_qr_image": state in {"awaiting_scan", "awaiting_confirm"} and expires_in > 0,
"credential_configured": bool(str(self.get_saved_music_u() or "").strip()),
"account": session.get("account"),
}
def _start_sync(self) -> dict[str, Any]:
jar = http.cookiejar.CookieJar()
opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
payload, _ = _request_weapi_json(
f"{self._base_url()}/weapi/login/qrcode/unikey",
{"type": 1, "csrf_token": ""},
opener=opener,
)
key = str(payload.get("unikey") or (payload.get("data") or {}).get("unikey") or "").strip()
if int(payload.get("code") or 0) != 200 or not key:
raise RuntimeError(f"网易云二维码申请失败 code={payload.get('code')}")
now = time.time()
return {
"state": "awaiting_scan",
"message": self._STATUS_MESSAGES["awaiting_scan"],
"created_at": now,
"expires_at": now + self.ttl_seconds,
"key": key,
# The NetEase client currently accepts the exact HTTP URL emitted by
# the official web login page. Using HTTPS shows an unsupported-login warning.
"qr_url": f"http://music.163.com/login?codekey={key}",
"jar": jar,
"opener": opener,
"account": None,
}
async def start(self) -> dict[str, Any]:
async with self._lock:
try:
self._session = await asyncio.to_thread(self._start_sync)
except Exception as exc:
self._session = {
"state": "failed",
"message": str(exc) or type(exc).__name__,
"expires_at": 0,
}
self.logger.warning("[网易云扫码登录] 二维码申请失败: %s", type(exc).__name__)
return self.snapshot()
self.logger.info("[网易云扫码登录] 二维码已生成,等待扫码")
return self.snapshot()
def _poll_sync(self, session: dict[str, Any]) -> dict[str, Any]:
url = f"{self._base_url()}/weapi/login/qrcode/client/login"
payload, response = _request_weapi_json(
url,
{"key": session["key"], "type": 1, "csrf_token": ""},
opener=session["opener"],
)
code = int(payload.get("code") or 0)
if code == 801:
return {"state": "awaiting_scan", "message": self._STATUS_MESSAGES["awaiting_scan"]}
if code == 802:
return {"state": "awaiting_confirm", "message": self._STATUS_MESSAGES["awaiting_confirm"]}
if code == 800:
return {"state": "expired", "message": self._STATUS_MESSAGES["expired"]}
if code != 803:
return {
"state": "failed",
"message": str(payload.get("message") or f"网易云扫码状态异常 code={code}"),
}
music_u = str(_cookie_values(session["jar"], response, payload).get("MUSIC_U") or "").strip()
if not music_u:
return {"state": "failed", "message": "扫码成功但响应中缺少 MUSIC_U"}
return {"state": "authenticated", "music_u": music_u}
async def _check_account(self, music_u: str) -> dict[str, Any]:
if self.account_checker:
result = self.account_checker(self._base_url(), music_u)
return await result if asyncio.iscoroutine(result) else result
resolver = NeteaseResolver(self.logger, api_base=self._base_url(), music_u=music_u)
return await resolver.account_status()
async def poll(self) -> dict[str, Any]:
async with self._lock:
if not self._session:
return self.snapshot()
state = str(self._session.get("state") or "idle")
if state in {"completed", "expired", "failed"}:
return self.snapshot()
if time.time() >= float(self._session.get("expires_at") or 0):
self._session.update(state="expired", message=self._STATUS_MESSAGES["expired"])
return self.snapshot()
try:
result = await asyncio.to_thread(self._poll_sync, self._session)
if result.get("state") != "authenticated":
self._session.update(result)
return self.snapshot()
music_u = str(result.get("music_u") or "")
account = await self._check_account(music_u)
if not account.get("authenticated"):
self._session.update(state="failed", message="扫码 Cookie 登录验证失败")
return self.snapshot()
callback_result = self.save_music_u(music_u, account)
if asyncio.iscoroutine(callback_result):
await callback_result
self._session.update(
state="completed",
message=self._STATUS_MESSAGES["completed"],
account={
"user_id": str(account.get("user_id") or ""),
"nickname": str(account.get("nickname") or "网易云用户"),
"vip_type": int(account.get("vip_type") or 0),
},
)
self.logger.info("[网易云扫码登录] 登录成功,MUSIC_U 已自动保存")
return self.snapshot()
except Exception as exc:
self._session.update(
state="failed",
message=str(exc) or f"扫码状态查询异常: {type(exc).__name__}",
)
self.logger.warning("[网易云扫码登录] 状态查询或保存失败: %s", type(exc).__name__)
return self.snapshot()
async def qr_png(self) -> bytes:
async with self._lock:
if not self._session or self._session.get("state") not in {"awaiting_scan", "awaiting_confirm"}:
raise RuntimeError("当前没有可用的网易云登录二维码")
if time.time() >= float(self._session.get("expires_at") or 0):
self._session.update(state="expired", message=self._STATUS_MESSAGES["expired"])
raise RuntimeError("网易云登录二维码已过期")
content = str(self._session.get("qr_url") or "")
return await asyncio.to_thread(_render_qr_png, content)
+224
View File
@@ -0,0 +1,224 @@
"""Resolve NetEase song IDs to short-lived playable audio URLs."""
from __future__ import annotations
import logging
import re
import urllib.parse
from typing import Any
import aiohttp
class NeteaseResolver:
def __init__(
self,
logger: logging.Logger,
*,
api_base: str = "https://music.163.com",
music_u: str = "",
):
self.logger = logger
self.api_base = api_base.rstrip("/")
self.music_u = str(music_u or "").strip()
self.last_error_code = ""
def update_api_base(self, api_base: str):
self.api_base = str(api_base or "https://music.163.com").rstrip("/")
def update_auth(self, music_u: str):
self.music_u = str(music_u or "").strip()
def _headers(self) -> dict[str, str]:
cookie = "os=pc; appver=2.9.8;"
if self.music_u:
cookie += f" MUSIC_U={self.music_u};"
return {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
"Referer": "https://music.163.com/",
"Cookie": cookie,
"X-Real-IP": "218.75.111.114",
"X-Forwarded-For": "218.75.111.114",
}
async def _probe_url(self, url: str) -> str:
timeout = aiohttp.ClientTimeout(total=12)
async with aiohttp.ClientSession(headers=self._headers(), timeout=timeout) as session:
try:
async with session.get(url, allow_redirects=True, headers={"Range": "bytes=0-1"}) as resp:
content_type = str(resp.headers.get("Content-Type", "")).lower()
if resp.status in (200, 206) and ("audio" in content_type or "octet-stream" in content_type):
return str(resp.url)
except Exception:
return ""
return ""
async def _get_json(self, url: str) -> dict[str, Any] | None:
timeout = aiohttp.ClientTimeout(total=12)
try:
async with aiohttp.ClientSession(headers=self._headers(), timeout=timeout) as session:
async with session.get(url, allow_redirects=True) as resp:
if resp.status != 200:
return None
payload = await resp.json(content_type=None)
return payload if isinstance(payload, dict) else None
except Exception:
return None
async def account_status(self) -> dict[str, Any]:
if not self.music_u:
return {"authenticated": False}
payload = await self._get_json(f"{self.api_base}/api/nuser/account/get")
if not payload:
return {"authenticated": False}
account = payload.get("account")
profile = payload.get("profile")
if not isinstance(account, dict) or not account.get("id"):
return {"authenticated": False}
profile = profile if isinstance(profile, dict) else {}
return {
"authenticated": True,
"user_id": str(account.get("id") or ""),
"nickname": str(profile.get("nickname") or "网易云用户"),
"vip_type": int(account.get("vipType") or profile.get("vipType") or 0),
}
async def _fetch_player_entry(self, song_id: str) -> dict[str, Any] | None:
url = (
f"{self.api_base}/api/song/enhance/player/url"
f"?ids=%5B{urllib.parse.quote(song_id)}%5D&br=320000"
)
payload = await self._get_json(url)
entries = payload.get("data") if payload else None
if not isinstance(entries, list) or not entries or not isinstance(entries[0], dict):
return None
return entries[0]
@staticmethod
def _is_trial_entry(entry: dict[str, Any]) -> bool:
if entry.get("freeTrialInfo"):
return True
privilege = entry.get("freeTrialPrivilege")
if not isinstance(privilege, dict):
return False
return any(bool(privilege.get(key)) for key in ("resConsumable", "userConsumable", "listenType"))
@staticmethod
def playlist_id(value: str | int) -> str:
match = re.search(r"(?:playlist\?id=|\bid=)?(\d{5,})", str(value or ""))
return match.group(1) if match else ""
async def fetch_playlist(self, playlist: str | int) -> dict[str, Any] | None:
playlist_id = self.playlist_id(playlist)
if not playlist_id:
return None
timeout = aiohttp.ClientTimeout(total=15)
url = f"{self.api_base}/api/playlist/detail?id={urllib.parse.quote(playlist_id)}"
try:
async with aiohttp.ClientSession(headers=self._headers(), timeout=timeout) as session:
async with session.get(url, allow_redirects=True) as resp:
if resp.status != 200:
return None
payload = await resp.json(content_type=None)
except Exception as exc:
self.logger.warning(f"[mpv] 获取网易云歌单失败: {exc}")
return None
result = payload.get("result") if isinstance(payload, dict) else None
if not isinstance(result, dict):
return None
tracks = result.get("tracks")
if not isinstance(tracks, list):
return None
songs: list[dict[str, Any]] = []
for track in tracks:
if not isinstance(track, dict):
continue
song_id = self.playlist_id(track.get("id", ""))
if not song_id:
continue
artists_raw = track.get("artists") or track.get("ar") or []
artists = "/".join(
str(item.get("name") or "").strip()
for item in artists_raw
if isinstance(item, dict) and str(item.get("name") or "").strip()
)
album = track.get("album") or track.get("al") or {}
duration_ms = track.get("duration") or track.get("dt") or 0
try:
duration_sec = max(0, int(duration_ms) // 1000)
except (TypeError, ValueError):
duration_sec = 0
songs.append({
"id": song_id,
"name": str(track.get("name") or f"歌曲{song_id}"),
"artist": artists or "未知歌手",
"duration_sec": duration_sec,
"cover": str(album.get("picUrl") or "") if isinstance(album, dict) else "",
"source": "background_playlist",
"playlist_id": playlist_id,
})
if not songs:
return None
return {
"id": playlist_id,
"name": str(result.get("name") or f"歌单{playlist_id}"),
"songs": songs,
}
async def fetch_song_detail(self, song_id: str | int) -> dict[str, Any] | None:
clean_id = self.playlist_id(song_id)
if not clean_id:
return None
timeout = aiohttp.ClientTimeout(total=12)
url = f"{self.api_base}/api/song/detail/?id={clean_id}&ids=[{clean_id}]"
try:
async with aiohttp.ClientSession(headers=self._headers(), timeout=timeout) as session:
async with session.get(url, allow_redirects=True) as resp:
payload = await resp.json(content_type=None)
except Exception:
return None
songs = payload.get("songs") if isinstance(payload, dict) else None
if not isinstance(songs, list) or not songs or not isinstance(songs[0], dict):
return None
track = songs[0]
album = track.get("album") or track.get("al") or {}
return {
"cover": str(album.get("picUrl") or album.get("blurPicUrl") or "") if isinstance(album, dict) else "",
}
async def resolve(self, song: dict[str, Any]) -> dict[str, Any] | None:
self.last_error_code = ""
song_id = "".join(ch for ch in str(song.get("id", "")) if ch.isdigit())
if not song_id:
self.last_error_code = "invalid_song_id"
return None
entry = await self._fetch_player_entry(song_id)
if entry:
if self._is_trial_entry(entry):
reason = "登录已失效或账号没有完整播放权益" if self.music_u else "未配置网易云登录"
self.logger.warning(
f"[mpv] 拒绝播放试听片段: {song.get('name')} ({song_id}){reason}"
)
self.last_error_code = "preview_only"
return None
player_url = str(entry.get("url") or "").strip()
if player_url:
resolved = await self._probe_url(player_url)
if resolved:
return {
"url": resolved,
"source": "netease.player.auth" if self.music_u else "netease.player",
"song_id": song_id,
}
# NetEase's public outer URL provides a short-lived CDN redirect for songs available to the current region/account tier.
outer = f"{self.api_base}/song/media/outer/url?id={urllib.parse.quote(song_id)}.mp3"
resolved = await self._probe_url(outer)
if not resolved:
self.last_error_code = "unavailable"
self.logger.warning(f"[mpv] 无法获取可播放地址: {song.get('name')} ({song_id})")
return None
return {
"url": resolved,
"source": "netease.outer",
"song_id": song_id,
}
+419
View File
@@ -0,0 +1,419 @@
"""兑换码的 SQLite 存储与原子领取逻辑。"""
from __future__ import annotations
import asyncio
import sqlite3
import threading
from contextlib import contextmanager
from datetime import UTC, datetime, timedelta, timezone
from pathlib import Path
from typing import Any
BEIJING_TZ = timezone(timedelta(hours=8), name="Asia/Shanghai")
def normalize_code(value: str) -> str:
return str(value or "").strip().casefold()
def utc_now() -> str:
return datetime.now(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z")
def parse_beijing_datetime(value: str) -> str:
raw = str(value or "").strip()
if not raw:
raise ValueError("生效时间和失效时间不能为空")
try:
parsed = datetime.fromisoformat(raw.replace("Z", "+00:00"))
except ValueError as exc:
raise ValueError("时间格式无效") from exc
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=BEIJING_TZ)
return parsed.astimezone(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z")
def to_beijing_datetime(value: str | None) -> str:
if not value:
return ""
parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=UTC)
return parsed.astimezone(BEIJING_TZ).isoformat(timespec="minutes")
_SCHEMA = """
CREATE TABLE IF NOT EXISTS redemption_codes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
code_normalized TEXT NOT NULL UNIQUE,
code_display TEXT NOT NULL,
points INTEGER NOT NULL CHECK(points > 0),
starts_at_utc TEXT NOT NULL,
ends_at_utc TEXT NOT NULL,
max_redemptions INTEGER CHECK(max_redemptions IS NULL OR max_redemptions > 0),
redeemed_count INTEGER NOT NULL DEFAULT 0 CHECK(redeemed_count >= 0),
enabled INTEGER NOT NULL DEFAULT 1 CHECK(enabled IN (0, 1)),
deleted_at_utc TEXT,
created_at_utc TEXT NOT NULL,
updated_at_utc TEXT NOT NULL,
CHECK(starts_at_utc < ends_at_utc)
);
CREATE TABLE IF NOT EXISTS redemption_records (
id INTEGER PRIMARY KEY AUTOINCREMENT,
code_id INTEGER NOT NULL REFERENCES redemption_codes(id),
code_display TEXT NOT NULL,
platform TEXT NOT NULL,
platform_user_id TEXT NOT NULL,
display_name TEXT NOT NULL,
points INTEGER NOT NULL,
balance_before INTEGER NOT NULL,
balance_after INTEGER,
status TEXT NOT NULL DEFAULT 'pending',
redeemed_at_utc TEXT NOT NULL,
completed_at_utc TEXT,
UNIQUE(code_id, platform, platform_user_id)
);
CREATE INDEX IF NOT EXISTS idx_redemption_codes_active
ON redemption_codes(deleted_at_utc, enabled, starts_at_utc, ends_at_utc);
CREATE INDEX IF NOT EXISTS idx_redemption_records_code_time
ON redemption_records(code_id, redeemed_at_utc DESC);
CREATE INDEX IF NOT EXISTS idx_redemption_records_user_time
ON redemption_records(platform, platform_user_id, redeemed_at_utc DESC);
"""
class RedemptionCodeStore:
"""用短事务处理兑换码,避免并发超领和重复领取。
当提供 stats_store 时,所有读写都复用 stats_store 的唯一连接与串行写队列,
消除多连接写同一文件导致的 "database is locked" 争用。
"""
def __init__(self, database_path: str | Path, stats_store: Any | None = None):
self.database_path = Path(database_path)
self._store = stats_store
self._lock = threading.RLock()
self.database_path.parent.mkdir(parents=True, exist_ok=True)
with self._connect() as connection:
connection.executescript(_SCHEMA)
connection.commit()
@contextmanager
def _connect(self):
connection = sqlite3.connect(self.database_path, timeout=5.0)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA journal_mode=WAL")
connection.execute("PRAGMA synchronous=NORMAL")
connection.execute("PRAGMA foreign_keys=ON")
connection.execute("PRAGMA busy_timeout=5000")
try:
yield connection
finally:
connection.close()
def _run(self, func, *args):
"""在 stats_store 单连接(或自有连接)上执行 func(connection, *args)。"""
if self._store is not None:
return self._store.execute(func, *args)
return asyncio.to_thread(self._run_direct, func, args)
def _run_direct(self, func, args):
with self._lock, self._connect() as connection:
return func(connection, *args)
@staticmethod
def _serialize_code(row: sqlite3.Row) -> dict[str, Any]:
max_redemptions = row["max_redemptions"]
redeemed_count = int(row["redeemed_count"] or 0)
return {
"id": int(row["id"]),
"code": row["code_display"],
"points": int(row["points"]),
"starts_at": to_beijing_datetime(row["starts_at_utc"]),
"ends_at": to_beijing_datetime(row["ends_at_utc"]),
"max_redemptions": int(max_redemptions) if max_redemptions is not None else None,
"redeemed_count": redeemed_count,
"remaining_count": max(0, int(max_redemptions) - redeemed_count) if max_redemptions is not None else None,
"enabled": bool(row["enabled"]),
"created_at": to_beijing_datetime(row["created_at_utc"]),
"updated_at": to_beijing_datetime(row["updated_at_utc"]),
}
@staticmethod
def _serialize_record(row: sqlite3.Row) -> dict[str, Any]:
return {
"id": int(row["id"]),
"code_id": int(row["code_id"]),
"code": row["code_display"],
"platform": row["platform"],
"uid": row["platform_user_id"],
"uname": row["display_name"],
"points": int(row["points"]),
"balance_before": int(row["balance_before"]),
"balance_after": int(row["balance_after"]) if row["balance_after"] is not None else None,
"status": row["status"],
"redeemed_at": to_beijing_datetime(row["redeemed_at_utc"]),
"completed_at": to_beijing_datetime(row["completed_at_utc"]),
}
async def create_code(
self,
*,
code: str,
points: int,
starts_at: str,
ends_at: str,
max_redemptions: int | None,
enabled: bool = True,
) -> dict[str, Any]:
return await self._run(
self._create_code,
code,
points,
starts_at,
ends_at,
max_redemptions,
enabled,
)
def _create_code(
self,
connection,
code: str,
points: int,
starts_at: str,
ends_at: str,
max_redemptions: int | None,
enabled: bool,
) -> dict[str, Any]:
display = str(code or "").strip()
normalized = normalize_code(display)
if not normalized:
raise ValueError("兑换码不能为空")
if len(display) > 100:
raise ValueError("兑换码不能超过100个字符")
points = int(points)
if points <= 0:
raise ValueError("兑换积分必须是正整数")
if max_redemptions in ("", None):
max_value = None
else:
max_value = int(max_redemptions)
if max_value <= 0:
raise ValueError("总兑换次数必须是正整数,或留空表示不限")
starts_utc = parse_beijing_datetime(starts_at)
ends_utc = parse_beijing_datetime(ends_at)
if starts_utc >= ends_utc:
raise ValueError("失效时间必须晚于生效时间")
now = utc_now()
try:
cursor = connection.execute(
"""INSERT INTO redemption_codes (
code_normalized, code_display, points, starts_at_utc, ends_at_utc,
max_redemptions, enabled, created_at_utc, updated_at_utc
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(normalized, display, points, starts_utc, ends_utc, max_value, int(bool(enabled)), now, now),
)
connection.commit()
except sqlite3.IntegrityError as exc:
connection.rollback()
if "code_normalized" in str(exc) or "UNIQUE constraint" in str(exc):
raise ValueError("兑换码已存在") from exc
raise
row = connection.execute("SELECT * FROM redemption_codes WHERE id=?", (cursor.lastrowid,)).fetchone()
assert row is not None
return self._serialize_code(row)
async def list_codes(self) -> list[dict[str, Any]]:
return await self._run(self._list_codes)
def _list_codes(self, connection) -> list[dict[str, Any]]:
rows = connection.execute(
"SELECT * FROM redemption_codes WHERE deleted_at_utc IS NULL ORDER BY id DESC"
).fetchall()
return [self._serialize_code(row) for row in rows]
async def list_records(self, *, code_id: int | None = None, limit: int = 500) -> list[dict[str, Any]]:
return await self._run(self._list_records, code_id, limit)
def _list_records(self, connection, code_id: int | None, limit: int) -> list[dict[str, Any]]:
limit = max(1, min(2000, int(limit)))
if code_id is None:
rows = connection.execute(
"SELECT * FROM redemption_records ORDER BY id DESC LIMIT ?", (limit,)
).fetchall()
else:
rows = connection.execute(
"SELECT * FROM redemption_records WHERE code_id=? ORDER BY id DESC LIMIT ?",
(int(code_id), limit),
).fetchall()
return [self._serialize_record(row) for row in rows]
async def set_enabled(self, code_id: int, enabled: bool) -> dict[str, Any]:
return await self._run(self._set_enabled, code_id, enabled)
def _set_enabled(self, connection, code_id: int, enabled: bool) -> dict[str, Any]:
cursor = connection.execute(
"UPDATE redemption_codes SET enabled=?, updated_at_utc=? WHERE id=? AND deleted_at_utc IS NULL",
(int(bool(enabled)), utc_now(), int(code_id)),
)
if cursor.rowcount != 1:
connection.rollback()
raise ValueError("兑换码不存在")
connection.commit()
row = connection.execute("SELECT * FROM redemption_codes WHERE id=?", (int(code_id),)).fetchone()
assert row is not None
return self._serialize_code(row)
async def delete_code(self, code_id: int) -> None:
await self._run(self._delete_code, code_id)
def _delete_code(self, connection, code_id: int) -> None:
now = utc_now()
cursor = connection.execute(
"""UPDATE redemption_codes
SET enabled=0, deleted_at_utc=?, updated_at_utc=?
WHERE id=? AND deleted_at_utc IS NULL""",
(now, now, int(code_id)),
)
if cursor.rowcount != 1:
connection.rollback()
raise ValueError("兑换码不存在")
connection.commit()
async def reserve(
self,
message: str,
*,
platform: str,
platform_user_id: str,
display_name: str,
balance_before: int,
) -> dict[str, Any]:
return await self._run(
self._reserve,
message,
platform,
platform_user_id,
display_name,
balance_before,
)
def _reserve(
self,
connection,
message: str,
platform: str,
platform_user_id: str,
display_name: str,
balance_before: int,
) -> dict[str, Any]:
normalized = normalize_code(message)
if not normalized:
return {"status": "unknown"}
now = utc_now()
connection.execute("BEGIN IMMEDIATE")
try:
code = connection.execute(
"SELECT * FROM redemption_codes WHERE code_normalized=? AND deleted_at_utc IS NULL",
(normalized,),
).fetchone()
if code is None:
connection.rollback()
return {"status": "unknown"}
if not bool(code["enabled"]):
connection.rollback()
return {"status": "disabled"}
if now < code["starts_at_utc"]:
connection.rollback()
return {"status": "not_started"}
if now >= code["ends_at_utc"]:
connection.rollback()
return {"status": "expired"}
existing = connection.execute(
"""SELECT status FROM redemption_records
WHERE code_id=? AND platform=? AND platform_user_id=?""",
(code["id"], platform, platform_user_id),
).fetchone()
if existing is not None:
connection.rollback()
return {"status": "already_redeemed"}
max_redemptions = code["max_redemptions"]
if max_redemptions is not None and int(code["redeemed_count"]) >= int(max_redemptions):
connection.rollback()
return {"status": "exhausted"}
cursor = connection.execute(
"""INSERT INTO redemption_records (
code_id, code_display, platform, platform_user_id, display_name,
points, balance_before, status, redeemed_at_utc
) VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?)""",
(
code["id"], code["code_display"], platform, platform_user_id,
display_name, code["points"], int(balance_before), now,
),
)
updated = connection.execute(
"""UPDATE redemption_codes
SET redeemed_count=redeemed_count+1, updated_at_utc=?
WHERE id=? AND (max_redemptions IS NULL OR redeemed_count < max_redemptions)""",
(now, code["id"]),
)
if updated.rowcount != 1:
connection.rollback()
return {"status": "exhausted"}
connection.commit()
return {
"status": "reserved",
"record_id": int(cursor.lastrowid),
"code_id": int(code["id"]),
"code": code["code_display"],
"points": int(code["points"]),
}
except sqlite3.IntegrityError:
connection.rollback()
return {"status": "already_redeemed"}
except Exception:
connection.rollback()
raise
async def finalize(self, record_id: int, balance_after: int) -> None:
await self._run(self._finalize, record_id, balance_after)
def _finalize(self, connection, record_id: int, balance_after: int) -> None:
cursor = connection.execute(
"""UPDATE redemption_records
SET status='completed', balance_after=?, completed_at_utc=?
WHERE id=? AND status='pending'""",
(int(balance_after), utc_now(), int(record_id)),
)
if cursor.rowcount != 1:
connection.rollback()
raise ValueError("兑换记录不存在或已经完成")
connection.commit()
async def cancel(self, record_id: int) -> None:
await self._run(self._cancel, record_id)
def _cancel(self, connection, record_id: int) -> None:
connection.execute("BEGIN IMMEDIATE")
try:
row = connection.execute(
"SELECT code_id FROM redemption_records WHERE id=? AND status='pending'",
(int(record_id),),
).fetchone()
if row is None:
connection.rollback()
return
connection.execute("DELETE FROM redemption_records WHERE id=?", (int(record_id),))
connection.execute(
"""UPDATE redemption_codes
SET redeemed_count=MAX(0, redeemed_count-1), updated_at_utc=?
WHERE id=?""",
(utc_now(), int(row["code_id"])),
)
connection.commit()
except Exception:
connection.rollback()
raise
+1733
View File
File diff suppressed because it is too large Load Diff
+127
View File
@@ -0,0 +1,127 @@
"""TTS 状态监控窗口
独立小黑窗显示 TTS 引擎状态、最近合成文本和耗时。
启动参数:
tts_monitor.py --state-file data/tts_state.json
"""
import argparse
import json
import os
import sys
import time
from pathlib import Path
APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
sys.path.insert(0, str(APP_DIR))
from core.runtime_paths import DATA_DIR, ensure_runtime_dirs
def clear():
os.system("cls" if os.name == "nt" else "clear")
def default_state() -> dict:
return {
"enabled": False,
"provider": "none",
"model_loaded": False,
"model_name": "",
"last_text": "",
"last_duration_ms": 0,
"last_error": "",
"total_synthesized": 0,
"total_errors": 0,
"recent_events": [],
"updated_at": 0,
}
def load_state(state_file: Path) -> dict:
if not state_file.exists():
return default_state()
try:
data = json.loads(state_file.read_text(encoding="utf-8"))
if isinstance(data, dict):
state = default_state()
state.update(data)
return state
except Exception:
pass
return default_state()
def fmt_duration(ms: int) -> str:
if ms <= 0:
return "-"
if ms < 1000:
return f"{ms}ms"
return f"{ms/1000:.2f}s"
def render(state: dict):
clear()
print("=" * 60)
print(" TTS 状态监控")
print("=" * 60)
print()
enabled = state.get("enabled", False)
provider = state.get("provider", "none")
model_loaded = state.get("model_loaded", False)
print(f" TTS 启用: {'' if enabled else ''}")
print(f" 引擎: {provider}")
print(f" 模型加载: {'完成' if model_loaded else '未加载/加载中'}")
print(f" 模型名: {state.get('model_name', '') or '-'}")
print()
print(f" 累计合成: {state.get('total_synthesized', 0)}")
print(f" 累计失败: {state.get('total_errors', 0)}")
print(f" 上次合成: {state.get('last_text', '') or '-'}")
print(f" 合成耗时: {fmt_duration(state.get('last_duration_ms', 0))}")
print()
last_error = state.get("last_error", "")
if last_error:
print(f" [错误] {last_error}")
print()
print("-" * 60)
print(" 最近事件")
print("-" * 60)
events = state.get("recent_events", [])
if not events:
print(" (无)")
else:
for ev in events[-8:]:
ts = ev.get("time", "")
msg = ev.get("msg", "")
print(f" {ts} {msg}")
print()
print(" 按 Ctrl+C 关闭本窗口")
def main():
ensure_runtime_dirs()
parser = argparse.ArgumentParser(description="TTS 状态监控")
parser.add_argument("--state-file", default=str(DATA_DIR / "tts_state.json"), help="TTS 状态文件路径")
args = parser.parse_args()
state_file = Path(args.state_file)
print("等待 TTS 状态更新...")
time.sleep(0.5)
try:
while True:
state = load_state(state_file)
render(state)
time.sleep(0.5)
except KeyboardInterrupt:
print("\nTTS 监控已关闭")
sys.exit(0)
if __name__ == "__main__":
main()
+10
View File
@@ -0,0 +1,10 @@
@echo off
cd /d "%~dp0"
net session >nul 2>&1
if errorlevel 1 (
powershell -Command "Start-Process '%~f0' -Verb RunAs -WorkingDirectory '%~dp0'"
exit /b
)
powershell -NoProfile -ExecutionPolicy Bypass -File "%~dp0scripts\build.ps1" %*
+313
View File
@@ -0,0 +1,313 @@
{
"_说明": "BetterGI弹幕联动配置。改完保存即可,程序支持热说明请重启生效。",
"bilibili": {
"_说明": "room_id 可填短号。请在浏览器登录B站后,把 SESSDATA、bili_jct、buvid3 填入 cookie 字段(或分别填入 sessdata/bili_jct/buvid3),才能监听弹幕和回复。",
"room_id": 1871977513,
"cookie": "",
"sessdata": "",
"bili_jct": "",
"buvid3": "",
"cookie_auto_refresh_enabled": true,
"cookie_check_interval_hours": 6
},
"bettergi": {
"_说明": "BetterGI.exe 的完整路径。work_dir 留空则自动取 exe 所在目录。",
"exe_path": "C:\\Program Files\\BetterGI\\BetterGI.exe",
"work_dir": ""
},
"global": {
"_说明": "default_cooldown: 冷却秒数防刷屏。admin_uids: 管理员UID绕过冷却。log_level: DEBUG/INFO/WARNING。restart_mode: gentle温和(执行中跳过)/aggressive激进(先杀BetterGI再重启,立即响应,适合直播)。",
"default_cooldown": 30,
"admin_uids": [
1123997326
],
"log_level": "INFO",
"restart_mode": "aggressive",
"log_file": "danmu_bettergi.log"
},
"rules": [],
"queue": {
"default_group": "薄荷",
"data_dir": "data",
"initial_points": 5,
"signin_points_min": 5,
"signin_points_max": 10,
"signin_reset_hour": 4,
"max_points": 30,
"points_per_minute": 1,
"admin_window_seconds": 90
},
"broadcast": {
"tts": {
"faster-qwen3-tts": {
"device": "cuda",
"model_name_or_path": "vendor/tts-model",
"language": "Chinese",
"ref_audio": "data/ref_audio.wav",
"ref_text": "凯茨莱茵家族的迪奥娜小姐,货物我确实收下了,再次感谢您选择狛荷屋",
"xvec_only": true,
"non_streaming_mode": true,
"chunk_size": 8,
"append_silence": true,
"streaming": false,
"cpu_threads": 3,
"cpu_affinity_count": 6,
"process_priority": "below_normal"
}
},
"enable_danmu_reply": true,
"enable_system_danmu": true,
"enable_tts": true,
"danmu_interval_sec": 1,
"tts_provider": "faster-qwen3-tts",
"tts_categories": {
"signin": false,
"queue": true,
"song_request": true,
"login": true,
"execution": true,
"points": true,
"help": true,
"reset": true,
"system": true,
"gift": true
},
"tts_queue": {
"max_pending": 8,
"playback_max_pending": 2,
"max_age_sec": 40,
"urgent_max_age_sec": 60,
"low_priority_max_age_sec": 30,
"warmup_on_start": true,
"rebuild_before_live": true
},
"gift_thanks": {
"enabled": true,
"tts": true,
"danmu": false,
"template": "感谢{uname}送出的{num}个{gift_name}",
"merge_window_sec": 2,
"dedupe_window_sec": 15,
"max_pending": 50
}
},
"frontend": {
"background_image": "/uploads/background.jpg?t=1783666624",
"background_opacity": 0.65,
"background_blur": 0,
"background_fit": "cover",
"theme": "classic"
},
"music_monitor": {
"platform": "netease",
"targets": [
"网易云音乐",
"Netease",
"CloudMusic",
"cloudmusic"
],
"allow_all": false,
"interval_sec": 1,
"holdover_ms": 1500,
"prefer_playing": true,
"keep_last_when_none": true,
"cover_enabled": true,
"auto_resume_enabled": false,
"auto_resume_interval_sec": 3,
"auto_resume_stall_sec": 10,
"extra_filter": "",
"request_player": {
"enabled": true,
"cost_points": 1,
"commands": [
"点歌",
"dg"
],
"api_base": "https://music.163.com",
"play_url_template": "https://music.163.com/#/song?id={id}",
"auto_open": true,
"play_when_idle": true,
"handoff_lead_sec": 1.2,
"max_duration_sec": 600,
"dedupe_history": true,
"dedupe_cooldown_sec": 3600,
"clear_on_start": true,
"play_method": "mpv",
"mpv_exe": "vendor/mpv/mpv.exe",
"mpv_stall_seconds": 12,
"fallback_to_netease": false,
"netease_music_u": "",
"background_playlist_enabled": true,
"background_playlist_url": "",
"background_playlist_refresh_sec": 3600,
"background_playlist_retry_sec": 30,
"cdp_port": 9222,
"auto_launch_cdp": true,
"client_process": "cloudmusic.exe",
"client_exe": "",
"search_hotkey": "ctrl+f",
"play_enter_count": 2,
"ui_wait_sec": 0.6,
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
}
},
"system": {
"enable_startup_shortcut": true,
"startup_bat": "run.bat",
"live_start_time": "08:30",
"live_end_time": "23:30",
"auto_reboot_enabled": false,
"auto_reboot_time": "23:30",
"reboot_after_stop_enabled": true,
"reboot_after_stop_delay_sec": 60,
"launch_bilibili_live_enabled": true,
"launch_bilibili_live_time": "08:50",
"bilibili_live_exe": "C:\\Program Files\\bililive\\livehime\\livehime.exe",
"launch_genshin_enabled": true,
"launch_genshin_time": "08:50",
"genshin_exe": "C:\\Program Files\\miHoYo Launcher\\games\\Genshin Impact Game\\YuanShen.exe",
"bilibili_push_enabled": true,
"bilibili_push_time": "09:00",
"bilibili_push_window_keyword": "直播姬",
"bilibili_push_click_x_ratio": 0.787,
"bilibili_push_click_y_ratio": 0.927,
"bilibili_stop_push_enabled": true,
"bilibili_stop_push_time": "23:00",
"bilibili_stop_push_click_x_ratio": 0.787,
"bilibili_stop_push_click_y_ratio": 0.927,
"bilibili_stop_push_confirm_enter": true
},
"commands": {
"queue": {
"enabled": true,
"aliases": [
"排队"
],
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
},
"signin": {
"enabled": true,
"aliases": [
"签到"
],
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
},
"login": {
"enabled": true,
"aliases": [
"上号"
],
"allowed_roles": [
"super_admin",
"pending_operator"
]
},
"confirm_yes": {
"enabled": true,
"aliases": [
"是"
],
"allowed_roles": [
"super_admin",
"pending_operator"
]
},
"confirm_no": {
"enabled": true,
"aliases": [
"不是"
],
"allowed_roles": [
"super_admin",
"pending_operator"
]
},
"run": {
"enabled": true,
"aliases": [
"执行",
"跑",
"开始"
],
"allowed_roles": [
"super_admin",
"active_operator"
]
},
"leave": {
"enabled": true,
"aliases": [
"退出",
"退出排队",
"取消排队"
],
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
},
"reset": {
"enabled": true,
"aliases": [
"重置"
],
"allowed_roles": [
"super_admin",
"pending_operator"
]
},
"points": {
"enabled": true,
"aliases": [
"积分"
],
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
},
"queue_list": {
"enabled": true,
"aliases": [
"队列"
],
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
},
"help": {
"enabled": true,
"aliases": [
"帮助"
],
"allowed_roles": [
"super_admin",
"active_operator",
"pending_operator",
"viewer"
]
}
}
}
Binary file not shown.
+12
View File
@@ -0,0 +1,12 @@
<!doctype html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>直播联动后台</title>
</head>
<body>
<div id="app"></div>
<script type="module" src="/src/main.js"></script>
</body>
</html>
+1303
View File
File diff suppressed because it is too large Load Diff
+16
View File
@@ -0,0 +1,16 @@
{
"name": "live-streaming-admin",
"private": true,
"version": "0.1.0",
"type": "module",
"scripts": {
"dev": "vite --host 127.0.0.1 --port 5173",
"build": "vite build"
},
"dependencies": {
"vue": "^3.5.0"
},
"devDependencies": {
"vite": "^7.0.0"
}
}
File diff suppressed because it is too large Load Diff
+799
View File
@@ -0,0 +1,799 @@
:root {
color-scheme: light;
--bg: #f5f7f6;
--panel: #ffffff;
--panel-soft: #f9faf9;
--line: #dfe6e2;
--line-strong: #cbd6d0;
--text: #15201b;
--muted: #66746d;
--green: #237b58;
--green-soft: #e8f4ee;
--blue: #326fd1;
--blue-soft: #eaf1ff;
--red: #c94444;
--red-soft: #fff0f0;
--amber: #a66a00;
--shadow: 0 8px 24px rgba(24, 36, 30, 0.08);
font-family: Inter, "Segoe UI", "Microsoft YaHei", system-ui, sans-serif;
}
* { box-sizing: border-box; }
body {
margin: 0;
background: var(--bg);
color: var(--text);
}
.auth-shell {
min-height: 100vh;
display: grid;
place-items: center;
padding: 24px;
}
.auth-card {
width: min(460px, 100%);
padding: 28px;
border: 1px solid var(--line);
border-radius: 14px;
background: var(--panel);
box-shadow: var(--shadow);
}
.auth-card h1 {
margin-bottom: 10px;
}
button, input, textarea, select {
font: inherit;
}
button {
height: 34px;
border: 1px solid var(--line-strong);
border-radius: 7px;
padding: 0 12px;
background: #fff;
color: var(--text);
cursor: pointer;
}
button:hover { border-color: #9fb1a8; }
button:disabled { opacity: .55; cursor: not-allowed; }
.primary {
color: #fff;
background: var(--green);
border-color: var(--green);
}
.secondary {
background: var(--panel-soft);
}
.danger {
color: var(--red);
background: var(--red-soft);
border-color: #efc3c3;
}
.shell {
display: grid;
grid-template-columns: 236px minmax(0, 1fr);
min-height: 100vh;
}
.sidebar {
position: sticky;
top: 0;
height: 100vh;
padding: 22px 16px;
background: #0f1d17;
color: #edf5f0;
}
.brand {
padding: 4px 6px 18px;
border-bottom: 1px solid rgba(255,255,255,.1);
}
.brand-title {
font-size: 20px;
font-weight: 800;
}
.brand-sub {
margin-top: 4px;
color: rgba(237,245,240,.62);
font-size: 12px;
}
.nav {
display: grid;
gap: 6px;
margin-top: 18px;
}
.nav button {
justify-content: flex-start;
width: 100%;
height: 38px;
border: 0;
color: rgba(237,245,240,.72);
background: transparent;
text-align: left;
}
.nav button.active,
.nav button:hover {
color: #fff;
background: rgba(255,255,255,.1);
}
.sidebar-footer {
position: absolute;
left: 22px;
right: 22px;
bottom: 22px;
display: flex;
gap: 8px;
align-items: center;
color: rgba(237,245,240,.72);
font-size: 13px;
}
.status-dot {
width: 8px;
height: 8px;
border-radius: 999px;
background: #9ca3af;
}
.status-dot.ok { background: #43c383; box-shadow: 0 0 0 4px rgba(67,195,131,.14); }
.status-dot.danger { background: #ef6666; box-shadow: 0 0 0 4px rgba(239,102,102,.14); }
.main {
padding: 24px 28px 48px;
}
.topbar {
display: flex;
justify-content: space-between;
align-items: flex-start;
margin-bottom: 18px;
}
h1, h2, p { margin: 0; }
h1 { font-size: 24px; }
h2 { font-size: 16px; margin-bottom: 14px; }
.topbar p, .muted { color: var(--muted); font-size: 13px; margin-top: 5px; }
.top-actions, .actions-row, .table-actions, .panel-head {
display: flex;
align-items: center;
gap: 8px;
}
.panel-head { justify-content: space-between; margin-bottom: 12px; }
.panel-head h2 { margin-bottom: 0; }
.message-stack {
position: fixed;
top: 18px;
left: 50%;
transform: translateX(-50%);
z-index: 100;
display: grid;
gap: 8px;
width: min(360px, calc(100vw - 32px));
pointer-events: none;
}
.message {
padding: 10px 12px;
border: 1px solid var(--line);
border-radius: 8px;
background: rgba(255, 255, 255, .98);
box-shadow: var(--shadow);
color: var(--text);
font-size: 13px;
line-height: 1.5;
overflow-wrap: anywhere;
}
.message.success {
color: var(--green);
border-color: #bfdfce;
background: rgba(232, 244, 238, .98);
}
.message.error {
color: var(--red);
border-color: #efc3c3;
background: rgba(255, 240, 240, .98);
}
.stack { display: grid; gap: 14px; }
.metrics {
display: grid;
grid-template-columns: repeat(4, minmax(0, 1fr));
gap: 14px;
}
.metric, .panel {
background: var(--panel);
border: 1px solid var(--line);
border-radius: 8px;
box-shadow: var(--shadow);
}
.metric { padding: 16px; }
.metric span { display: block; color: var(--muted); font-size: 13px; }
.metric strong { display: block; margin-top: 8px; font-size: 26px; }
.panel { padding: 16px; }
.form-grid {
display: grid;
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 14px;
}
.coordinate-grid {
display: grid;
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 10px;
}
.redemption-create-grid {
display: grid;
grid-template-columns: repeat(3, minmax(0, 1fr));
gap: 0 12px;
}
.redemption-enabled {
min-height: 34px;
align-self: end;
margin-bottom: 12px;
color: var(--text);
}
.warning-text {
margin: 8px 0 12px;
color: var(--red);
font-size: 13px;
line-height: 1.6;
}
.danger-panel {
border-color: #efc3c3;
background: #fffafa;
}
label {
display: grid;
gap: 6px;
color: var(--muted);
font-size: 13px;
margin-bottom: 12px;
}
input, textarea, select {
width: 100%;
min-height: 34px;
border: 1px solid var(--line-strong);
border-radius: 7px;
padding: 7px 9px;
background: #fff;
color: var(--text);
}
textarea { min-height: 78px; resize: vertical; }
input:focus, textarea:focus, select:focus {
outline: 2px solid rgba(35,123,88,.18);
border-color: var(--green);
}
.inline {
display: flex;
align-items: center;
gap: 8px;
}
.inline input { width: auto; min-height: auto; }
.file-row input { padding: 6px; }
.bilibili-qr-login,
.netease-qr-login {
display: grid;
gap: 14px;
margin-top: 16px;
padding-top: 16px;
border-top: 1px solid var(--line);
}
.bilibili-qr-login .panel-head,
.netease-qr-login .panel-head {
align-items: center;
margin-bottom: 0;
}
.bilibili-qr-login .panel-head > div,
.netease-qr-login .panel-head > div,
.bilibili-qr-state,
.bilibili-account-state {
display: grid;
gap: 4px;
}
.qr-status {
color: var(--muted);
font-size: 12px;
}
.qr-status.awaiting_confirm,
.qr-status.completed {
color: var(--green);
}
.qr-status.expired,
.qr-status.failed {
color: var(--red);
}
.bilibili-qr-body {
display: grid;
grid-template-columns: 196px minmax(0, 1fr);
align-items: center;
gap: 18px;
}
.bilibili-qr-image {
width: 196px;
height: 196px;
border: 1px solid var(--line);
border-radius: 6px;
background: #fff;
object-fit: contain;
}
.bilibili-qr-state span,
.bilibili-account-state span {
color: var(--muted);
font-size: 13px;
}
.info-grid {
display: grid;
grid-template-columns: 120px minmax(0, 1fr);
gap: 10px 14px;
font-size: 14px;
}
.info-grid label {
margin: 0;
color: var(--muted);
}
.info-grid span {
min-width: 0;
overflow-wrap: anywhere;
}
.url-list {
display: grid;
gap: 8px;
}
.url-item {
display: grid;
grid-template-columns: 88px minmax(0, 1fr);
gap: 10px;
align-items: center;
min-height: 34px;
padding: 8px 10px;
border: 1px solid var(--line);
border-radius: 7px;
background: var(--panel-soft);
}
.url-item span {
color: var(--muted);
font-size: 13px;
}
.url-item a {
min-width: 0;
color: var(--blue);
overflow-wrap: anywhere;
text-decoration: none;
}
table {
width: 100%;
border-collapse: collapse;
font-size: 14px;
}
th, td {
padding: 10px 8px;
border-bottom: 1px solid var(--line);
text-align: left;
}
th {
color: var(--muted);
font-size: 12px;
font-weight: 700;
}
.search { width: 220px; }
.empty {
color: var(--muted);
padding: 12px 0;
}
.pill {
display: inline-flex;
align-items: center;
min-width: 86px;
height: 24px;
justify-content: center;
border-radius: 999px;
padding: 0 10px;
font-size: 12px;
font-weight: 700;
color: var(--muted);
background: #eef2ef;
}
.pill.running { color: var(--green); background: var(--green-soft); }
.pill.failed, .pill.degraded { color: var(--red); background: var(--red-soft); }
.pill.reconnecting, .pill.starting { color: var(--amber); background: #fff7e8; }
.pill.stopped { color: #6b7280; background: #f1f3f2; }
.error-cell {
max-width: 360px;
color: var(--red);
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.row-admin { background: var(--green-soft); }
.tag-admin {
display: inline-block;
margin-left: 6px;
padding: 1px 6px;
border-radius: 999px;
font-size: 11px;
font-weight: 700;
color: #fff;
background: var(--green);
vertical-align: middle;
}
.rule-row {
display: grid;
grid-template-columns: 1fr 132px 1.3fr 92px 96px 64px;
gap: 8px;
align-items: center;
margin-bottom: 8px;
}
.rule-row .inline { margin: 0; color: var(--text); }
.logs-grid {
display: grid;
grid-template-columns: 1fr 1fr;
gap: 14px;
}
.log {
min-height: 520px;
max-height: 620px;
overflow: auto;
margin: 0;
padding: 12px;
border: 1px solid var(--line);
border-radius: 8px;
background: #101814;
color: #dbe8e0;
font-size: 12px;
line-height: 1.6;
white-space: pre-wrap;
}
.raw-panel { display: grid; gap: 12px; }
.raw-editor {
min-height: 620px;
font-family: "Cascadia Code", Consolas, monospace;
font-size: 13px;
}
.cmd-grid {
display: grid;
gap: 8px;
}
.cmd-block {
padding: 12px;
border: 1px solid var(--line);
border-radius: 10px;
background: var(--panel-soft);
}
.cmd-row {
display: grid;
grid-template-columns: 160px minmax(0, 1fr);
gap: 10px;
align-items: center;
}
.cmd-toggle {
margin: 0;
color: var(--text);
}
.cmd-label {
font-weight: 600;
}
.cmd-arg-hint {
font-size: 11px;
color: var(--muted);
}
.cmd-aliases {
font-size: 13px;
}
.role-chip-row {
display: flex;
flex-wrap: wrap;
gap: 8px;
}
.role-chip-row.compact {
margin-top: 10px;
}
.role-chip {
display: inline-flex;
align-items: center;
gap: 8px;
margin: 0;
padding: 8px 10px;
border: 1px solid var(--line);
border-radius: 999px;
background: #fff;
color: var(--text);
}
.role-chip input {
width: auto;
min-height: auto;
}
.rule-card {
padding: 12px;
border: 1px solid var(--line);
border-radius: 10px;
background: var(--panel-soft);
}
.song-current {
display: grid;
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 12px;
margin-bottom: 16px;
}
.song-current > div {
display: grid;
gap: 4px;
padding: 12px;
border: 1px solid var(--line);
border-radius: 10px;
background: var(--panel-soft);
}
.song-current-card {
grid-template-columns: 72px minmax(0, 1fr) !important;
align-items: center;
}
.song-cover {
width: 72px;
height: 72px;
border-radius: 10px;
object-fit: cover;
background: #eef2f7;
}
.song-cover-placeholder {
display: grid;
place-items: center;
color: var(--muted);
font-size: 28px;
}
.song-current-text {
display: grid;
gap: 4px;
min-width: 0;
}
.song-current strong {
font-size: 16px;
}
.song-current small,
.small {
font-size: 12px;
color: var(--muted);
}
.list-block {
display: grid;
gap: 10px;
}
.list-row {
display: flex;
align-items: center;
gap: 10px;
justify-content: space-between;
padding: 10px 12px;
border: 1px solid var(--line);
border-radius: 8px;
background: var(--panel-soft);
}
.list-row strong {
min-width: 0;
}
.wide-search {
width: min(420px, 100%);
}
.mini-check {
display: inline-flex;
align-items: center;
gap: 6px;
margin: 0 8px 4px 0;
font-size: 12px;
color: var(--text);
}
.mini-check input {
width: auto;
min-height: auto;
}
@media (max-width: 1240px) {
.metrics, .form-grid, .logs-grid { grid-template-columns: 1fr; }
.redemption-create-grid { grid-template-columns: repeat(2, minmax(0, 1fr)); }
}
@media (max-width: 760px) {
html,
body {
max-width: 100%;
overflow-x: hidden;
}
.message-stack {
top: 12px;
left: 12px;
right: 12px;
transform: none;
width: auto;
}
.shell {
grid-template-columns: 1fr;
max-width: 100vw;
}
.sidebar {
position: sticky;
z-index: 20;
height: auto;
max-width: 100vw;
padding: 14px 12px;
}
.brand {
padding: 0 2px 12px;
}
.brand-title {
font-size: 18px;
}
.nav {
display: flex;
gap: 6px;
margin-top: 12px;
overflow-x: auto;
padding-bottom: 2px;
}
.nav button {
flex: 0 0 auto;
width: auto;
height: 34px;
padding: 0 10px;
}
.sidebar-footer {
position: static;
margin-top: 10px;
}
.main {
max-width: 100vw;
padding: 16px 12px 32px;
}
.topbar {
flex-direction: column;
gap: 12px;
}
.top-actions {
display: grid;
grid-template-columns: 1fr 1fr;
width: 100%;
}
.top-actions button {
flex: 1;
}
.metrics {
grid-template-columns: repeat(2, minmax(0, 1fr));
gap: 10px;
}
.metric {
padding: 13px;
}
.metric strong {
font-size: 22px;
}
.panel {
padding: 14px;
overflow-x: auto;
}
.info-grid {
grid-template-columns: 92px minmax(0, 1fr);
}
.url-item {
grid-template-columns: 1fr;
gap: 3px;
}
.rule-row {
grid-template-columns: 1fr;
}
.song-current {
grid-template-columns: 1fr;
}
.redemption-create-grid {
grid-template-columns: 1fr;
}
.bilibili-qr-body {
grid-template-columns: 1fr;
}
.bilibili-qr-image {
width: min(196px, 100%);
height: auto;
aspect-ratio: 1;
}
.actions-row,
.table-actions,
.panel-head {
flex-wrap: wrap;
}
.search {
width: 100%;
}
.raw-editor,
.log {
min-height: 360px;
}
}
+28
View File
@@ -0,0 +1,28 @@
import { defineConfig } from "vite";
export default defineConfig({
base: "/admin/",
resolve: {
alias: {
vue: "vue/dist/vue.esm-bundler.js"
}
},
define: {
__VUE_OPTIONS_API__: true,
__VUE_PROD_DEVTOOLS__: false,
__VUE_PROD_HYDRATION_MISMATCH_DETAILS__: false
},
build: {
outDir: "../../web/admin",
emptyOutDir: true,
sourcemap: false,
assetsInlineLimit: 0,
rollupOptions: {
output: {
entryFileNames: "assets/[name]-[hash].js",
chunkFileNames: "assets/[name]-[hash].js",
assetFileNames: "assets/[name]-[hash][extname]"
}
}
}
});
Binary file not shown.

After

Width:  |  Height:  |  Size: 764 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 503 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.5 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 608 B

+284
View File
@@ -0,0 +1,284 @@
(async function () {
// ========================================
// 扫码上号 v3
// 流程: 状态判定 → (已登录则先退出) → A0检测tap → A1选号 → A2等登录+点中心 → A3完成
// 状态文件: status.txt (登录中 / 已登录)
// ========================================
setGameMetrics(1920, 1080, 1);
// ---------- 状态文件 ----------
const STATUS_FILE = "status.txt";
function writeStatus(text) {
try {
file.writeTextSync(STATUS_FILE, text);
} catch (e) {
log.warn("写入状态文件失败: " + e);
}
}
// 脚本启动 → 写入"登录中"
writeStatus("登录中");
// ---------- 加载图像资源 ----------
const tapMat = file.readImageMatSync("assets/tap.png");
const a0PhoneMat = file.readImageMatSync("assets/a0_phone.png");
const loggedInMat = file.readImageMatSync("assets/btn_logged_in_real.png");
const paimonMat = file.readImageMatSync("assets/paimon_menu.png");
const exitDoorMat = file.readImageMatSync("assets/btn_exit_door.png");
const preLoginNoticeMat = file.readImageMatSync("assets/pre_login_notice.png");
// ---------- 工具函数 ----------
/**
* 在截屏中匹配指定图片, 返回识别结果(含 x/y/中心坐标)
*/
function findImageMatch(mat, x, y, w, h) {
const cap = captureGameRegion();
try {
const ro = RecognitionObject.TemplateMatch(mat, x || 0, y || 0, w || 1920, h || 1080);
const r = cap.find(ro);
if (r.isExist()) {
return { x: r.x, y: r.y, w: r.width, h: r.height };
}
return null;
} catch (e) {
log.error("图像识别失败: " + e);
return null;
} finally {
cap.dispose();
}
}
/**
* 严格匹配:用于 A0 入口判定,避免地图页面/传送点/圆形 UI 被低阈值误识别。
*/
function findImageMatchStrict(mat, x, y, w, h) {
const thresholds = [0.88, 0.84];
for (const t of thresholds) {
const cap = captureGameRegion();
try {
const ro = RecognitionObject.TemplateMatch(mat, x || 0, y || 0, w || 1920, h || 1080);
ro.threshold = t;
ro.Use3Channels = true;
const r = cap.find(ro);
if (r.isExist()) {
log.info(`严格找到图标 (阈值=${t}, x=${r.x}, y=${r.y})`);
return { x: r.x, y: r.y, w: r.width, h: r.height };
}
} catch (e) {
// 忽略, 继续下一个阈值
} finally {
cap.dispose();
}
}
return null;
}
/**
* 多阈值尝试匹配, 提高识别率
*/
function findImageMatchRobust(mat, x, y, w, h) {
const thresholds = [0.7, 0.6, 0.5];
for (const t of thresholds) {
const cap = captureGameRegion();
try {
const ro = RecognitionObject.TemplateMatch(mat, x || 0, y || 0, w || 1920, h || 1080);
ro.threshold = t;
ro.Use3Channels = true;
const r = cap.find(ro);
if (r.isExist()) {
log.info(`找到图标 (阈值=${t}, x=${r.x}, y=${r.y})`);
return { x: r.x, y: r.y, w: r.width, h: r.height };
}
} catch (e) {
// 忽略, 继续下一个阈值
} finally {
cap.dispose();
}
}
return null;
}
function clickImage(found) {
click(Math.round(found.x + found.w / 2), Math.round(found.y + found.h / 2));
}
async function waitForImage(mat, timeout, interval) {
interval = interval || 500;
const start = Date.now();
while (Date.now() - start < timeout) {
const found = findImageMatch(mat);
if (found) return found;
await sleep(interval);
}
return null;
}
// ========================================
// 前置处理:检测登录前的提示图标
// 检测到则点击 (1830, 985)1秒后点击 (1100, 675),再1秒后进入初始判定;
// 未检测到则直接进入初始判定,不影响原流程。
// ========================================
log.info("===== 前置检测: pre_login_notice 图标 =====");
const preNoticeFound = findImageMatchRobust(preLoginNoticeMat);
if (preNoticeFound) {
log.info(`检测到 pre_login_notice 图标 (x=${preNoticeFound.x}, y=${preNoticeFound.y})`);
log.info("点击 (1830, 985)...");
click(1830, 985);
await sleep(1000);
log.info("点击 (1100, 675)...");
click(1100, 675);
await sleep(1000);
} else {
log.info("未检测到 pre_login_notice 图标,跳过前置处理");
}
// ========================================
// 初始状态判定:先检测 tap + 手机图标
// 1. 同时存在:已经在开门页,直接进入 A0 选号
// 2. 不同时存在:检测派蒙头像;若无派蒙,则按 ESC 后重检,最多 8 次
// 3. 8 次仍无派蒙:再严格检测一次开门页;开门页也不存在才判定登录失败
// 4. 找到派蒙:说明已在游戏内,先退出到开门页,再进入 A0
// ========================================
log.info("===== 初始状态判定: 严格检测 tap 图标 + 手机图标 =====");
// A0 入口图标不能全屏低阈值搜索,否则地图/传送点等 UI 容易误判。
// 这里限定在开门页中下区域,并使用高阈值严格匹配。
const A0_X = 420;
const A0_Y = 120;
const A0_W = 1080;
const A0_H = 820;
let tapFound = findImageMatchStrict(tapMat, A0_X, A0_Y, A0_W, A0_H);
let phoneFound = findImageMatchStrict(a0PhoneMat, A0_X, A0_Y, A0_W, A0_H);
if (!tapFound || !phoneFound) {
log.info("未同时检测到 tap 图标和手机图标 → 开始检测派蒙头像");
let paimonFound = findImageMatchRobust(paimonMat);
for (let i = 0; !paimonFound && i < 8; i++) {
log.info(`${i + 1}/8 次未检测到派蒙头像,按 ESC 后重试...`);
keyPress("Escape");
await sleep(1000);
paimonFound = findImageMatchRobust(paimonMat);
}
if (!paimonFound) {
log.info("按 ESC 检测 8 次后仍未检测到派蒙头像 → 再严格检测一次开门页");
tapFound = findImageMatchStrict(tapMat, A0_X, A0_Y, A0_W, A0_H);
phoneFound = findImageMatchStrict(a0PhoneMat, A0_X, A0_Y, A0_W, A0_H);
if (!tapFound || !phoneFound) {
log.error("未检测到派蒙头像,最终也未同时检测到 tap 图标和手机图标,判定登录失败");
writeStatus("登录失败");
return;
}
log.info("最终严格检测到 tap 图标和手机图标 → 当前已在开门页,继续扫码流程");
}
if (paimonFound) {
// ========================================
// 已登录流程: 退出到开门页
// 找到派蒙→ESC关菜单→点55,1010→找exit_door→点击→等10s
// ========================================
log.info("检测到派蒙头像 → 已登录状态, 执行退出流程");
log.info(`找到派蒙头像 (x=${paimonFound.x}, y=${paimonFound.y}), 按 ESC 打开菜单...`);
keyPress("Escape");
await sleep(1000);
log.info("点击 (55, 1010) 打开派蒙菜单...");
click(55, 1010);
await sleep(1500);
log.info("查找退出图标 btn_exit_door...");
const doorFound = await waitForImage(exitDoorMat, 10 * 1000, 500);
if (doorFound) {
log.info("找到退出图标, 点击退出...");
clickImage(doorFound);
} else {
log.warn("未找到退出图标, 继续后续流程");
}
log.info("等待 8 秒, 等待退回到开门页...");
await sleep(8 * 1000);
log.info("退出流程完成, 重新严格检测 tap 图标 + 手机图标");
tapFound = findImageMatchStrict(tapMat, A0_X, A0_Y, A0_W, A0_H);
phoneFound = findImageMatchStrict(a0PhoneMat, A0_X, A0_Y, A0_W, A0_H);
if (!tapFound || !phoneFound) {
log.error("退出后仍未同时检测到 tap 图标和手机图标,判定登录失败");
writeStatus("登录失败");
return;
}
}
}
log.info("同时检测到 tap 图标和手机图标 → 准备进入扫码流程");
// ========================================
// A1: 进入扫码页
// 点击选号坐标 → 确认进入
// ========================================
log.info("===== A1: 进入扫码页 =====");
log.info("点击 (660, 250) 选择米游社账号...");
click(660, 250);
await sleep(750);
log.info("点击 (828, 691) 确认进入...");
click(828, 691);
log.info("A1 完成, 进入等待登录...");
// ========================================
// A2: 等待登录完成 → 点击画面进入游戏
// 等待 btn_logged_in_real 出现 (80秒超时)
// 找到后每秒点击 (960, 800), 每3秒检测派蒙头像
// ========================================
log.info("===== A2: 等待登录完成 =====");
log.info("等待扫码登录... 超时 80 秒");
const loginResult = await waitForImage(loggedInMat, 80 * 1000, 1000);
if (!loginResult) {
log.error("登录超时! 80 秒内未检测到登录完成图标");
writeStatus("登录失败");
return;
}
log.info(`检测到登录完成图标! (x=${loginResult.x}, y=${loginResult.y})`);
log.info("点击画面进入游戏...");
const CENTER_X = 960;
const CENTER_Y = 800;
let enteredGame = false;
for (let i = 0; i < 60; i++) {
click(CENTER_X, CENTER_Y);
await sleep(1000);
// 每 3 秒检测一次是否进入游戏 (匹配派蒙头像)
if (i % 3 === 2) {
const found = findImageMatch(paimonMat);
if (found) {
log.info(`检测到派蒙头像! 已进入游戏 (x=${found.x}, y=${found.y})`);
enteredGame = true;
break;
}
}
}
if (!enteredGame) {
log.error("60秒内未检测到最后的派蒙头像, 判定登录失败");
writeStatus("登录失败");
return;
}
// ========================================
// A3: 登录流程完成, 退出脚本
// ========================================
log.info("===== A3: 登录流程完成! =====");
log.info("扫码上号成功, 脚本退出");
// 写入"已登录"后按一次 ESC,收起菜单/退出可能残留的界面
writeStatus("已登录");
await sleep(500);
keyPress("Escape");
})();
+16
View File
@@ -0,0 +1,16 @@
{
"manifest_version": 1,
"name": "扫码上号",
"version": "1.0",
"bgi_version": "0.48.0",
"description": "直播用扫码上号脚本:判断在大世界则退出到开门页面→点击扫码登录→等待扫码→检测验证码",
"authors": [
{
"name": "BGI直播项目",
"links": "https://github.com/"
}
],
"settings_ui": "settings.json",
"main": "main.js",
"saved_files": []
}
+83
View File
@@ -0,0 +1,83 @@
[
{
"name": "timeoutEsc",
"type": "input-text",
"label": "步骤1: ESC循环超时(秒)",
"default": "60"
},
{
"name": "escInterval",
"type": "input-text",
"label": "步骤1: ESC按键间隔(毫秒)",
"default": "2000"
},
{
"name": "paimonBtnX",
"type": "input-text",
"label": "步骤3: 派蒙按钮 X(默认55)",
"default": "55"
},
{
"name": "paimonBtnY",
"type": "input-text",
"label": "步骤3: 派蒙按钮 Y(默认1010)",
"default": "1010"
},
{
"name": "timeoutMenu",
"type": "input-text",
"label": "步骤3: 派蒙菜单展开超时(秒)",
"default": "10"
},
{
"name": "timeoutExit",
"type": "input-text",
"label": "步骤4: 退出图标超时(秒)",
"default": "10"
},
{
"name": "timeoutGate",
"type": "input-text",
"label": "步骤6: 开门页面等待超时(秒)",
"default": "30"
},
{
"name": "useOcr",
"type": "checkbox",
"label": "使用OCR识别文本(否则用固定坐标)",
"default": true
},
{
"type": "separator"
},
{
"name": "confirmExitX",
"type": "input-text",
"label": "确认退出按钮 X(固定坐标模式)",
"default": "830"
},
{
"name": "confirmExitY",
"type": "input-text",
"label": "确认退出按钮 Y(固定坐标模式)",
"default": "600"
},
{
"name": "qrLoginX",
"type": "input-text",
"label": "扫码登录按钮 X(固定坐标模式)",
"default": "1720"
},
{
"name": "qrLoginY",
"type": "input-text",
"label": "扫码登录按钮 Y(固定坐标模式)",
"default": "970"
},
{
"name": "timeoutQr",
"type": "input-text",
"label": "扫码等待超时(秒)",
"default": "120"
}
]
+26
View File
@@ -0,0 +1,26 @@
# -*- coding: utf-8 -*-
# 默认依赖:可正常启动排队系统、Web 后台、弹幕监听、音乐监听等基础功能。
# TTS/GPU 依赖不要放在默认安装里,否则 Python 3.13 下 torch 没有可用 wheel 会导致整套项目安装失败。
websockets>=12.0,<13.0
brotli>=1.1.0
aiohttp>=3.9.0
pygame>=2.5.0
pywin32>=306;platform_system=="Windows"
winrt-Windows.Media.Control>=3.0
winrt-Windows.Foundation>=3.0
winrt-Windows.Storage.Streams>=3.0
uiautomation>=2.0.20;platform_system=="Windows"
pycaw>=20251023;platform_system=="Windows"
comtypes>=1.4.16;platform_system=="Windows"
soundfile>=0.12.0
sounddevice>=0.5.0
qrcode[pil]>=8.0,<9.0
pycryptodome>=3.20,<4.0
huggingface_hub>=0.20.0
pyinstaller>=6.11,<7.0
# 可选:启用 faster-qwen3-tts 时再安装以下依赖,建议使用 Python 3.10-3.12 + CUDA 对应 torch。
# faster-qwen3-tts>=0.1.0
# torch>=2.0.0
# torchaudio>=2.0.0
# transformers>=4.40.0
+20
View File
@@ -0,0 +1,20 @@
@echo off
title BetterGI Live Streaming Service
cd /d "%~dp0"
set PYTHON=%~dp0.venv-tts\Scripts\python.exe
if not exist "%PYTHON%" (
echo Python virtual environment not found: %PYTHON%
echo Please run setup first or ask WorkBuddy to repair dependencies.
pause
exit /b 1
)
echo Starting Live-streaming service
echo Admin page: http://127.0.0.1:5191/admin
echo Data page: http://127.0.0.1:5191/data
echo Press Ctrl+C to stop
echo.
"%PYTHON%" app\main.py --role all --host 0.0.0.0 --port 5191
pause
+14
View File
@@ -0,0 +1,14 @@
import winreg
SUBKEY = r"Software\Microsoft\Windows\CurrentVersion\Run"
NAME = "BetterGI弹幕排队系统"
VALUE = 'cmd.exe /d /k ""C:\\Users\\Administrator\\Desktop\\Live-streaming\\run.bat""'
key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, SUBKEY, 0, winreg.KEY_SET_VALUE)
winreg.SetValueEx(key, NAME, 0, winreg.REG_SZ, VALUE)
winreg.CloseKey(key)
key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, SUBKEY, 0, winreg.KEY_READ)
value, _ = winreg.QueryValueEx(key, NAME)
winreg.CloseKey(key)
print("WINREG WRITE+READ OK:", value)
+209
View File
@@ -0,0 +1,209 @@
[CmdletBinding()]
param(
[switch]$SkipInstall,
[switch]$FullTts,
[switch]$Lite,
[switch]$NoAdmin
)
$ErrorActionPreference = "Stop"
[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false)
$OutputEncoding = [System.Text.UTF8Encoding]::new($false)
$Root = Resolve-Path (Join-Path $PSScriptRoot "..")
$AdminRoot = Join-Path $Root "frontend\admin"
$DistRoot = Join-Path $Root "dist"
$AppDist = Join-Path $DistRoot "LiveStreaming"
$BuildRoot = Join-Path $Root "build"
$PyBuild = Join-Path $BuildRoot "pyinstaller"
function Write-Step([string]$Message) {
Write-Host ""
Write-Host "==> $Message" -ForegroundColor Cyan
}
function Invoke-Checked([string]$File, [string[]]$Arguments) {
& $File @Arguments
if ($LASTEXITCODE -ne 0) {
throw "$File $($Arguments -join ' ') failed with exit code $LASTEXITCODE."
}
}
function Get-ProjectPython {
$candidates = @(
$env:LIVE_STREAMING_PYTHON,
"E:\Programs\Anaconda3\envs\Live-streaming\python.exe",
"python"
) | Where-Object { $_ }
foreach ($candidate in $candidates) {
try {
$version = & $candidate -c "import sys; print(sys.version.split()[0])" 2>$null
if ($LASTEXITCODE -eq 0 -and $version) {
Write-Host "Python: $candidate ($version)"
return $candidate
}
} catch {
}
}
throw "No usable Python found. Install Python 3.10+, set LIVE_STREAMING_PYTHON, or use the Live-streaming conda env."
}
$Python = Get-ProjectPython
$PythonPrefix = (& $Python -c "import sys; print(sys.prefix)").Trim()
if ($LASTEXITCODE -ne 0 -or -not $PythonPrefix) {
throw "Unable to resolve Python prefix from $Python."
}
$CondaBin = Join-Path $PythonPrefix "Library\bin"
if (Test-Path $CondaBin) {
$env:PATH = "$CondaBin;$env:PATH"
}
if (-not $SkipInstall) {
Write-Step "Install Python runtime/build dependencies"
Invoke-Checked $Python @("-m", "pip", "install", "-r", (Join-Path $Root "requirements.txt"))
}
Write-Step "Build Vue admin dist"
if (-not (Test-Path (Join-Path $AdminRoot "node_modules"))) {
Invoke-Checked "npm" @("--prefix", $AdminRoot, "install")
}
Invoke-Checked "npm" @("--prefix", $AdminRoot, "run", "build")
Write-Step "Prepare PyInstaller output"
# 确保没有残留进程占用 dist 目录
$killed = Get-Process -Name "LiveStreaming" -ErrorAction SilentlyContinue
if ($killed) {
$killed | Stop-Process -Force -ErrorAction SilentlyContinue
$killed | Wait-Process -Timeout 5 -ErrorAction SilentlyContinue
Start-Sleep -Seconds 2
Write-Host " Stopped LiveStreaming.exe before build"
}
function Remove-WithRetry([string]$Path) {
if (-not (Test-Path $Path)) { return }
$retries = 3
while ($retries -gt 0) {
try {
Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
return
} catch {
$retries--
if ($retries -eq 0) {
Write-Host " WARNING: Cannot remove $Path, skipping" -ForegroundColor Yellow
} else {
Write-Host " Retry remove $Path ($retries left)..." -ForegroundColor Yellow
Start-Sleep -Seconds 2
}
}
}
}
Remove-WithRetry $AppDist
Remove-WithRetry $PyBuild
New-Item -ItemType Directory -Force -Path $DistRoot | Out-Null
New-Item -ItemType Directory -Force -Path $PyBuild | Out-Null
$excludeModules = @(
"tkinter",
"pytest",
"IPython",
"jupyter",
"notebook",
"matplotlib.tests",
"numpy.tests",
"pandas.tests",
"scipy.tests",
"torch.testing"
)
if ($Lite) {
$excludeModules += @(
"torch",
"torchaudio",
"torchvision",
"transformers",
"faster_qwen3_tts",
"dots_tts"
)
} else {
$excludeModules += @(
"torch.distributed",
"torch.utils.tensorboard",
"torchvision",
"torchaudio._internal",
"torch.utils.benchmark",
"torch.testing",
"torch.onnx",
"tensorboard",
"tensorboardX"
)
}
$pyiArgs = @(
"--noconfirm",
"--onedir",
"--name", "LiveStreaming",
"--distpath", $DistRoot,
"--workpath", $PyBuild,
"--specpath", $BuildRoot,
"--paths", (Join-Path $Root "app"),
"--hidden-import", "danmu_queue",
"--hidden-import", "music_monitor",
"--hidden-import", "tts_monitor",
"--hidden-import", "core.runtime_paths"
)
foreach ($module in $excludeModules) {
$pyiArgs += @("--exclude-module", $module)
}
$condaDlls = @("ffi.dll", "liblzma.dll", "libbz2.dll", "libexpat.dll")
if (Test-Path $CondaBin) {
foreach ($dll in $condaDlls) {
$dllPath = Join-Path $CondaBin $dll
if (Test-Path $dllPath) {
$pyiArgs += @("--add-binary", "$dllPath;.")
}
}
}
if (-not $NoAdmin) {
$pyiArgs += "--uac-admin"
}
$pyiArgs += (Join-Path $Root "app\main.py")
Write-Step "Package LiveStreaming.exe"
# PyInstaller 隔离模式与 torch DLL 冲突,设置环境变量跳过 CUDA 加载
$env:CUDA_VISIBLE_DEVICES = ""
$env:PYTORCH_NVFUSER_DISABLE = "1"
$SpecFile = Join-Path $BuildRoot "LiveStreaming.spec"
if (Test-Path $SpecFile) {
Write-Host " Reusing spec (incremental) ..."
Invoke-Checked $Python @("-m", "PyInstaller", "--noconfirm", $SpecFile)
} else {
Invoke-Checked $Python (@("-m", "PyInstaller") + $pyiArgs)
}
Write-Step "Create runtime directories"
$ExePath = Join-Path $AppDist "LiveStreaming.exe"
if (-not (Test-Path $ExePath)) {
throw "PyInstaller finished without creating $ExePath."
}
New-Item -ItemType Directory -Force -Path (Join-Path $AppDist "data") | Out-Null
New-Item -ItemType Directory -Force -Path (Join-Path $AppDist "logs") | Out-Null
Copy-Item -LiteralPath (Join-Path $Root "web") -Destination (Join-Path $AppDist "web") -Recurse -Force
Copy-Item -LiteralPath (Join-Path $Root "config") -Destination (Join-Path $AppDist "config") -Recurse -Force
Copy-Item -LiteralPath (Join-Path $Root "integrations") -Destination (Join-Path $AppDist "integrations") -Recurse -Force
Copy-Item -LiteralPath (Join-Path $Root "vendor") -Destination (Join-Path $AppDist "vendor") -Recurse -Force
Copy-Item -LiteralPath (Join-Path $Root "README.md") -Destination (Join-Path $AppDist "README.md") -Force
Write-Host ""
Write-Host "Build complete: $AppDist\LiveStreaming.exe" -ForegroundColor Green
if ($Lite) {
Write-Host "Lite build: TTS/GPU dependencies excluded."
} else {
Write-Host "Full build with TTS support."
}
+275
View File
@@ -0,0 +1,275 @@
# -*- coding: utf-8 -*-
"""
补充处理之前跳过的项目(有子目录的多作者/多版本食材与特产)。
规则:
- 子目录名含"效率版/高效率/中效率/效率路线/效率" → 用该效率版子目录的 json 创建配置组
- 其余 → 创建空配置组(projects=[]),用户手动添加路线 json
用法:
python fill_skipped_groups.py
"""
import os
import json
import sys
BASE = r"C:\Program Files\BetterGI\BetterGI\User"
AUTOPATHING = os.path.join(BASE, "AutoPathing")
SCRIPTGROUP = os.path.join(BASE, "ScriptGroup")
IGNORE_NAMES = {"desktop.ini", "icon.ico", "Thumbs.db"}
REGION_ORDER = ["蒙德", "璃月", "稻妻", "须弥", "枫丹", "纳塔", "挪德卡莱"]
CATEGORIES = ["地方特产", "食材与炼金"]
def region_sort_key(name):
try:
return (REGION_ORDER.index(name), name)
except ValueError:
return (len(REGION_ORDER), name)
def load_config_template():
template_path = os.path.join(SCRIPTGROUP, "子探测单元.json")
with open(template_path, "r", encoding="utf-8") as f:
return json.load(f)["config"]
def get_next_group_index():
max_index = 0
for fn in os.listdir(SCRIPTGROUP):
if not fn.lower().endswith(".json"):
continue
fp = os.path.join(SCRIPTGROUP, fn)
try:
with open(fp, "r", encoding="utf-8") as f:
data = json.load(f)
idx = data.get("index", 0)
if isinstance(idx, int) and idx > max_index:
max_index = idx
except Exception:
continue
return max_index + 1
def detect_structure(category_path):
has_any_direct_json = False
for entry in os.listdir(category_path):
entry_path = os.path.join(category_path, entry)
if not os.path.isdir(entry_path):
continue
sub_entries = [e for e in os.listdir(entry_path) if e not in IGNORE_NAMES]
has_json = any(
fn.lower().endswith(".json") and os.path.isfile(os.path.join(entry_path, fn))
for fn in sub_entries
)
if has_json:
has_any_direct_json = True
break
return "two" if has_any_direct_json else "three"
def list_items(category_path, structure):
result = []
if structure == "two":
for item in os.listdir(category_path):
ip = os.path.join(category_path, item)
if not os.path.isdir(ip):
continue
result.append((item, item, ip))
result.sort(key=lambda x: x[0])
else:
for region in os.listdir(category_path):
rp = os.path.join(category_path, region)
if not os.path.isdir(rp):
continue
for item in os.listdir(rp):
ip = os.path.join(rp, item)
if not os.path.isdir(ip):
continue
result.append(((region, item), item, ip))
result.sort(key=lambda x: (region_sort_key(x[0][0]), x[0][1]))
return result
def analyze_item(item_path):
entries = os.listdir(item_path)
subdirs = [e for e in entries
if os.path.isdir(os.path.join(item_path, e)) and e not in IGNORE_NAMES]
json_files = [e for e in entries
if e.lower().endswith(".json")
and os.path.isfile(os.path.join(item_path, e))
and e not in IGNORE_NAMES]
return subdirs, sorted(json_files)
def list_jsons_in_dir(dir_path):
"""列出某目录下直接包含的 json 文件(排除系统文件)。"""
if not os.path.isdir(dir_path):
return []
result = []
for e in os.listdir(dir_path):
if e in IGNORE_NAMES:
continue
if e.lower().endswith(".json") and os.path.isfile(os.path.join(dir_path, e)):
result.append(e)
return sorted(result)
def pick_efficiency_subdir(subdirs):
"""按优先级选取效率版子目录,返回子目录名或 None。"""
priorities = ["效率版", "高效率", "中效率", "效率路线", "效率"]
for keyword in priorities:
for sd in subdirs:
if keyword in sd:
return sd
return None
def build_group(group_index, name, folder_name, json_files, config_template):
projects = []
for i, jf in enumerate(json_files, 1):
projects.append({
"name": jf,
"folderName": folder_name,
"jsScriptSettingsObject": None,
"index": i,
"type": "Pathing",
"status": "Enabled",
"schedule": "Daily",
"runNum": 1,
"allowJsNotification": True,
"allowJsHTTPHash": ""
})
return {
"index": group_index,
"name": name,
"config": config_template,
"projects": projects
}
def main():
config_template = load_config_template()
report_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"补充报告_跳过项目.txt"
)
eff_created = [] # (category, group_key, name, eff_subdir, route_count, file_path)
empty_created = [] # (category, group_key, name, subdirs, file_path)
write_errors = []
group_index = get_next_group_index()
print(f"[信息] 起始 group index: {group_index}")
for category in CATEGORIES:
category_path = os.path.join(AUTOPATHING, category)
if not os.path.isdir(category_path):
continue
structure = detect_structure(category_path)
items = list_items(category_path, structure)
for group_key, name, item_path in items:
subdirs, json_files = analyze_item(item_path)
if not subdirs:
continue # 无子目录,之前已创建,跳过
# 构造 folderName
if structure == "two":
folder_name = f"{category}\\{name}"
else:
region, _ = group_key
folder_name = f"{category}\\{region}\\{name}"
# 检查是否已存在同名配置组(避免覆盖之前创建的)
out_path = os.path.join(SCRIPTGROUP, f"{name}.json")
if os.path.exists(out_path):
# 已存在,跳过
continue
eff_subdir = pick_efficiency_subdir(subdirs)
if eff_subdir:
# 用效率版子目录的 json
eff_path = os.path.join(item_path, eff_subdir)
# 效率版子目录可能直接含 json,也可能再含子目录(一般直接含)
eff_jsons = list_jsons_in_dir(eff_path)
if not eff_jsons:
# 效率版子目录下无直接 json,创建空配置组
group = build_group(group_index, name, folder_name, [], config_template)
try:
with open(out_path, "w", encoding="utf-8") as f:
json.dump(group, f, ensure_ascii=False, indent=2)
empty_created.append((category, group_key, name, subdirs, out_path))
group_index += 1
except Exception as e:
write_errors.append((name, str(e)))
continue
group = build_group(group_index, name, folder_name, eff_jsons, config_template)
try:
with open(out_path, "w", encoding="utf-8") as f:
json.dump(group, f, ensure_ascii=False, indent=2)
eff_created.append((category, group_key, name, eff_subdir, len(eff_jsons), out_path))
group_index += 1
except Exception as e:
write_errors.append((name, str(e)))
else:
# 无效率版 → 创建空配置组
group = build_group(group_index, name, folder_name, [], config_template)
try:
with open(out_path, "w", encoding="utf-8") as f:
json.dump(group, f, ensure_ascii=False, indent=2)
empty_created.append((category, group_key, name, subdirs, out_path))
group_index += 1
except Exception as e:
write_errors.append((name, str(e)))
# 生成报告
lines = []
lines.append("=" * 70)
lines.append("BetterGI 跳过项目补充处理报告")
lines.append("=" * 70)
lines.append("")
lines.append(f"【使用效率版创建】共 {len(eff_created)}")
lines.append("-" * 70)
cur_cat = None
for category, gk, name, eff_sd, cnt, fp in eff_created:
if category != cur_cat:
cur_cat = category
lines.append(f"\n[{category}]")
lines.append(f" {name}")
lines.append(f" 效率版: {eff_sd} ({cnt} 条路线)")
lines.append("")
lines.append(f"【创建空配置组(需手动添加路线)】共 {len(empty_created)}")
lines.append("-" * 70)
cur_cat = None
for category, gk, name, subdirs, fp in empty_created:
if category != cur_cat:
cur_cat = category
lines.append(f"\n[{category}]")
lines.append(f" {name}")
lines.append(f" 候选子目录: {subdirs}")
lines.append("")
if write_errors:
lines.append("【写入错误】")
lines.append("-" * 70)
for name, err in write_errors:
lines.append(f" {name}: {err}")
lines.append("")
lines.append("=" * 70)
report = "\n".join(lines)
with open(report_path, "w", encoding="utf-8") as f:
f.write(report)
print(report)
print(f"\n报告已保存: {report_path}")
if __name__ == "__main__":
main()
+313
View File
@@ -0,0 +1,313 @@
# -*- coding: utf-8 -*-
"""
为 BetterGI AutoPathing 分类目录批量生成 ScriptGroup 配置组。
支持两类结构:
- 两层(分类/食材):如"食材与炼金/<食材>/xxx.json" folderName = "<分类>\\<食材>"
- 三层(分类/地区/特产):如"地方特产/<地区>/<特产>/xxx.json" folderName = "<分类>\\<地区>\\<特产>"
规则:
- 食材/特产目录下若无子目录(json 直接平铺)→ 单一来源路线 → 生成配置组
- 食材/特产目录下若有子目录(多作者/多版本路线)→ 跳过,记录到未加入清单
用法:
python generate_scriptgroups.py <分类名>
例: python generate_scriptgroups.py 食材与炼金
python generate_scriptgroups.py 地方特产
输出:
- 在 ScriptGroup 目录下为每个可创建的项目生成 <名称>.json
- 在项目目录下生成 生成报告_<分类名>.txt
"""
import os
import json
import sys
BASE = r"C:\Program Files\BetterGI\BetterGI\User"
AUTOPATHING = os.path.join(BASE, "AutoPathing")
SCRIPTGROUP = os.path.join(BASE, "ScriptGroup")
# 非路线文件,扫描时排除
IGNORE_NAMES = {"desktop.ini", "icon.ico", "Thumbs.db"}
# 地区显示顺序(仅用于排序,地方特产用)
REGION_ORDER = ["蒙德", "璃月", "稻妻", "须弥", "枫丹", "纳塔", "挪德卡莱"]
def region_sort_key(name):
try:
return (REGION_ORDER.index(name), name)
except ValueError:
return (len(REGION_ORDER), name)
def load_config_template():
"""读取 子探测单元.json 作为 config 模板。"""
template_path = os.path.join(SCRIPTGROUP, "子探测单元.json")
if not os.path.exists(template_path):
print(f"[错误] 找不到模板文件: {template_path}")
sys.exit(1)
with open(template_path, "r", encoding="utf-8") as f:
return json.load(f)["config"]
def get_next_group_index():
"""扫描 ScriptGroup 下所有 json,返回下一个可用 index。"""
max_index = 0
if not os.path.isdir(SCRIPTGROUP):
return 1
for fn in os.listdir(SCRIPTGROUP):
if not fn.lower().endswith(".json"):
continue
fp = os.path.join(SCRIPTGROUP, fn)
try:
with open(fp, "r", encoding="utf-8") as f:
data = json.load(f)
idx = data.get("index", 0)
if isinstance(idx, int) and idx > max_index:
max_index = idx
except Exception:
continue
return max_index + 1
def detect_structure(category_path):
"""检测目录结构:返回 'two' (分类/食材) 或 'three' (分类/地区/特产)。
判定规则:扫描所有第二层目录,只要有一个直接包含 json 文件,
就判为两层结构(单作者食材直接放 json)。只有当所有第二层目录
都只含子目录(无直接 json)时,才判为三层(地区→特产)。
"""
has_any_direct_json = False
for entry in os.listdir(category_path):
entry_path = os.path.join(category_path, entry)
if not os.path.isdir(entry_path):
continue
sub_entries = [e for e in os.listdir(entry_path) if e not in IGNORE_NAMES]
has_json = any(
fn.lower().endswith(".json") and os.path.isfile(os.path.join(entry_path, fn))
for fn in sub_entries
)
if has_json:
has_any_direct_json = True
break
return "two" if has_any_direct_json else "three"
def list_items(category_path, structure):
"""返回 [(group_key, item_name, item_path)] 列表。
- two: group_key=item_name, item_path=分类/食材
- three: group_key=(region, item_name), item_path=分类/地区/特产
"""
result = []
if structure == "two":
for item in os.listdir(category_path):
ip = os.path.join(category_path, item)
if not os.path.isdir(ip):
continue
result.append((item, item, ip))
result.sort(key=lambda x: x[0])
else:
for region in os.listdir(category_path):
rp = os.path.join(category_path, region)
if not os.path.isdir(rp):
continue
for item in os.listdir(rp):
ip = os.path.join(rp, item)
if not os.path.isdir(ip):
continue
result.append(((region, item), item, ip))
result.sort(key=lambda x: (region_sort_key(x[0][0]), x[0][1]))
return result
def analyze_item(item_path):
"""分析项目目录,返回 (subdirs, json_files)。"""
entries = os.listdir(item_path)
subdirs = [e for e in entries
if os.path.isdir(os.path.join(item_path, e)) and e not in IGNORE_NAMES]
json_files = [e for e in entries
if e.lower().endswith(".json")
and os.path.isfile(os.path.join(item_path, e))
and e not in IGNORE_NAMES]
return subdirs, sorted(json_files)
def classify_subdirs(subdirs):
"""对有子目录的项目分类,返回 (类型, 说明)。"""
authors = []
for sd in subdirs:
if "@" in sd:
author = sd.split("@", 1)[1]
authors.append(author)
else:
authors.append(None)
unique_authors = {a for a in authors if a is not None}
has_unmarked = any(a is None for a in authors)
if len(unique_authors) == 0:
return "版本说明类", "所有子目录均无作者标识(版本/路线说明)"
if has_unmarked:
return "混合多路线", f"含无作者标识的版本目录;作者: {sorted(unique_authors)}"
if len(unique_authors) == 1:
return "同一作者多版本", f"作者: {next(iter(unique_authors))}"
return "多作者", f"作者: {sorted(unique_authors)}"
def build_group(group_index, name, folder_name, json_files, config_template):
"""构造一个配置组 dict。"""
projects = []
for i, jf in enumerate(json_files, 1):
projects.append({
"name": jf,
"folderName": folder_name,
"jsScriptSettingsObject": None,
"index": i,
"type": "Pathing",
"status": "Enabled",
"schedule": "Daily",
"runNum": 1,
"allowJsNotification": True,
"allowJsHTTPHash": ""
})
return {
"index": group_index,
"name": name,
"config": config_template,
"projects": projects
}
def main():
if len(sys.argv) < 2:
print("用法: python generate_scriptgroups.py <分类名>")
print("例: python generate_scriptgroups.py 食材与炼金")
sys.exit(1)
category = sys.argv[1]
category_path = os.path.join(AUTOPATHING, category)
if not os.path.isdir(category_path):
print(f"[错误] 分类目录不存在: {category_path}")
sys.exit(1)
config_template = load_config_template()
report_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
f"生成报告_{category}.txt"
)
structure = detect_structure(category_path)
print(f"[信息] 分类 '{category}' 结构: {'两层(分类/食材)' if structure == 'two' else '三层(分类/地区/特产)'}")
items = list_items(category_path, structure)
print(f"[信息] 共扫描到 {len(items)} 个项目")
created = [] # (group_key, name, route_count, file_path)
skipped = [] # (group_key, name, subdirs, 类型, 说明)
write_errors = [] # (name, error)
group_index = get_next_group_index()
print(f"[信息] 起始 group index: {group_index}")
for group_key, name, item_path in items:
subdirs, json_files = analyze_item(item_path)
if not subdirs:
# 无子目录 → 创建配置组
if not json_files:
skipped.append((group_key, name, [], "空目录", "目录下无路线文件"))
continue
# 构造 folderName
if structure == "two":
folder_name = f"{category}\\{name}"
else:
region, _ = group_key
folder_name = f"{category}\\{region}\\{name}"
group = build_group(group_index, name, folder_name, json_files, config_template)
out_path = os.path.join(SCRIPTGROUP, f"{name}.json")
try:
with open(out_path, "w", encoding="utf-8") as f:
json.dump(group, f, ensure_ascii=False, indent=2)
created.append((group_key, name, len(json_files), out_path))
group_index += 1
except PermissionError as e:
write_errors.append((name, f"权限不足: {e}"))
except Exception as e:
write_errors.append((name, str(e)))
else:
# 有子目录 → 跳过
kind, desc = classify_subdirs(subdirs)
skipped.append((group_key, name, subdirs, kind, desc))
# 生成报告
lines = []
lines.append("=" * 70)
lines.append(f"BetterGI {category} 配置组生成报告")
lines.append("=" * 70)
lines.append("")
lines.append(f"【已创建配置组】共 {len(created)}")
lines.append("-" * 70)
def group_label(gk):
if structure == "two":
return gk
else:
return f"{gk[0]} / {gk[1]}"
# 按分组打印已创建
if structure == "two":
for gk, name, cnt, fp in created:
lines.append(f" {name} ({cnt} 条路线) -> {os.path.basename(fp)}")
else:
cur_region = None
for gk, name, cnt, fp in created:
region = gk[0]
if region != cur_region:
cur_region = region
lines.append(f"\n[{region}]")
lines.append(f" {name} ({cnt} 条路线) -> {os.path.basename(fp)}")
lines.append("")
lines.append(f"【未加入配置组(需手动添加)】共 {len(skipped)}")
lines.append("-" * 70)
if structure == "two":
for gk, name, subdirs, kind, desc in skipped:
lines.append(f" {name}")
lines.append(f" 类型: {kind}")
lines.append(f" 说明: {desc}")
if subdirs:
lines.append(f" 子目录: {subdirs}")
lines.append("")
else:
cur_region = None
for gk, name, subdirs, kind, desc in skipped:
region = gk[0]
if region != cur_region:
cur_region = region
lines.append(f"\n[{region}]")
lines.append(f" {name}")
lines.append(f" 类型: {kind}")
lines.append(f" 说明: {desc}")
if subdirs:
lines.append(f" 子目录: {subdirs}")
lines.append("")
if write_errors:
lines.append("【写入错误】")
lines.append("-" * 70)
for name, err in write_errors:
lines.append(f" {name}: {err}")
lines.append("")
lines.append("=" * 70)
report = "\n".join(lines)
with open(report_path, "w", encoding="utf-8") as f:
f.write(report)
print(report)
print(f"\n报告已保存: {report_path}")
if __name__ == "__main__":
main()
+127
View File
@@ -0,0 +1,127 @@
"""Preview the built admin UI against a fake local backend.
This script does not connect to Bilibili or BetterGI. It is only for checking
the admin page in a local browser while developing the Vue UI.
"""
from __future__ import annotations
import asyncio
import logging
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "app"))
from danmu_queue import Config, ServiceRegistry, WebServer # noqa: E402
class FakeRunner:
def is_bgi_running(self):
return False
async def kill_bgi(self):
return None
class FakeQueueManager:
def __init__(self):
self.state = {
"queue": [10001, 10002],
"current_admin_uid": 10001,
"current_group": "薄荷",
"group_start_time": None,
"admin_window_end": None,
"default_running": False,
"login_status": "confirming",
}
def _save(self):
return None
def leave_queue(self, uid):
if uid in self.state["queue"]:
self.state["queue"].remove(uid)
return {"success": True, "was_running": False}
class FakeUserManager:
def __init__(self):
self.users = {
"10001": {"uname": "测试用户A", "points": 18, "last_signin_date": ""},
"10002": {"uname": "测试用户B", "points": 7, "last_signin_date": ""},
}
async def _save(self):
return None
async def add_points(self, uid, points):
user = self.users.setdefault(str(uid), {"uname": f"用户{uid}", "points": 0})
user["points"] += int(points)
return user["points"]
class FakeSongRequestManager:
def __init__(self):
self.state = {
"queue": [
{"id": "1", "name": "测试歌曲", "artist": "测试歌手", "uname": "测试用户A"}
]
}
def remove_request(self, index=None, song_id=""):
queue = self.state.setdefault("queue", [])
if queue:
return {"success": True, "removed": queue.pop(0)}
return {"success": False, "msg": "empty"}
def clear_requests(self):
count = len(self.state.get("queue", []))
self.state["queue"] = []
return count
class FakeLogMonitor:
def set_current_group(self, group):
return None
class FakeSystem:
def __init__(self):
self.runner = FakeRunner()
self.queue_mgr = FakeQueueManager()
self.user_mgr = FakeUserManager()
self.song_request_mgr = FakeSongRequestManager()
self.log_monitor = FakeLogMonitor()
self.handler = type("FakeHandler", (), {"recent_danmu": [
{"uname": "测试用户A", "text": "排队"},
{"uname": "测试用户B", "text": "点歌 测试歌曲"},
]})()
self.broadcaster = None
self.health = ServiceRegistry()
self.health.set("主程序", ServiceRegistry.RUNNING, "preview")
self.health.set("Web后台服务", ServiceRegistry.RUNNING, "preview")
self.health.set("直播监听", ServiceRegistry.RECONNECTING, "preview reconnect")
def apply_config(self):
self.health.set("配置", ServiceRegistry.RUNNING, "preview save")
async def _start_default_group(self):
return None
async def main():
cfg = Config(str(ROOT / "config" / "config_queue.json"))
logger = logging.getLogger("preview_admin")
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler(sys.stdout))
server = WebServer(FakeSystem(), cfg, logger, port=5190, host="127.0.0.1")
server_task = asyncio.create_task(server.start())
await asyncio.sleep(0.2)
print("Preview: http://127.0.0.1:5190/admin")
await server_task
if __name__ == "__main__":
asyncio.run(main())
+73
View File
@@ -0,0 +1,73 @@
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from types import SimpleNamespace
from unittest.mock import patch
from app.danmu_queue import BgiLogMonitor, WebServer
class BgiLogDisplayFormatTests(unittest.TestCase):
def test_formats_multi_instance_log_as_single_readable_line(self):
raw_lines = [
"[23:55:25.160] [INF] [Primary:S1:P16720:T1785426828552] "
"BetterGenshinImpact.GameTask.TaskTriggerDispatcher\n",
"游戏已退出,BetterGI 自动停止截图器\n",
"\n",
]
self.assertEqual(
BgiLogMonitor.format_display_lines(raw_lines),
["[23:55:25] 游戏已退出,BetterGI 自动停止截图器"],
)
def test_formats_legacy_log_header_the_same_way(self):
raw_lines = [
"[16:39:12.328] [INF] BetterGenshinImpact.GameTask.TaskTriggerDispatcher\n",
"游戏已退出,BetterGI 自动停止截图器\n",
]
self.assertEqual(
BgiLogMonitor.format_display_lines(raw_lines),
["[16:39:12] 游戏已退出,BetterGI 自动停止截图器"],
)
def test_preserves_multiline_message_inside_one_log_entry(self):
raw_lines = [
"[10:00:00.001] [ERR] [Primary:S1:P1:T2] BetterGI.Component\n",
"第一行\n",
"第二行\n",
"[10:00:01.999] [INF] BetterGI.OtherComponent\n",
"下一条\n",
]
self.assertEqual(
BgiLogMonitor.format_display_lines(raw_lines),
["[10:00:00] 第一行\n第二行", "[10:00:01] 下一条"],
)
def test_frontend_reader_falls_back_to_latest_log_after_midnight(self):
with TemporaryDirectory() as temp_dir:
log_dir = Path(temp_dir) / "log"
log_dir.mkdir()
(log_dir / "better-genshin-impact20260730.log").write_text(
"[23:55:25.160] [INF] [Primary:S1:P16720:T1785426828552] "
"BetterGenshinImpact.GameTask.TaskTriggerDispatcher\n"
"游戏已退出,BetterGI 自动停止截图器\n",
encoding="utf-8",
)
server = WebServer.__new__(WebServer)
server.config = SimpleNamespace(bettergi_work_dir=temp_dir)
with patch("app.danmu_queue.datetime") as mocked_datetime:
mocked_datetime.now.return_value.strftime.return_value = "20260731"
result = server._read_bgi_log(50)
self.assertEqual(
result,
["[23:55:25] 游戏已退出,BetterGI 自动停止截图器"],
)
if __name__ == "__main__":
unittest.main()
+156
View File
@@ -0,0 +1,156 @@
import logging
import unittest
import urllib.error
from unittest.mock import AsyncMock, Mock, patch
from app.bilibili_cookie_refresh import BilibiliQrLogin, _request_json
class _JsonResponse:
def __init__(self, payload=b'{"code": 0}'):
self.payload = payload
def read(self):
return self.payload
class BilibiliRequestTests(unittest.TestCase):
def test_get_retries_transient_connection_failures(self):
opener = Mock()
response = _JsonResponse()
opener.open.side_effect = [
urllib.error.URLError(ConnectionRefusedError(10061, "connection refused")),
urllib.error.URLError(ConnectionRefusedError(10061, "connection refused")),
response,
]
with patch("app.bilibili_cookie_refresh.time.sleep") as sleep:
payload, actual_response = _request_json(
"https://passport.bilibili.com/test",
opener=opener,
retries=2,
)
self.assertEqual(payload, {"code": 0})
self.assertIs(actual_response, response)
self.assertEqual(opener.open.call_count, 3)
self.assertEqual([call.args[0] for call in sleep.call_args_list], [0.4, 0.8])
def test_exhausted_retries_hide_low_level_network_error(self):
opener = Mock()
opener.open.side_effect = urllib.error.URLError(
ConnectionRefusedError(10061, "connection refused")
)
with patch("app.bilibili_cookie_refresh.time.sleep"), self.assertRaisesRegex(
RuntimeError,
"^连接B站登录服务失败,请稍后重试$",
) as raised:
_request_json(
"https://passport.bilibili.com/test",
opener=opener,
retries=2,
)
self.assertEqual(opener.open.call_count, 3)
self.assertNotIn("10061", str(raised.exception))
class _CredentialStore:
def __init__(self):
self.refresh_token = ""
def save_refresh_token(self, refresh_token):
self.refresh_token = refresh_token
def is_configured(self):
return bool(self.refresh_token)
class BilibiliQrLoginTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.store = _CredentialStore()
self.updated_cookies = []
self.on_logged_in = AsyncMock()
self.login = BilibiliQrLogin(
credential_store=self.store,
update_cookie=self.updated_cookies.append,
logger=logging.getLogger("test-bilibili-qr"),
on_logged_in=self.on_logged_in,
)
@staticmethod
def _generate_response():
return ({
"code": 0,
"data": {
"url": "https://passport.bilibili.com/h5-app/passport/login/scan?navhide=1",
"qrcode_key": "private-qrcode-key",
},
}, None)
async def test_start_keeps_qrcode_key_on_server(self):
with patch("app.bilibili_cookie_refresh._request_json", return_value=self._generate_response()):
result = await self.login.start()
self.assertEqual(result["state"], "awaiting_scan")
self.assertTrue(result["has_qr_image"])
self.assertNotIn("qr_key", result)
self.assertNotIn("qr_url", result)
self.assertFalse(result["credential_configured"])
async def test_poll_reports_scanned_without_saving_credentials(self):
with patch("app.bilibili_cookie_refresh._request_json", return_value=self._generate_response()):
await self.login.start()
with patch("app.bilibili_cookie_refresh._request_json", return_value=({
"code": 0,
"data": {"code": 86090, "message": "二维码已扫码未确认"},
}, None)):
result = await self.login.poll()
self.assertEqual(result["state"], "awaiting_confirm")
self.assertEqual(self.store.refresh_token, "")
self.assertEqual(self.updated_cookies, [])
self.on_logged_in.assert_not_awaited()
async def test_success_saves_refresh_token_and_updates_cookie(self):
with patch("app.bilibili_cookie_refresh._request_json", return_value=self._generate_response()):
await self.login.start()
login_url = (
"https://www.bilibili.com/?SESSDATA=session-value&bili_jct=csrf-value"
"&DedeUserID=123456"
)
def request_side_effect(url, **_kwargs):
if "qrcode/poll" in url:
return ({
"code": 0,
"data": {
"code": 0,
"url": login_url,
"refresh_token": "refresh-token-value",
},
}, None)
if "x/web-interface/nav" in url:
return ({
"code": 0,
"data": {"isLogin": True, "mid": 123456, "uname": "测试账号"},
}, None)
raise AssertionError(f"unexpected url: {url}")
with patch("app.bilibili_cookie_refresh._request_json", side_effect=request_side_effect):
result = await self.login.poll()
self.assertEqual(result["state"], "completed")
self.assertEqual(result["account"]["uname"], "测试账号")
self.assertTrue(result["credential_configured"])
self.assertNotIn("refresh_token", result)
self.assertEqual(self.store.refresh_token, "refresh-token-value")
self.assertEqual(self.updated_cookies[0]["SESSDATA"], "session-value")
self.assertEqual(self.updated_cookies[0]["bili_jct"], "csrf-value")
self.on_logged_in.assert_awaited_once()
if __name__ == "__main__":
unittest.main()
+256
View File
@@ -0,0 +1,256 @@
import logging
import os
import unittest
from unittest.mock import patch
from app.faster_qwen_worker import (
DEFAULT_CPU_AFFINITY_COUNT,
DEFAULT_CPU_THREADS,
MAX_NEW_TOKENS,
FasterQwenWorkerClient,
FasterQwenWorkerTimeout,
_generation_kwargs,
_configure_worker_environment,
)
class _FakeConnection:
def __init__(self, responses):
self.responses = list(responses)
self.sent = []
self.closed = False
def poll(self, _timeout):
return bool(self.responses)
def recv(self):
return self.responses.pop(0)
def send(self, message):
self.sent.append(message)
def close(self):
self.closed = True
class _FakeChildConnection:
def close(self):
return None
class _FakeProcess:
_next_pid = 4000
def __init__(self):
type(self)._next_pid += 1
self.pid = type(self)._next_pid
self.alive = False
self.terminated = False
self.killed = False
def start(self):
self.alive = True
def is_alive(self):
return self.alive
def join(self, timeout=None):
return None
def terminate(self):
self.terminated = True
self.alive = False
def kill(self):
self.killed = True
self.alive = False
class _FakeContext:
def __init__(self, connections):
self.connections = list(connections)
self.processes = []
def Pipe(self, duplex=True):
self.assert_duplex = duplex
return self.connections.pop(0), _FakeChildConnection()
def Process(self, **_kwargs):
process = _FakeProcess()
self.processes.append(process)
return process
def _ready(pid):
return {
"type": "ready",
"pid": pid,
"load_ms": 100,
"warmup_ms": 50,
"max_new_tokens": MAX_NEW_TOKENS,
}
def _probe_worker_main(connection, _settings):
connection.send({
"type": "ready",
"pid": os.getpid(),
"load_ms": 1,
"warmup_ms": 1,
"max_new_tokens": MAX_NEW_TOKENS,
})
while True:
message = connection.recv()
if message.get("command") == "stop":
break
if message.get("command") == "synthesize":
audio = b"RIFF-spawn-probe"
connection.send({
"type": "result",
"request_id": message["request_id"],
"audio": audio,
"duration_ms": 2,
"bytes": len(audio),
})
connection.close()
class FasterQwenWorkerTests(unittest.TestCase):
def test_worker_environment_limits_native_thread_pools(self):
names = (
"OMP_NUM_THREADS",
"MKL_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"NUMEXPR_NUM_THREADS",
"VECLIB_MAXIMUM_THREADS",
"BLIS_NUM_THREADS",
)
with patch.dict(os.environ, {}, clear=False):
threads = _configure_worker_environment({})
self.assertEqual(threads, DEFAULT_CPU_THREADS)
for name in names:
self.assertEqual(os.environ[name], str(DEFAULT_CPU_THREADS))
self.assertEqual(os.environ["TOKENIZERS_PARALLELISM"], "false")
def test_default_worker_resource_limits_are_four_threads_and_eight_cores(self):
self.assertEqual(DEFAULT_CPU_THREADS, 4)
self.assertEqual(DEFAULT_CPU_AFFINITY_COUNT, 8)
def test_worker_environment_accepts_bounded_override(self):
with patch.dict(os.environ, {}, clear=False):
self.assertEqual(_configure_worker_environment({"cpu_threads": 4}), 4)
self.assertEqual(_configure_worker_environment({"cpu_threads": 99}), 8)
def test_generation_is_hard_limited_to_384_tokens(self):
kwargs = _generation_kwargs(
{"language": "Chinese", "non_streaming_mode": True},
"测试",
)
self.assertEqual(MAX_NEW_TOKENS, 384)
self.assertEqual(kwargs["max_new_tokens"], 384)
def test_client_returns_worker_audio(self):
connection = _FakeConnection([
_ready(4100),
{
"type": "result",
"request_id": "placeholder",
"audio": b"RIFF-audio",
"duration_ms": 1234,
"bytes": 10,
},
])
context = _FakeContext([connection])
client = FasterQwenWorkerClient(
{},
logging.getLogger("test-faster-qwen-worker"),
context=context,
)
client.ensure_ready()
request_id = "fixed-request-id"
connection.responses[0]["request_id"] = request_id
with patch("app.faster_qwen_worker.uuid.uuid4") as make_uuid:
make_uuid.return_value.hex = request_id
audio, metadata = client.synthesize("测试")
self.assertEqual(audio, b"RIFF-audio")
self.assertEqual(metadata["duration_ms"], 1234)
self.assertEqual(len(context.processes), 1)
client.close()
def test_timeout_terminates_worker_and_starts_a_prewarmed_replacement(self):
first_connection = _FakeConnection([_ready(4200)])
replacement_connection = _FakeConnection([_ready(4300)])
context = _FakeContext([first_connection, replacement_connection])
client = FasterQwenWorkerClient(
{},
logging.getLogger("test-faster-qwen-timeout"),
synthesis_timeout_seconds=120,
context=context,
)
with self.assertRaises(FasterQwenWorkerTimeout):
client.synthesize("会超时的播报")
self.assertEqual(len(context.processes), 2)
self.assertTrue(context.processes[0].terminated)
self.assertTrue(context.processes[1].is_alive())
self.assertEqual(client.worker_pid, context.processes[1].pid)
client.close()
def test_real_spawned_worker_round_trip(self):
client = FasterQwenWorkerClient(
{},
logging.getLogger("test-faster-qwen-spawn"),
startup_timeout_seconds=20,
process_target=_probe_worker_main,
)
audio, metadata = client.synthesize("测试子进程")
self.assertEqual(audio, b"RIFF-spawn-probe")
self.assertEqual(metadata["duration_ms"], 2)
self.assertGreater(client.worker_pid, 0)
client.close()
def test_startup_failure_backoff_blocks_immediate_retry(self):
from app.faster_qwen_worker import (
STARTUP_FAILURE_BACKOFF_SECONDS,
FasterQwenWorkerError,
)
# 启动永远超时(poll 返回 False),触发启动失败
timeout_connection = _FakeConnection([])
timeout_connection.poll = lambda _timeout: False
context = _FakeContext([timeout_connection, _FakeConnection([_ready(4400)])])
client = FasterQwenWorkerClient(
{},
logging.getLogger("test-faster-qwen-backoff"),
startup_timeout_seconds=0.01,
context=context,
)
with self.assertRaises(FasterQwenWorkerError):
client.ensure_ready()
self.assertGreater(client._next_start_after, 0.0)
# 退避期内立即重试应直接报退避错误,且不创建新进程
with self.assertRaises(FasterQwenWorkerError) as ctx:
client.ensure_ready()
self.assertIn("退避", str(ctx.exception))
self.assertEqual(len(context.processes), 1)
# 退避期过后允许重新启动
client._next_start_after = 0.0
client.ensure_ready()
self.assertEqual(len(context.processes), 2)
self.assertEqual(client._next_start_after, 0.0)
self.assertGreater(STARTUP_FAILURE_BACKOFF_SECONDS, 0)
client.close()
if __name__ == "__main__":
unittest.main()
+77
View File
@@ -0,0 +1,77 @@
import sqlite3
import tempfile
import unittest
from pathlib import Path
from app.danmu_queue import bilibili_gift_cny_values
from app.stats_store import StatsStore
class BilibiliGiftValueTests(unittest.TestCase):
def test_gold_coin_conversion_uses_one_thousand_per_cny(self):
unit_value, total_value = bilibili_gift_cny_values("gold", 100, 1)
self.assertAlmostEqual(unit_value, 0.1)
self.assertAlmostEqual(total_value, 0.1)
def test_multi_quantity_conversion_preserves_total(self):
unit_value, total_value = bilibili_gift_cny_values("gold", 5000, 2)
self.assertAlmostEqual(unit_value, 2.5)
self.assertAlmostEqual(total_value, 5.0)
def test_silver_coin_has_no_cny_value(self):
self.assertEqual(bilibili_gift_cny_values("silver", 5000, 1), (0.0, 0.0))
def test_historical_repair_is_idempotent_and_rebuilds_aggregate(self):
with tempfile.TemporaryDirectory() as temp_dir:
path = Path(temp_dir) / "statistics.sqlite3"
connection = sqlite3.connect(path)
StatsStore._migrate(connection)
fields = {
"event_id": "gift-1",
"platform": "bilibili",
"platform_user_id": "123",
"occurred_at_utc": "2026-07-28T00:00:00Z",
"business_date": "2026-07-28",
"gift_id": "1",
"gift_name": "灯牌",
"quantity": 1,
"unit_value": 10.0,
"total_value": 10.0,
"currency": "CNY",
"payload_json": '{"coin_type":"gold","raw_total_coin":100,"value_rule":"gold_battery_10_to_cny_1_v1"}',
}
StatsStore._insert(connection, "gift_events", fields)
StatsStore._aggregate_gift(connection, fields)
legacy_fields = dict(fields)
legacy_fields.update({
"event_id": "gift-legacy",
"occurred_at_utc": "2026-07-28T00:01:00Z",
"payload_json": '{"coin_type":"gold","raw_total_coin":100,"value_migration":"bilibili_coin_to_cny_v1"}',
})
StatsStore._insert(connection, "gift_events", legacy_fields)
StatsStore._aggregate_gift(connection, legacy_fields)
connection.commit()
first = StatsStore._repair_gift_value_history(connection)
second = StatsStore._repair_gift_value_history(connection)
event = connection.execute(
"SELECT unit_value, total_value, json_extract(payload_json, '$.value_rule') "
"FROM gift_events WHERE event_id='gift-1'"
).fetchone()
aggregate = connection.execute(
"SELECT quantity, total_value FROM gift_aggregates"
).fetchone()
connection.close()
self.assertEqual(first, 2)
self.assertEqual(second, 0)
self.assertAlmostEqual(event[0], 0.1)
self.assertAlmostEqual(event[1], 0.1)
self.assertEqual(event[2], "bilibili_gold_coin_1000_to_cny_1_v2")
self.assertEqual(aggregate[0], 2)
self.assertAlmostEqual(aggregate[1], 0.2)
if __name__ == "__main__":
unittest.main()
+242
View File
@@ -0,0 +1,242 @@
import asyncio
import logging
import tempfile
import time
import unittest
from datetime import datetime, timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
from app.danmu_queue import CommandHandler, QueueManager, QueueSystem
class IdleDefaultQueueStateTests(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
self.manager = QueueManager(self.temp_dir.name, logging.getLogger("idle-default-test"))
self.manager.state["has_user_finished_once"] = True
def tearDown(self):
self.temp_dir.cleanup()
def _start_default_with_waiting_admin(self, uid=101):
self.manager.state.update({
"default_running": True,
"current_group": "薄荷",
"current_group_run_id": "default-run",
"group_start_time": datetime.now().isoformat(),
"billing_uid": None,
})
self.manager.join_queue(uid)
def test_default_completion_preserves_newly_joined_admin(self):
self._start_default_with_waiting_admin()
original_window = self.manager.state["admin_window_end"]
result = self.manager.group_finished("default-run")
self.assertTrue(result["accepted"])
self.assertTrue(result["was_default_running"])
self.assertTrue(result["need_default"])
self.assertEqual(self.manager.state["queue"], [101])
self.assertEqual(self.manager.state["current_admin_uid"], 101)
self.assertEqual(self.manager.state["admin_window_end"], original_window)
self.assertFalse(self.manager.state["default_running"])
self.assertIsNone(self.manager.state["current_group"])
def test_admin_window_can_expire_while_default_group_runs(self):
self._start_default_with_waiting_admin()
self.manager.state["admin_window_end"] = time.time() - 1
result = self.manager.check_admin_window_timeout()
self.assertTrue(result["timeout"])
self.assertEqual(result["kicked_uid"], 101)
self.assertEqual(self.manager.state["queue"], [])
self.assertTrue(self.manager.state["default_running"])
self.assertEqual(self.manager.state["current_group"], "薄荷")
def test_waiting_admin_leave_preserves_default_group(self):
self._start_default_with_waiting_admin()
result = self.manager.leave_queue(101)
self.assertTrue(result["success"])
self.assertTrue(self.manager.state["default_running"])
self.assertEqual(self.manager.state["current_group"], "薄荷")
self.assertEqual(self.manager.state["current_group_run_id"], "default-run")
class IdleDefaultFlowTests(unittest.IsolatedAsyncioTestCase):
@staticmethod
def _make_system(manager, runner):
system = QueueSystem.__new__(QueueSystem)
system.queue_mgr = manager
system.user_mgr = SimpleNamespace(
users={"101": {"uname": "tester"}},
get_points=lambda _uid: 10,
)
system.config = SimpleNamespace(default_group="薄荷")
system.runner = runner
system.log_monitor = SimpleNamespace(set_current_group=MagicMock())
system.login_monitor = SimpleNamespace(reset=MagicMock())
system.stats_store = None
system.logger = logging.getLogger("idle-default-race-test")
system.broadcast = AsyncMock()
system._default_start_lock = asyncio.Lock()
return system
@staticmethod
def _make_handler(system):
handler = CommandHandler.__new__(CommandHandler)
handler.config = SimpleNamespace(admin_uids=[])
handler.queue_mgr = system.queue_mgr
handler.user_mgr = system.user_mgr
handler.runner = system.runner
handler.log_monitor = system.log_monitor
handler.login_monitor = system.login_monitor
handler.logger = system.logger
handler.stats_store = None
handler.system = system
handler.broadcast = AsyncMock()
return handler
async def test_login_serializes_with_default_group_start(self):
with tempfile.TemporaryDirectory() as temp_dir:
manager = QueueManager(temp_dir, logging.getLogger("idle-default-login-race"))
manager.state["has_user_finished_once"] = True
manager.join_queue(101)
manager.state.update({
"default_running": True,
"current_group": "薄荷",
"current_group_run_id": "default-run",
"group_start_time": datetime.now().isoformat(),
})
manager._save()
kill_started = asyncio.Event()
allow_kill = asyncio.Event()
async def blocked_kill(**_kwargs):
kill_started.set()
await allow_kill.wait()
return True
runner = SimpleNamespace(
kill_bgi=AsyncMock(side_effect=blocked_kill),
start_groups=AsyncMock(return_value=True),
)
system = self._make_system(manager, runner)
handler = self._make_handler(system)
login_task = asyncio.create_task(handler._cmd_login(101, "tester"))
await kill_started.wait()
default_task = asyncio.create_task(system._start_default_group())
await asyncio.sleep(0)
allow_kill.set()
await asyncio.gather(login_task, default_task)
self.assertEqual(manager.state["login_status"], "logining")
self.assertFalse(manager.state["default_running"])
self.assertIsNone(manager.state["current_group"])
runner.start_groups.assert_awaited_once_with(["扫码上号"])
async def test_login_watchdog_runs_even_if_default_state_leaks(self):
with tempfile.TemporaryDirectory() as temp_dir:
manager = QueueManager(temp_dir, logging.getLogger("login-watchdog-default-leak"))
manager.join_queue(101)
manager.state.update({
"login_status": "logining",
"login_started_at": (datetime.now() - timedelta(seconds=300)).isoformat(),
"default_running": True,
"current_group": "薄荷",
})
manager._save()
system = self._make_system(manager, SimpleNamespace())
system._handle_login_timeout = AsyncMock()
with patch(
"app.danmu_queue.asyncio.sleep",
side_effect=[None, asyncio.CancelledError()],
):
with self.assertRaises(asyncio.CancelledError):
await system._check_login_watchdog_loop()
system._handle_login_timeout.assert_awaited_once_with(101, "tester")
async def test_user_group_completion_starts_default_while_next_admin_waits(self):
with tempfile.TemporaryDirectory() as temp_dir:
manager = QueueManager(temp_dir, logging.getLogger("idle-default-flow-test"))
manager.state.update({
"queue": [1, 2],
"current_admin_uid": 1,
"current_group": "晶蝶",
"current_group_run_id": "user-run",
"group_start_time": (datetime.now() - timedelta(seconds=200)).isoformat(),
"login_status": "logged_in",
"billing_uid": 1,
"billing_started_at": datetime.now().isoformat(),
"billing_last_at": datetime.now().isoformat(),
"has_user_finished_once": True,
})
manager._save()
system = QueueSystem.__new__(QueueSystem)
system.queue_mgr = manager
system.user_mgr = SimpleNamespace(
users={"1": {"uname": "first"}, "2": {"uname": "second"}},
get_points=lambda _uid: 10,
)
system.log_monitor = SimpleNamespace(set_current_group=MagicMock())
system.logger = logging.getLogger("idle-default-flow-test")
system.broadcast = AsyncMock()
system._start_default_group = AsyncMock()
await system._on_group_finished("晶蝶", "user-run")
self.assertEqual(manager.state["queue"], [2])
self.assertEqual(manager.state["current_admin_uid"], 2)
self.assertIsNone(manager.state["login_status"])
system._start_default_group.assert_awaited_once_with()
async def test_default_group_can_start_with_waiting_admin(self):
with tempfile.TemporaryDirectory() as temp_dir:
manager = QueueManager(temp_dir, logging.getLogger("idle-default-start-test"))
manager.state["has_user_finished_once"] = True
manager.join_queue(101)
original_window = manager.state["admin_window_end"]
system = QueueSystem.__new__(QueueSystem)
system.queue_mgr = manager
system.config = SimpleNamespace(default_group="薄荷")
system.user_mgr = SimpleNamespace(
users={"101": {"uname": "队首昵称"}},
_is_masked_uname=lambda name: (not name) or ("*" in name),
resolve_uname=AsyncMock(return_value=""),
)
system.runner = SimpleNamespace(
kill_bgi=AsyncMock(),
start_groups=AsyncMock(return_value=True),
)
system.log_monitor = SimpleNamespace(set_current_group=MagicMock())
system.stats_store = None
system.logger = logging.getLogger("idle-default-start-test")
system.broadcast = AsyncMock()
system._default_start_lock = asyncio.Lock()
await system._start_default_group()
self.assertTrue(manager.state["default_running"])
self.assertEqual(manager.state["current_group"], "薄荷")
self.assertEqual(manager.state["current_admin_uid"], 101)
self.assertEqual(manager.state["admin_window_end"], original_window)
self.assertIsNone(manager.state["billing_uid"])
system.runner.start_groups.assert_awaited_once_with(["薄荷"])
# 播报文案应使用昵称而非纯数字 UID
broadcast_texts = [str(call.args[0]) for call in system.broadcast.await_args_list]
self.assertTrue(any("队首昵称" in t for t in broadcast_texts))
self.assertFalse(any("队首101" in t for t in broadcast_texts))
if __name__ == "__main__":
unittest.main()
+73
View File
@@ -0,0 +1,73 @@
import asyncio
import json
import logging
import time
import unittest
from unittest.mock import AsyncMock, patch
from app.mpv_player import MpvPlayer
class _FakePipe:
def __init__(self):
self.closed = False
self.response = b""
def write(self, request):
payload = json.loads(request.decode("utf-8"))
self.response = json.dumps({
"request_id": payload["request_id"],
"error": "success",
"data": payload["command"][1],
}).encode("utf-8") + b"\n"
def readline(self):
response, self.response = self.response, b""
return response
def close(self):
self.closed = True
class MpvPipeTests(unittest.TestCase):
def test_pipe_connection_is_reused_across_commands(self):
player = MpvPlayer("mpv.exe", logging.getLogger("test-mpv"))
fake_pipe = _FakePipe()
with patch("builtins.open", return_value=fake_pipe) as open_pipe:
self.assertEqual(player._pipe_request_sync(["get_property", "duration"], 1), "duration")
self.assertEqual(player._pipe_request_sync(["get_property", "time-pos"], 2), "time-pos")
open_pipe.assert_called_once()
player._reset_pipe_sync()
self.assertTrue(fake_pipe.closed)
class MpvMaintainTests(unittest.IsolatedAsyncioTestCase):
async def test_maintain_returns_snapshot_for_caller_reuse(self):
player = MpvPlayer("mpv.exe", logging.getLogger("test-mpv-maintain"))
player.current_url = "https://example.invalid/audio"
player.desired_state = "playing"
player.started_at = time.time() - 10
player.last_progress_at = time.time()
state = {
"playing": True,
"paused": False,
"idle": False,
"eof": False,
"path": player.current_url,
"current": {"progress": 10.0, "duration": 100.0},
}
with patch.object(player, "running", return_value=True), patch.object(
player, "snapshot", AsyncMock(return_value=state)
) as snapshot:
result = await player.maintain()
self.assertEqual(result["action"], "none")
self.assertIs(result["snapshot"], state)
snapshot.assert_awaited_once()
if __name__ == "__main__":
unittest.main()
+136
View File
@@ -0,0 +1,136 @@
import logging
import unittest
from unittest.mock import AsyncMock, patch
from app.netease_qr_login import NeteaseQrLogin, _weapi_form
class _Headers:
def __init__(self, set_cookie=None):
self.set_cookie = list(set_cookie or [])
def get_all(self, name, default=None):
if name.lower() == "set-cookie":
return self.set_cookie
return default or []
class _Response:
def __init__(self, set_cookie=None):
self.headers = _Headers(set_cookie)
class NeteaseQrLoginTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.saved_music_u = ""
self.save_callback = AsyncMock(side_effect=self._save)
self.account_checker = AsyncMock(return_value={
"authenticated": True,
"user_id": "10086",
"nickname": "测试账号",
"vip_type": 11,
})
self.login = NeteaseQrLogin(
api_base=lambda: "https://music.163.com",
get_saved_music_u=lambda: self.saved_music_u,
save_music_u=self.save_callback,
account_checker=self.account_checker,
logger=logging.getLogger("test-netease-qr"),
)
async def _save(self, music_u, _account):
self.saved_music_u = music_u
async def _start(self):
with patch(
"app.netease_qr_login._request_weapi_json",
return_value=({"code": 200, "unikey": "private-key"}, _Response()),
):
return await self.login.start()
def test_weapi_form_encrypts_payload(self):
result = _weapi_form(
{"type": 1, "csrf_token": ""},
secret_key="0123456789abcdef",
)
self.assertEqual(set(result), {"params", "encSecKey"})
self.assertEqual(len(result["encSecKey"]), 256)
self.assertNotIn("csrf_token", result["params"])
async def test_start_keeps_key_on_server(self):
with patch(
"app.netease_qr_login._request_weapi_json",
return_value=({"code": 200, "unikey": "private-key"}, _Response()),
) as request_json:
result = await self.login.start()
self.assertEqual(result["state"], "awaiting_scan")
self.assertTrue(result["has_qr_image"])
self.assertNotIn("key", result)
self.assertNotIn("qr_url", result)
self.assertFalse(result["credential_configured"])
self.assertTrue(request_json.call_args.args[0].endswith("/weapi/login/qrcode/unikey"))
self.assertEqual(request_json.call_args.args[1], {"type": 1, "csrf_token": ""})
self.assertEqual(
self.login._session["qr_url"],
"http://music.163.com/login?codekey=private-key",
)
async def test_poll_reports_scan_and_confirmation_states(self):
await self._start()
with patch(
"app.netease_qr_login._request_weapi_json",
side_effect=[
({"code": 801, "message": "等待扫码"}, _Response()),
({"code": 802, "message": "待确认"}, _Response()),
],
) as request_json:
waiting = await self.login.poll()
confirming = await self.login.poll()
self.assertEqual(waiting["state"], "awaiting_scan")
self.assertEqual(confirming["state"], "awaiting_confirm")
self.assertTrue(
request_json.call_args_list[0].args[0].endswith(
"/weapi/login/qrcode/client/login"
)
)
self.assertEqual(
request_json.call_args_list[0].args[1],
{"key": "private-key", "type": 1, "csrf_token": ""},
)
self.save_callback.assert_not_awaited()
async def test_success_validates_and_auto_saves_music_u(self):
await self._start()
response = _Response(["MUSIC_U=private-music-u; Path=/; HttpOnly; SameSite=None"])
with patch(
"app.netease_qr_login._request_weapi_json",
return_value=({"code": 803, "message": "授权登录成功"}, response),
):
result = await self.login.poll()
self.assertEqual(result["state"], "completed")
self.assertTrue(result["credential_configured"])
self.assertEqual(result["account"]["nickname"], "测试账号")
self.assertNotIn("music_u", result)
self.account_checker.assert_awaited_once_with("https://music.163.com", "private-music-u")
self.save_callback.assert_awaited_once()
self.assertEqual(self.saved_music_u, "private-music-u")
async def test_success_without_music_u_is_rejected(self):
await self._start()
with patch(
"app.netease_qr_login._request_weapi_json",
return_value=({"code": 803}, _Response()),
):
result = await self.login.poll()
self.assertEqual(result["state"], "failed")
self.assertIn("缺少 MUSIC_U", result["message"])
self.save_callback.assert_not_awaited()
if __name__ == "__main__":
unittest.main()
+41
View File
@@ -0,0 +1,41 @@
import logging
import unittest
from unittest.mock import AsyncMock
from app.netease_resolver import NeteaseResolver
class NeteaseResolverTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.resolver = NeteaseResolver(logging.getLogger("test-netease"))
async def test_rejects_trial_audio_instead_of_playing_preview(self):
self.resolver._fetch_player_entry = AsyncMock(return_value={
"url": "https://example.test/trial.mp3",
"freeTrialInfo": {"start": 0, "end": 15000},
})
self.resolver._probe_url = AsyncMock(return_value="https://example.test/trial.mp3")
result = await self.resolver.resolve({"id": "123456", "name": "VIP歌曲"})
self.assertIsNone(result)
self.assertEqual(self.resolver.last_error_code, "preview_only")
self.resolver._probe_url.assert_not_awaited()
async def test_uses_authenticated_full_player_url(self):
self.resolver.update_auth("valid-cookie")
self.resolver._fetch_player_entry = AsyncMock(return_value={
"url": "https://example.test/full.mp3",
"freeTrialInfo": None,
"freeTrialPrivilege": {"resConsumable": False, "userConsumable": False},
})
self.resolver._probe_url = AsyncMock(return_value="https://cdn.test/full.mp3")
result = await self.resolver.resolve({"id": "123456", "name": "VIP歌曲"})
self.assertEqual(result["url"], "https://cdn.test/full.mp3")
self.assertEqual(result["source"], "netease.player.auth")
if __name__ == "__main__":
unittest.main()
+105
View File
@@ -0,0 +1,105 @@
import sqlite3
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from app.stats_store import StatsStore, business_date
class StatsStoreDurationRepairTests(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
self.database_path = Path(self.temp_dir.name) / "statistics.sqlite3"
self.connection = sqlite3.connect(self.database_path)
StatsStore._migrate(self.connection)
self.store = StatsStore(self.database_path)
def tearDown(self):
self.connection.close()
self.temp_dir.cleanup()
def _insert_bilibili(self, connection_id, started, ended=None, reason=None):
self.connection.execute(
"INSERT INTO bilibili_connections ("
"connection_id, connected_at_utc, disconnected_at_utc, business_date, status, disconnect_reason"
") VALUES (?, ?, ?, ?, ?, ?)",
(connection_id, started, ended, business_date(started), "failed" if ended else "connected", reason),
)
StatsStore._upsert_live_session(self.connection, {
"session_id": connection_id,
"kind": "bilibili_connection",
"started_at_utc": started,
"ended_at_utc": ended,
"status": "failed" if ended else "connected",
"source": "bilibili_connections",
})
def test_daily_query_merges_overlapping_lifecycles(self):
self._insert_bilibili("b1", "2026-07-16T00:00:00Z", "2026-07-16T02:00:00Z")
self._insert_bilibili("b2", "2026-07-16T01:00:00Z", "2026-07-16T03:00:00Z")
self.connection.execute(
"INSERT INTO group_runs (group_run_id, started_at_utc, ended_at_utc, business_date, status) "
"VALUES ('g1', '2026-07-16T00:00:00Z', '2026-07-16T02:00:00Z', '2026-07-16', 'completed')"
)
self.connection.execute(
"INSERT INTO group_runs (group_run_id, started_at_utc, ended_at_utc, business_date, status) "
"VALUES ('g2', '2026-07-16T01:00:00Z', '2026-07-16T04:00:00Z', '2026-07-16', 'completed')"
)
self.connection.commit()
row = self.store._query_daily_sync("2026-07-16", "2026-07-16")[0]
self.assertEqual(row["bilibili_connection_duration_ms"], 3 * 60 * 60 * 1000)
self.assertEqual(row["group_run_total_duration_ms"], 4 * 60 * 60 * 1000)
def test_daily_query_splits_at_beijing_0400_boundary(self):
self.connection.execute(
"INSERT INTO group_runs (group_run_id, started_at_utc, ended_at_utc, business_date, status) "
"VALUES ('g1', '2026-07-16T19:30:00Z', '2026-07-16T20:30:00Z', '2026-07-16', 'completed')"
)
self.connection.commit()
rows = self.store._query_daily_sync("2026-07-16", "2026-07-17")
by_day = {row["business_date"]: row for row in rows}
self.assertEqual(by_day["2026-07-16"]["group_run_total_duration_ms"], 30 * 60 * 1000)
self.assertEqual(by_day["2026-07-17"]["group_run_total_duration_ms"], 30 * 60 * 1000)
def test_reconciliation_ends_stale_connection_at_next_start(self):
self._insert_bilibili("b1", "2026-07-16T00:00:00Z")
self._insert_bilibili("b2", "2026-07-16T00:05:00Z")
self.connection.commit()
with patch.object(StatsStore, "_pid_is_running", return_value=False):
StatsStore._reconcile_stale_lifecycles(self.connection)
first = self.connection.execute(
"SELECT disconnected_at_utc, disconnect_reason FROM bilibili_connections WHERE connection_id='b1'"
).fetchone()
self.assertEqual(first, ("2026-07-16T00:05:00Z", "startup_reconciliation"))
def test_historical_repair_is_idempotent_and_rebuilds_daily_duration(self):
self._insert_bilibili(
"b1", "2026-07-16T00:00:00Z", "2026-07-17T00:00:00Z", "startup_reconciliation"
)
self._insert_bilibili("b2", "2026-07-16T00:05:00Z", "2026-07-16T00:10:00Z")
self.connection.commit()
first = StatsStore._repair_startup_reconciliation_history(self.connection)
second = StatsStore._repair_startup_reconciliation_history(self.connection)
ended = self.connection.execute(
"SELECT disconnected_at_utc FROM bilibili_connections WHERE connection_id='b1'"
).fetchone()[0]
duration = self.connection.execute(
"SELECT SUM(duration_ms) FROM live_session_daily_durations WHERE session_id='b1'"
).fetchone()[0]
self.assertEqual(first["bilibili_connections"], 1)
self.assertEqual(second["bilibili_connections"], 0)
self.assertEqual(ended, "2026-07-16T00:05:00Z")
self.assertEqual(duration, 5 * 60 * 1000)
if __name__ == "__main__":
unittest.main()
+159
View File
@@ -0,0 +1,159 @@
import asyncio
import logging
import unittest
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock, patch
from app.danmu_queue import SystemScheduler
class _FakeConfig:
def __init__(self, system_cfg):
self.system_cfg = system_cfg
class RebootAfterStopTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.logger = logging.getLogger("test-system-scheduler")
self.scheduler = SystemScheduler(
_FakeConfig({
"reboot_after_stop_enabled": True,
"reboot_after_stop_delay_sec": 60,
}),
self.logger,
)
@patch("app.danmu_queue.os.name", "nt")
@patch("app.danmu_queue.subprocess.Popen")
def test_schedules_one_minute_reboot_only_once(self, popen):
stopped_at = datetime(2026, 7, 23, 23, 0)
self.assertTrue(self.scheduler._schedule_reboot_after_stop(stopped_at))
self.assertFalse(self.scheduler._schedule_reboot_after_stop(stopped_at))
popen.assert_called_once()
command = popen.call_args.args[0]
self.assertEqual(command[:5], ["shutdown", "/r", "/t", "60", "/c"])
async def test_failed_stop_still_schedules_reboot(self):
self.scheduler._live_occurrences = MagicMock(
return_value=[("stop", datetime.now().replace(second=0, microsecond=0))]
)
self.scheduler.stop_bilibili_live = AsyncMock(return_value=False)
self.scheduler._schedule_reboot_after_stop = MagicMock()
self.scheduler._end_live_statistics = MagicMock()
async def stop_after_first_sleep(_seconds):
self.scheduler._stop = True
with patch("app.danmu_queue.asyncio.sleep", side_effect=stop_after_first_sleep):
await self.scheduler.run()
self.scheduler._schedule_reboot_after_stop.assert_called_once()
async def test_stop_exception_still_schedules_reboot(self):
self.scheduler._live_occurrences = MagicMock(
return_value=[("stop", datetime.now().replace(second=0, microsecond=0))]
)
self.scheduler.stop_bilibili_live = AsyncMock(side_effect=RuntimeError("click failed"))
self.scheduler._schedule_reboot_after_stop = MagicMock()
self.scheduler._end_live_statistics = MagicMock()
async def stop_after_first_sleep(_seconds):
self.scheduler._stop = True
with patch("app.danmu_queue.asyncio.sleep", side_effect=stop_after_first_sleep):
await self.scheduler.run()
self.scheduler._schedule_reboot_after_stop.assert_called_once()
class StartupLiveCompensationTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.logger = logging.getLogger("test-startup-live-compensation")
self.scheduler = SystemScheduler(
_FakeConfig({
"live_start_time": "09:00",
"live_end_time": "23:00",
"launch_bilibili_live_enabled": True,
"launch_genshin_enabled": True,
"bilibili_push_enabled": True,
}),
self.logger,
)
self.start_at = datetime.now() - timedelta(hours=1)
self.end_at = datetime.now() + timedelta(hours=1)
self.scheduler._startup_window_active = MagicMock(return_value=True)
self.scheduler._is_exe_running = MagicMock(return_value=False)
self.scheduler._launch_exe = AsyncMock(side_effect=[True, True])
self.scheduler.push_bilibili_live = AsyncMock(return_value=True)
async def test_launches_apps_after_one_minute_then_pushes_after_another(self):
with patch("app.danmu_queue.asyncio.sleep", new_callable=AsyncMock) as sleep:
await self.scheduler._run_startup_live_compensation(self.start_at, self.end_at)
self.assertEqual(
[call.args[0] for call in sleep.await_args_list],
[60, 60],
)
self.assertEqual(
[call.args for call in self.scheduler._launch_exe.await_args_list],
[("B站直播姬", "bilibili_live_exe"), ("原神", "genshin_exe")],
)
self.scheduler.push_bilibili_live.assert_awaited_once()
async def test_skips_push_when_livehime_was_already_running(self):
self.scheduler._is_exe_running = MagicMock(return_value=True)
self.scheduler._launch_exe = AsyncMock(side_effect=[False, False])
with patch("app.danmu_queue.asyncio.sleep", new_callable=AsyncMock):
await self.scheduler._run_startup_live_compensation(self.start_at, self.end_at)
self.scheduler.push_bilibili_live.assert_not_awaited()
async def test_cancels_before_push_when_live_window_ends(self):
self.scheduler._startup_window_active = MagicMock(side_effect=[True, False])
with patch("app.danmu_queue.asyncio.sleep", new_callable=AsyncMock):
await self.scheduler._run_startup_live_compensation(self.start_at, self.end_at)
self.scheduler.push_bilibili_live.assert_not_awaited()
async def test_scheduling_marks_prepare_and_start_to_avoid_duplicate_clicks(self):
with patch.object(self.scheduler, "_run_startup_live_compensation", new_callable=AsyncMock):
scheduled = self.scheduler._schedule_startup_live_compensation(self.start_at, self.end_at)
await asyncio.sleep(0)
if self.scheduler._startup_compensation_task:
await self.scheduler._startup_compensation_task
self.assertTrue(scheduled)
self.assertFalse(self.scheduler._schedule_startup_live_compensation(self.start_at, self.end_at))
self.assertIn(
self.scheduler._event_key("prepare_tts", self.start_at - timedelta(minutes=15)),
self.scheduler._triggered_events,
)
self.assertIn(
self.scheduler._event_key("prepare", self.start_at - timedelta(minutes=10)),
self.scheduler._triggered_events,
)
self.assertIn(
self.scheduler._event_key("start", self.start_at),
self.scheduler._triggered_events,
)
def test_current_live_window_supports_cross_midnight_schedule(self):
self.scheduler.config.system_cfg["live_start_time"] = "23:00"
self.scheduler.config.system_cfg["live_end_time"] = "02:00"
current = datetime(2026, 7, 31, 1, 0)
window = self.scheduler._current_live_window(current)
self.assertEqual(
window,
(datetime(2026, 7, 30, 23, 0), datetime(2026, 7, 31, 2, 0)),
)
if __name__ == "__main__":
unittest.main()
+200
View File
@@ -0,0 +1,200 @@
import asyncio
import logging
import time
import unittest
from unittest.mock import patch
import numpy as np
from app.danmu_queue import Broadcaster, TTSEngine, _BoundedPriorityQueue, _StreamingAudioBuffer
class BoundedPriorityQueueTests(unittest.IsolatedAsyncioTestCase):
async def test_returns_high_priority_first(self):
queue = _BoundedPriorityQueue(4)
await queue.put({"priority": 2, "sequence": 1, "name": "normal"})
await queue.put({"priority": 0, "sequence": 2, "name": "urgent"})
self.assertEqual((await queue.get())["name"], "urgent")
self.assertEqual((await queue.get())["name"], "normal")
async def test_replaces_worst_pending_item_when_full(self):
queue = _BoundedPriorityQueue(2)
await queue.put({"priority": 3, "sequence": 1, "name": "low-old"})
await queue.put({"priority": 3, "sequence": 2, "name": "low-new"})
accepted, dropped = await queue.put(
{"priority": 0, "sequence": 3, "name": "urgent"}
)
self.assertTrue(accepted)
self.assertEqual(dropped["name"], "low-new")
self.assertEqual((await queue.get())["name"], "urgent")
self.assertEqual((await queue.get())["name"], "low-old")
async def test_rejects_new_item_when_it_is_not_more_valuable(self):
queue = _BoundedPriorityQueue(1)
await queue.put({"priority": 0, "sequence": 1, "name": "urgent"})
accepted, dropped = await queue.put(
{"priority": 3, "sequence": 2, "name": "low"}
)
self.assertFalse(accepted)
self.assertIsNone(dropped)
self.assertEqual((await queue.get())["name"], "urgent")
class StreamingAudioBufferTests(unittest.TestCase):
def test_error_unblocks_consumer(self):
buffer = _StreamingAudioBuffer()
error = RuntimeError("failed")
buffer.finish(error)
self.assertTrue(buffer.ready.is_set())
self.assertIs(buffer.error, error)
self.assertIs(buffer.chunks.get_nowait(), _StreamingAudioBuffer.END)
class StreamingPreparationTests(unittest.IsolatedAsyncioTestCase):
async def test_prepare_returns_after_first_chunk_before_generation_finishes(self):
generation_gate = asyncio.Event()
class FakeStreamingEngine:
streaming = True
async def _synthesize_to_buffer(self, text, buffer):
buffer.put(np.zeros(128, dtype=np.float32), 24000)
await generation_gate.wait()
buffer.finish()
tts = TTSEngine.__new__(TTSEngine)
tts._engine = FakeStreamingEngine()
tts._synthesis_lock = asyncio.Lock()
tts._log_event = lambda _message: None
prepared = await tts.prepare_request({"request_id": "test", "text": "测试"})
self.assertEqual(prepared["kind"], "stream")
self.assertFalse(prepared["synth_task"].done())
generation_gate.set()
await prepared["synth_task"]
class PipelineOverlapTests(unittest.IsolatedAsyncioTestCase):
async def test_next_item_synthesizes_while_previous_item_is_playing(self):
first_play_started = asyncio.Event()
release_first_play = asyncio.Event()
second_synthesis_started = asyncio.Event()
class FakeTTS:
async def prepare_request(self, request):
if request["request_id"] == "second":
second_synthesis_started.set()
return {"kind": "audio", "audio": request["request_id"].encode()}
def mark_playing(self, request):
return None
async def play_prepared(self, prepared):
if prepared["audio"] == b"first":
first_play_started.set()
await release_first_play.wait()
return True
def complete_request(self, request):
return None
def fail_request(self, request, error, **kwargs):
return None
broadcaster = Broadcaster.__new__(Broadcaster)
broadcaster._stop = False
broadcaster._tts_warmup_done = asyncio.Event()
broadcaster._tts_warmup_done.set()
broadcaster._tts_pending = _BoundedPriorityQueue(4)
broadcaster._tts_playback_queue = asyncio.Queue(maxsize=2)
broadcaster.tts = FakeTTS()
broadcaster.logger = logging.getLogger("tts-pipeline-test")
broadcaster._finish_tts_job = lambda *_args, **_kwargs: None
synth_worker = asyncio.create_task(broadcaster._tts_synthesis_loop())
play_worker = asyncio.create_task(broadcaster._tts_playback_loop())
expiry = time.monotonic() + 30
await broadcaster._tts_pending.put({
"priority": 1,
"sequence": 1,
"expires_at": expiry,
"tts_request": {"request_id": "first"},
})
await broadcaster._tts_pending.put({
"priority": 1,
"sequence": 2,
"expires_at": expiry,
"tts_request": {"request_id": "second"},
})
await asyncio.wait_for(first_play_started.wait(), timeout=1)
await asyncio.wait_for(second_synthesis_started.wait(), timeout=1)
release_first_play.set()
broadcaster._stop = True
synth_worker.cancel()
play_worker.cancel()
await asyncio.gather(synth_worker, play_worker, return_exceptions=True)
class TTSExpiryPolicyTests(unittest.TestCase):
def test_default_expiry_windows_match_configured_policy(self):
broadcaster = Broadcaster.__new__(Broadcaster)
broadcaster._tts_queue_cfg = {}
self.assertEqual(broadcaster._tts_policy("login"), (0, 60.0))
self.assertEqual(broadcaster._tts_policy("system"), (1, 40.0))
self.assertEqual(broadcaster._tts_policy("queue"), (2, 40.0))
self.assertEqual(broadcaster._tts_policy("points"), (3, 30.0))
class AudioPlaybackTests(unittest.IsolatedAsyncioTestCase):
async def test_pygame_mixer_is_initialized_once_and_reused(self):
class FakeChannel:
def get_busy(self):
return False
class FakeSound:
def __init__(self, *, file):
self.file = file
def play(self):
return FakeChannel()
class FakeMixer:
def __init__(self):
self.initialized = False
self.init_calls = 0
self.Sound = FakeSound
def get_init(self):
return self.initialized
def init(self):
self.initialized = True
self.init_calls += 1
def quit(self):
self.initialized = False
class FakePygame:
mixer = FakeMixer()
tts = TTSEngine.__new__(TTSEngine)
tts._pygame = None
with patch.dict("sys.modules", {"pygame": FakePygame}):
await tts._play_audio(b"first")
await tts._play_audio(b"second")
self.assertEqual(FakePygame.mixer.init_calls, 1)
if __name__ == "__main__":
unittest.main()
+35
View File
@@ -0,0 +1,35 @@
# Python
__pycache__/
*.pyc
*.swp
*.egg-info/
build/
dist/
.venv/
venv/
# Tooling caches
.pytest_cache/
.ruff_cache/
.mypy_cache/
.coverage
# Editors / OS
.DS_Store
.idea/
.vscode/
# Project outputs
infer_output/
smoke_accel/
downloaded_data
pretrained_models
tmp_*
temp_*
debug_*
apps/gradio/*.log
# Audio outputs (keep Gradio's bundled default prompts)
*wav
!apps/gradio/default_prompts/
!apps/gradio/default_prompts/*.wav
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2026 dots.tts Team, RedNote
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+437
View File
@@ -0,0 +1,437 @@
<p align="center">
<img src="assets/logo.png" alt="dots.tts" width="280">
</p>
<p align="center">
<a href="https://github.com/rednote-hilab/dots.tts"><img src="https://img.shields.io/badge/GitHub-rednote--hilab%2Fdots.tts-blue?logo=github" alt="GitHub"></a>
<a href="https://huggingface.co/collections/rednote-hilab/dotstts"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-dots.tts%20collection-yellow" alt="Hugging Face"></a>
<a href="https://arxiv.org/abs/2606.07080"><img src="https://img.shields.io/badge/arXiv-Report-b31b1b?logo=arxiv&logoColor=white" alt="arXiv"></a>
<a href="https://huggingface.co/spaces/rednote-hilab/dots.tts"><img src="https://img.shields.io/badge/Playground-Live-orange" alt="Playground"></a>
<a href="https://rednote-hilab.github.io/dots.tts-demo/"><img src="https://img.shields.io/badge/Demo%20Page-Live-red" alt="Demo Page"></a>
<a href="LICENSE"><img src="https://img.shields.io/badge/License-Apache%202.0-green" alt="License"></a>
</p>
**dots.tts** is a **2B-parameter fully continuous, end-to-end autoregressive (AR) text-to-speech system**. The backbone pairs a semantic encoder, an LLM, and an autoregressive flow-matching acoustic head over a **48 kHz** AudioVAE, with no discrete tokens anywhere in the pipeline.
dots.tts achieves the best average performance on **Seed-TTS-Eval**, with WERs of **0.94% / 1.30% / 6.60%** and SIM scores of **81.0 / 77.1 / 79.5** on the zh / en / zh-hard test sets, respectively. It further attains the **highest average speaker similarity (83.9)** on the 24-language **MiniMax multilingual** benchmark. Across other benchmarks, dots.tts also consistently demonstrates **open-source state-of-the-art performance**, exhibiting strong generation stability, voice cloning ability, and emotional expressiveness.
### News
* **[2026.06]** 🔥 We have released **dots.tts** — 2B fully continuous AR TTS, with pretrained / self-corrective-aligned / MeanFlow-distilled checkpoints and full inference & fine-tuning code under Apache-2.0.
---
## Contents
- [Quick Start](#-quick-start)
- [Installation](#installation)
- [Checkpoints](#checkpoints)
- [CLI](#cli)
- [Python API](#python-api)
- [Web Demo (Gradio)](#web-demo-gradio)
- [Fine-tuning](#fine-tuning)
- [MeanFlow Distillation](#meanflow-distillation)
- [Usage Tips](#-usage-tips)
- [Architecture](#-architecture)
- [Performance](#-performance)
- [Seed-TTS-Eval](#seed-tts-eval)
- [MiniMax Multilingual](#minimax-multilingual-24-languages)
- [CV3-Eval](#cv3-eval)
- [EmergentTTS-Eval](#emergenttts-eval)
- [Community Projects](#-community-projects)
- [Risks and Limitations](#%EF%B8%8F-risks-and-limitations)
- [Citation](#-citation)
- [License](#-license)
---
## 🚀 Quick Start
### Installation
We recommend creating a fresh conda environment first (Python 3.103.12):
```bash
conda create -n dots_tts python=3.10 -y
conda activate dots_tts
```
Then install from source:
```bash
python -m pip install --upgrade pip
python -m pip install -e . -c constraints/recommended.txt
```
For training / linting extras:
```bash
python -m pip install -e .[full] -c constraints/recommended.txt
```
The constraints file pins the recommended versions. To use other compatible
versions, omit `-c constraints/recommended.txt`; the compatibility ranges are
declared in `pyproject.toml`.
### Checkpoints
Three pretrained checkpoints are released on Hugging Face. All three share the same backbone — choose by the quality / inference-cost tradeoff:
| Model | Description | Recommended `--num-steps` |
|---|---|:---:|
| [`rednote-hilab/dots.tts-base`](https://huggingface.co/rednote-hilab/dots.tts-base) | Pretrained checkpoint. | `10``32` (default `10`) |
| [`rednote-hilab/dots.tts-soar`](https://huggingface.co/rednote-hilab/dots.tts-soar) | Self-corrective-aligned (SCA) checkpoint on top of `dots.tts-base`. Best voice cloning performance. | `10``32` (default `10`) |
| [`rednote-hilab/dots.tts-mf`](https://huggingface.co/rednote-hilab/dots.tts-mf) | MeanFlow-distilled student from `dots.tts-soar`. Recommended if you care about inference speed. | `4` |
Pass the repo id directly to `--model-name-or-path` (or `DotsTtsRuntime.from_pretrained`) — the snapshot is fetched on first use and cached locally.
### CLI
The package installs a `dots.tts` entry point:
```bash
# Continuation voice cloning (reference audio + transcript) — recommended, best SIM
dots.tts \
--model-name-or-path rednote-hilab/dots.tts-soar \
--text "Hello, this is a zero-shot voice cloning demonstration." \
--prompt-audio /path/to/reference.wav \
--prompt-text "The exact transcript of the reference audio." \
--num-steps 10 \
--output clone.wav
# X-vector-only voice cloning (reference audio only — timbre from speaker x-vector)
dots.tts \
--model-name-or-path rednote-hilab/dots.tts-soar \
--text "Hello, this is a zero-shot voice cloning demonstration." \
--prompt-audio /path/to/reference.wav \
--num-steps 10 \
--output clone.wav
# Random-voice sampling (no reference) — only meaningful with a fine-tuned
# single-speaker checkpoint
dots.tts \
--model-name-or-path rednote-hilab/dots.tts-soar \
--text "Hello, this is a quick speech synthesis test." \
--num-steps 10 \
--output output.wav
```
Common flags:
| Flag | Description | Default |
|------|-------------|---------|
| `--num-steps` | Flow-matching sampling steps (higher = better quality, lower = faster) | `10` |
| `--guidance-scale` | CFG scale (flow-matching only; MeanFlow has CFG fused into the student; values > 2 progressively amplify audio energy) | `1.2` |
| `--normalize-text` | Apply text normalization before inference (via [WeTextProcessing](https://github.com/wenet-e2e/WeTextProcessing)) | off |
| `--language` | Add an explicit language tag to the input text; accepts `none`, `auto_detect`, language codes such as `EN` / `ZH`, or names such as `english` / `chinese` | `none` |
| `--seed` | RNG seed (fixed seed → deterministic output) | `42` |
`dots.tts --help` lists the full set.
Notes:
- `--prompt-audio` selects the speaker voice — continuation cloning when paired with `--prompt-text`, x-vector-only cloning when used alone. Omitting `--prompt-audio` falls back to random-voice sampling, which is only meaningful on a fine-tuned single-speaker checkpoint.
- `--language` is useful for multilingual or code-switched text when you want to force the model-side language tag. For example, pass `--language EN` for English, `--language ZH` for Mandarin, `--language Cantonese` for Cantonese, or `--language auto_detect` to infer the tag from `--text`.
- Pass either a local model directory or a Hugging Face repo id.
### Python API
```python
from dots_tts.runtime import DotsTtsRuntime
import soundfile as sf
runtime = DotsTtsRuntime.from_pretrained(
"rednote-hilab/dots.tts-soar",
precision="bfloat16",
optimize=True, # torch.compile acceleration (warmup at load, faster steady-state)
)
result = runtime.generate(
text="Hello, this is a quick speech synthesis test.",
prompt_audio_path="/path/to/reference.wav",
prompt_text="The exact transcript of the reference audio.",
num_steps=10,
guidance_scale=1.2,
)
sf.write("output.wav", result["audio"].float().cpu().squeeze().numpy(), result["sample_rate"])
```
For low-latency playback or streaming to a client, use `generate_stream` instead — it yields audio chunks (`torch.Tensor`, shape `(1, samples)`) as they are produced. Arguments are identical to `generate`:
```python
import torch
stream = runtime.generate_stream(
text="Hello, this is a streaming speech synthesis test.",
prompt_audio_path="/path/to/reference.wav",
prompt_text="The exact transcript of the reference audio.",
num_steps=10,
guidance_scale=1.2,
)
chunks = []
for chunk in stream:
chunks.append(chunk.detach().float().cpu())
# handle_chunk(chunk) # push to a player / websocket / etc.
audio = torch.cat(chunks, dim=-1).squeeze().numpy()
sf.write("output_stream.wav", audio, runtime.sample_rate)
```
### Web Demo (Gradio)
```bash
python apps/gradio/app.py \
--model-name-or-path rednote-hilab/dots.tts-soar \
--optimize
```
Defaults to `http://0.0.0.0:7860`. With `--optimize` the first launch runs warmup (slower startup, faster steady-state).
### Fine-tuning
This repo exposes fine-tuning and MeanFlow distillation entry points. Fine-tune from a released checkpoint with:
```bash
accelerate launch scripts/train_dots_tts.py --config configs/dots_tts.yaml
```
`configs/dots_tts.yaml` is a smoke configuration that verifies the pipeline runs end-to-end on commodity hardware. Replace `train.pretrained_model_path`, `train_data.sources` / `val_data.sources`, `train.output_dir`, and `train.max_train_steps` with your own values to use it.
A helper script downloads LJSpeech-1.1-48kHz and emits a train/valid JSONL manifest for the smoke run:
```bash
python scripts/prepare_train_jsonl_manifest.py --output-dir downloaded_data
```
Manifest format — one JSON per line, minimum three fields:
```json
{"fid": "sample-0001", "audio": "/abs/path/to/audio.wav", "text": "hello world"}
```
### MeanFlow Distillation
MeanFlow distillation trains a MeanFlow DiT student against a frozen flow-matching teacher. The teacher can be the released SOAR checkpoint or any compatible flow-matching dots.tts checkpoint you have fine-tuned yourself.
To use SOAR as the teacher, download it first:
```bash
huggingface-cli download rednote-hilab/dots.tts-soar \
--local-dir pretrained_models/dots.tts-soar
```
Then launch distillation with the MeanFlow config:
```bash
accelerate launch \
--num_processes 2 \
--mixed_precision bf16 \
scripts/train_dots_tts_meanflow.py \
--config configs/dots_tts_meanflow.yaml \
--teacher-model-path pretrained_models/dots.tts-soar
```
To distill from your own fine-tuned teacher, pass that checkpoint instead:
```bash
accelerate launch \
--num_processes 2 \
--mixed_precision bf16 \
scripts/train_dots_tts_meanflow.py \
--config configs/dots_tts_meanflow.yaml \
--teacher-model-path /path/to/your_finetuned_teacher
```
`configs/dots_tts_meanflow.yaml` is a conservative smoke configuration that uses the same LJSpeech manifests produced by `scripts/prepare_train_jsonl_manifest.py`. Replace `train.pretrained_model_path`, `--teacher-model-path`, `train_data.sources` / `val_data.sources`, `train.output_dir`, and `train.max_train_steps` for your own distillation run.
By default, the script initializes the student from `train.pretrained_model_path`, adds the MeanFlow duration embedding, freezes the non-DiT modules, and trains `student.core.velocity_field_predictor`. MeanFlow does not run a separate CFG branch at inference time; the default `fused` mode distills the guided teacher target into the student. Training checkpoints save the MeanFlow student only; the frozen teacher is not written into the checkpoint model directory. Pass `--train-all-parameters` only if you want to update the full dots.tts model.
Common MeanFlow flags:
| Flag | Description | Default |
|------|-------------|---------|
| `--teacher-model-path` | Frozen flow-matching teacher directory. Defaults to `train.pretrained_model_path` if omitted. | `train.pretrained_model_path` |
| `--teacher-steps` | Teacher rollout steps used to build the distillation target. Higher is slower and usually stronger. | `8` |
| `--teacher-solver` | Teacher ODE solver: `euler`, `midpoint`, or `rk4`. | `euler` |
| `--cfg-distill-mode` | `fused` distills a guided teacher target into the student; `natural` trains on sampled conditional/unconditional masks without fusing CFG. | `fused` |
| `--distill-cfg-scale` | Extra CFG coefficient used when `--cfg-distill-mode fused` is enabled. It matches inference `guidance_scale` semantics: `teacher_cond + scale * (teacher_cond - teacher_uncond)`. | `1.2` |
| `--anchor-prob` | Probability of using a zero-duration anchor sample in MeanFlow training. | `0.5` |
| `--debug` | Print the first few batch summaries and gradient diagnostics. | off |
---
## 💡 Usage Tips
- **Keep the reference audio around 10s**. Longer audio won't yield better results.
- **`--prompt-text` should match what's actually spoken in the reference audio**. Mismatches degrade stability and may cause word-level errors.
- **Higher-quality references give better clones** — prefer a high sample rate, low background noise, no trailing noise, and natural-sounding speech.
- **Try different `--seed` values for prosody variation**. Each seed produces a different rhythm and intonation — resample a few times if the default doesn't feel right.
- **Increase `--num-steps` if quality isn't good enough**. More sampling steps trade compute for cleaner output and better expressiveness.
- **Force a pronunciation with Pinyin for polyphones.** Replace the character in the input text with its tone-marked pinyin — e.g. write `我生平不hào此道` to force `好` to be read as `hào`. Use tone-marked pinyin only (`hǎo`, `hào`, `bā`); numbered forms like `hao4` or `ha4o` are **not** recognized. Useful when reseeding doesn't fix a polyphone misread.
---
## 🏛 Architecture
A frozen **AudioVAE** encodes 48 kHz mono waveform into a continuous latent and decodes it back via a BigVGAN-style causal decoder. An **autoregressive backbone** predicts that latent one patch at a time, in three components:
- **Semantic encoder** — re-encodes each newly generated VAE patch into a compact embedding for the LLM, stripping high-variance acoustic detail.
- **LLM** — initialized from **Qwen2.5-1.5B-Base**, consumes BPE text directly (no phonemes), and emits one hidden state per audio step.
- **AR flow-matching head** — a DiT that conditions on the LLM hidden state and the AR prefix to denoise the next VAE patch, with a frozen CAM++ speaker x-vector as side input.
Two sequence layouts: *plain mode* places the full text as a prefix before the audio span (standard TTS); *[1T1A interleaved mode](scripts/example_double_streaming.py)* alternates one BPE token with one audio step, enabling low-latency streaming when driven by a duplex dialogue LLM. See the technical report for full architectural and training details.
---
## 📊 Performance
Baselines are taken from original publications or default-configuration open-source releases.
### Seed-TTS-Eval
Zero-shot, ~3 s reference prompt, scored by the benchmark's reference ASR and WavLM-SV similarity.
| Model | Params | test-en WER↓ / SIM↑ | test-zh WER↓ / SIM↑ | test-zh-hard WER↓ / SIM↑ | **Avg WER↓ / SIM↑** |
|---|---:|:---:|:---:|:---:|:---:|
| CosyVoice 3 | 1.5B | 2.22 / 72.0 | 1.12 / 78.1 | **5.83** / 75.8 | 3.06 / 75.3 |
| DiTAR | 0.6B | 1.69 / 73.5 | 1.02 / 75.3 | — | — |
| F5-TTS | 0.3B | 2.00 / 67.0 | 1.53 / 76.0 | 8.67 / 71.3 | 4.10 / 71.4 |
| FireRedTTS-2 | 1.5B | 1.95 / 66.5 | 1.14 / 73.6 | 8.98 / 70.3 | 4.02 / 70.1 |
| IndexTTS 2 | 1.5B | 2.23 / 70.6 | 1.03 / 76.5 | 7.12 / 75.5 | 3.46 / 74.2 |
| MegaTTS 3 | 0.5B | 2.79 / 77.1 | 1.52 / 79.0 | — | — |
| MiniMax-Speech | — | 1.65 / 69.2 | **0.83** / 78.3 | — | — |
| Qwen3-TTS | 1.7B | **1.23** / 71.7 | 1.22 / 77.0 | 6.76 / 74.8 | 3.07 / 74.5 |
| Seed-TTS | — | 2.25 / 76.2 | 1.12 / 79.6 | 7.59 / 77.6 | 3.65 / 77.8 |
| VibeVoice | 1.5B | 3.04 / 68.9 | 1.16 / 74.4 | — | — |
| VoxCPM 2 | 2B | 1.84 / 75.3 | 0.97 / 79.5 | 8.13 / 75.3 | 3.65 / 76.7 |
| **dots.tts (Pretrain)** | **2B** | 1.34 / 76.8 | 0.96 / 80.5 | 6.46 / 79.2 | **2.92** / 78.8 |
| **dots.tts (SCA)** | **2B** | 1.30 / **77.1** | 0.94 / **81.0** | 6.60 / **79.5** | 2.95 / **79.2** |
| **dots.tts (MF, NFE=4)** | **2B** | 1.29 / 76.2 | 0.94 / 80.0 | 6.60 / 78.5 | 2.94 / 78.2 |
### MiniMax Multilingual (24 languages)
Per-language WER / SIM on the MiniMax-Speech multilingual test set (100 utterances × 2 reference speakers per language). **Highest average SIM (83.9, SCA)**, with a dots.tts variant taking the per-language SIM lead outright on 19 of 24 languages and tying on 2 more. Content fidelity is on par with the strongest systems on high-resource / Western European splits, and trails on low-resource long-tail languages where SIM is still preserved.
<details>
<summary><b>Per-language WER / SIM (click to expand)</b></summary>
| Language | MiniMax | ElevenLabs | Fish-Audio S2 | VoxCPM 2 | **dots.tts (Pre.)** | **dots.tts (SCA)** | **dots.tts (MF$_4$)** |
|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| Arabic | **1.67** / 73.6 | **1.67** / 70.6 | 3.50 / 75.0 | 13.05 / **79.1** | 37.91 / 77.5 | 36.19 / **79.1** | 39.65 / 77.6 |
| Cantonese* | 34.11 / 77.8 | 51.51 / 67.0 | 30.67 / 80.5 | 38.58 / 83.5 | 37.91 / 84.7 | 42.32 / **85.0** | 37.82 / 84.0 |
| Chinese | 2.25 / 78.0 | 16.03 / 67.7 | **0.73** / 81.6 | 1.14 / **82.5** | 1.08 / 82.3 | 0.77 / **82.5** | 1.01 / 81.8 |
| Czech | 3.88 / 79.6 | **2.11** / 68.5 | 2.84 / 79.8 | 24.13 / 78.3 | 5.05 / 83.8 | 4.25 / **84.2** | 5.67 / 83.9 |
| Dutch | 1.14 / 73.8 | **0.80** / 68.0 | 0.99 / 73.0 | 0.91 / 80.8 | 1.20 / 81.4 | 1.39 / **82.2** | 1.30 / 82.1 |
| English | 2.16 / 75.6 | 2.34 / 61.3 | 1.62 / 79.7 | 2.29 / 85.4 | 1.06 / 86.9 | **1.03** / **87.5** | 1.09 / 86.9 |
| Finnish | 4.67 / 83.5 | 2.96 / 75.9 | 3.33 / 81.9 | **2.63** / **89.0** | 3.44 / 88.0 | 4.08 / 88.3 | 3.61 / 88.3 |
| French | 4.10 / 62.8 | 5.22 / 53.5 | **3.05** / 69.8 | 4.53 / 73.5 | 3.82 / 78.2 | 3.56 / **78.6** | 3.26 / 78.5 |
| German | 1.91 / 73.3 | 0.57 / 61.4 | **0.55** / 76.7 | 0.68 / 80.3 | 1.03 / 79.5 | 1.70 / **80.6** | 0.91 / 79.5 |
| Greek | 2.02 / 82.6 | **0.99** / 73.3 | 5.74 / 79.5 | 2.84 / 86.0 | 2.97 / **87.6** | 3.00 / **87.6** | 3.19 / 87.3 |
| Hindi | 6.96 / 81.8 | **5.83** / 73.0 | 14.64 / 82.1 | 19.70 / **85.6** | 14.32 / 84.5 | 14.24 / 84.7 | 14.75 / 84.8 |
| Indonesian | 1.24 / 72.9 | **1.06** / 66.0 | 1.46 / 76.3 | 1.08 / 80.0 | 2.71 / 80.8 | 2.96 / 80.8 | 3.91 / **81.2** |
| Italian | 1.54 / 69.9 | 1.74 / 57.9 | **1.27** / 74.7 | 1.56 / 78.0 | 3.16 / 84.5 | 3.12 / **84.7** | 2.16 / 84.3 |
| Japanese | 3.52 / 77.6 | 10.65 / 73.8 | **2.76** / 79.6 | 4.63 / 82.8 | 7.16 / 83.1 | 5.28 / **83.7** | 5.17 / 83.1 |
| Korean | 1.75 / 77.6 | 1.87 / 70.0 | **1.18** / 81.7 | 1.96 / 83.3 | 5.30 / 84.3 | 5.66 / 83.6 | 3.93 / **84.9** |
| Polish | 1.42 / 80.2 | **0.77** / 72.9 | 1.26 / 81.9 | 1.14 / **88.4** | 2.72 / 87.3 | 3.59 / 87.8 | 3.42 / 87.5 |
| Portuguese | 1.88 / 80.5 | 1.33 / 71.1 | **1.14** / 78.1 | 1.94 / 83.7 | 1.64 / 83.1 | 2.00 / **84.3** | 2.40 / 83.1 |
| Romanian | 2.88 / 80.9 | **1.35** / 69.9 | 10.74 / 73.3 | 21.58 / 79.7 | 3.36 / 86.2 | 3.87 / **87.1** | 3.38 / 86.1 |
| Russian | 4.28 / 76.1 | 3.88 / 67.6 | **2.40** / 79.0 | 3.63 / 81.1 | 3.64 / 83.0 | 4.28 / **83.2** | 4.42 / **83.2** |
| Spanish | 1.03 / 76.2 | 1.08 / 61.5 | 0.91 / 77.6 | 1.44 / 83.1 | 0.96 / 83.9 | 1.27 / **84.0** | **0.80** / **84.0** |
| Thai | **2.70** / 80.0 | 73.94 / 58.8 | 4.23 / 78.6 | 2.96 / 84.0 | 7.45 / 83.8 | 7.86 / 83.9 | 8.03 / **84.2** |
| Turkish | 1.52 / 77.9 | **0.70** / 59.6 | 0.87 / 83.5 | 0.82 / 87.1 | 5.45 / **87.4** | 4.96 / 87.3 | 6.20 / 86.8 |
| Ukrainian | 1.08 / 73.0 | **1.00** / 64.7 | 2.30 / 74.7 | 6.32 / 79.8 | 1.61 / 80.5 | 1.27 / **81.2** | 1.66 / 80.0 |
| Vietnamese | **0.88** / 74.3 | 73.42 / 36.9 | 7.41 / 74.0 | 3.31 / 80.6 | 3.85 / 80.7 | 3.89 / **81.6** | 5.43 / 80.5 |
| **Average** | **2.8** / 76.6 | 7.5 / 65.5 | 3.7 / 78.0 | 5.7 / 82.3 | 6.6 / 83.5 | 6.8 / **83.9** | 6.8 / 83.5 |
</details>
<sub>*Cantonese WER reflects an ASR-faithfulness floor common to all systems; SIM remains comparable.</sub>
### CV3-Eval
Hard-subset Chinese/English plus a cross-lingual voice-cloning split. **Takes the table top on hard-en (MF$_4$ at 4.37) and leads both cross-lingual SIM subsets (SCA at 75.0 / 72.8)**, with the post-trained variants bracketing the prior leader on the hardest English subset.
| Model | zh W↓ | en W↓ | hard-zh W↓ | hard-en W↓ | en→zh W↓ / S↑ | zh→en W↓ / S↑ |
|---|:---:|:---:|:---:|:---:|:---:|:---:|
| CosyVoice 2 | 4.08 | 6.32 | 12.58 | 11.96 | 13.50 / 63.3 | 6.47 / 64.3 |
| CosyVoice 3 (1.5B) | 3.91 | 4.99 | 9.77 | 10.55 | **8.01** / 66.9 | **4.32** / 66.4 |
| Fish-Audio S2 | **2.65** | **2.43** | 9.10 | 4.40 | — | — |
| VoxCPM 2 | 3.65 | 5.00 | **8.55** | 8.48 | — | — |
| **dots.tts (Pretrain)** | 3.51 | 5.24 | 9.69 | 5.99 | 10.88 / 74.6 | 4.97 / 71.9 |
| **dots.tts (SCA)** | 3.71 | 4.50 | 9.22 | 4.49 | 10.75 / **75.0** | 5.66 / **72.8** |
| **dots.tts (MF, NFE=4)** | 3.95 | 4.05 | 9.10 | **4.37** | 10.73 / 73.8 | 5.24 / 70.9 |
### EmergentTTS-Eval
Win-rate judged head-to-head against `gpt-4o-mini-tts` by Gemini-2.5-Pro-0506 across six expressiveness-oriented scenarios. **SCA takes the top Syntactic Complexity score in the table (65.7%) — above every closed-source system** — and Pretrain posts the **best Emotions score among open-source systems (72.7%)**.
| Model | Voice | WER↓ | Overall↑ | Emotions↑ | Paraling.↑ | Foreign↑ | C. Pron.↑ | Quest.↑ | Syntax↑ |
|---|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| Gemini-2.5-Flash-TTS\* | Zephyr | 10.39 | **70.7%** | **95.9%** | **91.3%** | 58.5% | 55.7% | **63.0%** | 57.9% |
| Gemini-2.5-Pro-TTS\* | Zephyr | 11.79 | 69.3% | 86.9% | 82.3% | 58.2% | **64.8%** | 61.3% | 61.8% |
| gpt-4o-audio-preview\* | Ballad | 11.87 | 65.2% | 88.8% | 82.1% | **60.2%** | 40.4% | 57.0% | 59.5% |
| gpt-4o-mini-tts\* | Alloy | 10.76 | 56.3% | 59.2% | 58.8% | 57.3% | 52.4% | 52.7% | 57.1% |
| *baseline: gpt-4o-mini-tts* | Alloy | 10.61 | 50.0% | — | — | — | — | — | — |
| **dots.tts (Pretrain)** | basic\_ref\_en | 10.86 | 49.2% | 72.7% | 54.7% | 39.5% | 18.0% | 48.4% | 58.4% |
| **dots.tts (MF4)** | basic\_ref\_en | 11.75 | 47.9% | 59.8% | 55.2% | 36.3% | 16.7% | 50.5% | 64.8% |
| **dots.tts (SCA)** | basic\_ref\_en | 10.45 | 47.6% | 63.9% | 52.7% | 39.4% | 16.4% | 47.0% | **65.7%** |
| Qwen3-TTS | basic\_ref\_en | 17.32 | 42.8% | 39.8% | 50.7% | 25.4% | 30.0% | 48.9% | 60.4% |
| HumeAI\* | — | 12.85 | 42.7% | 61.6% | 36.9% | 34.6% | 34.3% | 43.2% | 44.6% |
| Qwen3-TTS | Ryan | 19.65 | 42.3% | 60.5% | 62.7% | 17.1% | 9.8% | 56.4% | 43.0% |
| VoxCPM 2 | basic\_ref\_en | 11.84 | 41.1% | 42.3% | 44.1% | 33.3% | 18.6% | 53.4% | 52.3% |
| MiniMax/speech-02-hd\* | EN-narr | **10.02** | 36.6% | 40.9% | 34.3% | 34.3% | 16.3% | 47.3% | 43.9% |
| 11Labs Multilingual v2\* | Brian | 11.19 | 33.9% | 30.4% | 45.5% | 35.5% | 14.5% | 39.5% | 35.5% |
| F5-TTS | basic\_ref\_en | 16.47 | 15.3% | 26.8% | 21.6% | 1.8% | 1.4% | 14.8% | 23.8% |
<sub>\* Closed-source / commercial. Table shows a selected subset for brevity — for the full leaderboard, see [EmergentTTS-Eval-public](https://github.com/boson-ai/EmergentTTS-Eval-public/blob/main/LEADERBOARD_gemini-2.5-pro-05-06.md).</sub>
---
## 🤝 Community Projects
Third-party ports and integrations of dots.tts, maintained by the community.
| Project | Description | Maintainer |
|---|---|---|
| [dots-tts-mlx](https://github.com/sb1992/dots-tts-mlx) | Pure-MLX inference port for Apple Silicon (Python) | [@sb1992](https://github.com/sb1992) |
| [mlx-swift-dots-tts](https://github.com/sammcj/mlx-swift-dots-tts) | Native MLX Swift port for Apple Silicon (no Python runtime) | [@sammcj](https://github.com/sammcj) |
| [Dots-TTS-ComfyUI](https://github.com/Saganaki22/Dots-TTS-ComfyUI) | ComfyUI custom nodes for TTS, voice cloning, and Whisper transcription | [@Saganaki22](https://github.com/Saganaki22) |
---
## ⚠️ Risks and Limitations
- **Misuse risk.** High-fidelity zero-shot voice cloning can produce highly realistic synthetic speech. The released checkpoints are intended for research and authorized deployment. Do **not** use dots.tts for impersonation, fraud, or disinformation. Combine downstream use with consent-aware reference-audio policies, robust synthetic-speech detection, and content watermarking. Clearly mark AI-generated audio.
- **Low-resource WER gap.** A BPE backbone inherits the text LLM's language coverage at the cost of a higher data appetite. On script-divergent and under-represented languages (Arabic, Hindi, Turkish, Vietnamese) the WER gap visible on the MiniMax benchmark reflects this, and the same long tail surfaces on the Foreign Words and Complex Pronunciation scenarios of EmergentTTS-Eval. Speaker similarity is preserved across these languages.
- **Speech-heavy training.** Although the AudioVAE is trained at 48 kHz and is modality-agnostic in principle, the backbone is trained on a speech-heavy mixture. Singing and unified speech + sound generation are not covered in this release.
---
## 📖 Citation
If you find dots.tts useful, please consider citing the technical report and starring the repository.
```bibtex
@article{dotstts2026,
title = {dots.tts Technical Report},
author = {dots.tts Team},
year = {2026},
eprint = {2606.07080},
archivePrefix = {arXiv},
primaryClass = {cs.SD},
}
```
## 📄 License
dots.tts code and released checkpoints are licensed under [Apache-2.0](LICENSE).
## 🙏 Acknowledgments
- [Qwen2.5](https://github.com/QwenLM/Qwen2.5) — LLM backbone initialization.
- [DiTAR](https://arxiv.org/abs/2502.03930) and [ARDiT](https://arxiv.org/abs/2406.05551) — for the continuous-AR + per-patch diffusion design.
- [HoliTok](https://github.com/bovod-sjtu/HoliTok) — for the AudioVAE design.
- [BigVGAN](https://github.com/NVIDIA/BigVGAN) — for the vocoder design.
- [CAM++](https://github.com/alibaba-damo-academy/3D-Speaker) — for speaker x-vector encoder.
+1
View File
@@ -0,0 +1 @@
"""Application entrypoints for dots.tts."""
+1
View File
@@ -0,0 +1 @@
"""Gradio application for dots.tts."""
+663
View File
@@ -0,0 +1,663 @@
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
from typing import TYPE_CHECKING
REPO_ROOT = Path(__file__).resolve().parents[2]
SRC_ROOT = REPO_ROOT / "src"
for import_root in (REPO_ROOT, SRC_ROOT):
import_root_str = str(import_root)
if import_root_str not in sys.path:
sys.path.insert(0, import_root_str)
from apps.gradio.constants import ( # noqa: E402
DEFAULT_EXECUTION_MODE,
DEFAULT_GUIDANCE_SCALE,
DEFAULT_HOST,
DEFAULT_INPUT_TEXT,
DEFAULT_LOG_FILE,
DEFAULT_MAX_GENERATE_LENGTH,
DEFAULT_NUM_STEPS,
DEFAULT_ODE_METHOD,
DEFAULT_OUTPUT_DIR,
DEFAULT_OUTPUT_RETENTION,
DEFAULT_PORT,
DEFAULT_PRECISION,
DEFAULT_PROMPT_NAME,
DEFAULT_SEED,
DEFAULT_SPEAKER_SCALE,
)
if TYPE_CHECKING:
import gradio as gr
DEBUG_GRADIO_ENABLED = os.environ.get("DEBUG_GRADIO", "0") == "1"
PLAYGROUND_CSS = """
.gradio-container {
width: min(1600px, calc(100vw - 32px)) !important;
max-width: none !important;
margin: 0 auto !important;
padding-left: 0 !important;
padding-right: 0 !important;
}
.gradio-container,
.gradio-container .gradio-container {
--block-label-background-fill: #CCE5FF;
--block-label-text-color: #6666FF;
--block-label-border-color: #99c7ee;
--block-label-text-weight: 600;
--block-title-background-fill: #CCE5FF;
--block-title-text-color: #6666FF;
--block-title-border-color: #99c7ee;
--block-title-border-width: var(--block-label-border-width);
--block-title-radius: var(--block-label-radius);
--block-title-padding: var(--block-label-padding);
--block-title-text-size: var(--block-label-text-size);
--block-title-text-weight: 600;
}
.gradio-container label[data-testid="block-label"],
.gradio-container label[data-testid="block-label"] *,
.gradio-container span[data-testid="block-info"],
.gradio-container span[data-testid="block-info"] * {
background: #CCE5FF !important;
border-color: #99c7ee !important;
color: #6666FF !important;
fill: #6666FF !important;
font-family: Verdana, Geneva, "PingFang SC", "Hiragino Sans GB", "Microsoft YaHei", "Noto Sans CJK SC", sans-serif !important;
font-style: normal !important;
font-size: 0.78rem !important;
line-height: 1.2 !important;
letter-spacing: 0 !important;
text-transform: none !important;
}
.gradio-container label[data-testid="block-label"],
.gradio-container span[data-testid="block-info"],
.gradio-container [data-testid="block-title"],
.gradio-container .block-title {
border: var(--block-label-border-width) solid #99c7ee !important;
border-top: none !important;
border-left: none !important;
border-radius: var(--block-label-radius) !important;
box-shadow: var(--block-label-shadow) !important;
padding: var(--block-label-padding) !important;
}
.gradio-container label[data-testid="block-label"],
.gradio-container label[data-testid="block-label"] *,
.gradio-container span[data-testid="block-info"],
.gradio-container span[data-testid="block-info"] *,
.gradio-container [data-testid="block-title"],
.gradio-container [data-testid="block-title"] *,
.gradio-container .block-title,
.gradio-container .block-title * {
font-weight: 600 !important;
}
.gradio-container .block label > span,
.gradio-container .block label > span *,
.gradio-container .form label > span,
.gradio-container .form label > span *,
.gradio-container label > span:first-child,
.gradio-container label > span:first-child * {
font-weight: 600 !important;
}
.strong-label [data-testid="block-label"],
.strong-label [data-testid="block-label"] *,
.strong-label span[data-testid="block-info"],
.strong-label span[data-testid="block-info"] *,
.strong-label [data-testid="block-title"],
.strong-label [data-testid="block-title"] *,
.strong-label .block-label,
.strong-label .block-label *,
.strong-label .block-title,
.strong-label .block-title *,
.strong-label label > span:first-child,
.strong-label label > span:first-child * {
font-weight: 600 !important;
}
.gradio-container .info-text,
.gradio-container .info-text * {
font-weight: 400 !important;
}
.gradio-container input,
.gradio-container textarea,
.gradio-container select,
.gradio-container [role="textbox"],
.gradio-container [contenteditable="true"] {
font-weight: 400 !important;
}
.gradio-container label[data-testid="block-label"] > span:first-child {
display: none !important;
}
.generate-button {
background: #6666FF !important;
color: #ffffff !important;
border: 1px solid #5555ee !important;
font-family: Verdana, Geneva, sans-serif !important;
}
.generate-button:hover {
background: #5555ee !important;
}
#playground-banner {
padding: 0;
border-radius: 0;
margin-bottom: 18px;
background: transparent;
border: 0;
}
#playground-banner h1 {
margin: 0 0 4px 0;
font-size: 1.7rem;
font-weight: 700;
color: #0f172a;
letter-spacing: 0;
}
#playground-banner .subtitle {
margin: 0;
color: #1e293b;
font-size: 0.9rem;
}
.info-card {
padding: 14px 18px;
border-radius: 8px;
border: 1px solid #99c7ee;
border-left: 4px solid #2563eb;
background: transparent;
font-size: 0.86rem;
line-height: 1.55;
margin-bottom: 16px;
box-sizing: border-box;
color: #0f172a;
}
.info-card .card-title,
.info-card .notice-title {
display: block;
font-weight: 600;
font-size: 0.92rem;
color: #0f172a;
}
.info-card .card-title {
margin-bottom: 4px;
}
.info-card .notice-title {
margin-top: 8px;
margin-bottom: 4px;
}
.info-card ol,
.info-card ul {
margin: 0;
padding-left: 18px;
}
.info-card li {
margin: 2px 0;
}
.main-workspace {
gap: 18px !important;
align-items: stretch !important;
}
.prompt-column,
.synthesis-column {
gap: 14px !important;
}
.control-row,
.settings-slider-row {
gap: 14px !important;
}
.settings-card {
margin-top: 2px !important;
}
.generate-button {
margin-top: 2px !important;
width: 100% !important;
box-sizing: border-box !important;
flex: 0 0 auto !important;
min-height: 44px !important;
padding-top: 10px !important;
padding-bottom: 10px !important;
font-size: 1rem !important;
font-weight: 600 !important;
}
.output-audio {
flex: 0 0 auto !important;
min-height: 190px !important;
}
.output-audio audio {
width: 100% !important;
}
@media (max-width: 768px) {
.gradio-container {
width: calc(100vw - 20px) !important;
}
}
"""
def build_playground_theme(gr):
return gr.themes.Soft(
primary_hue="slate",
secondary_hue="slate",
neutral_hue="slate",
radius_size="md",
text_size="md",
spacing_size="md",
font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
)
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(description="dots.tts Gradio app.")
parser.add_argument("--host", default=DEFAULT_HOST, help="Server host")
parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Server port")
parser.add_argument(
"--execution-mode",
choices=("generate", "generate_stream"),
default=DEFAULT_EXECUTION_MODE,
help="Runtime execution mode fixed for the app",
)
parser.add_argument(
"--precision",
default=DEFAULT_PRECISION,
help="Inference precision fixed for the app runtime",
)
parser.add_argument(
"--optimize",
action="store_true",
help="Enable runtime optimize acceleration",
)
parser.add_argument(
"--model-name-or-path",
default=None,
help="Default model directory or Hugging Face repo id",
)
parser.add_argument(
"--output-dir",
default=str(DEFAULT_OUTPUT_DIR),
help="Directory for generated wav outputs",
)
parser.add_argument(
"--log-file",
default=str(DEFAULT_LOG_FILE),
help="Path to the Gradio log file",
)
parser.add_argument(
"--output-retention-count",
type=int,
default=DEFAULT_OUTPUT_RETENTION,
help="Maximum number of generated wav files to keep",
)
parser.add_argument(
"--max-generate-length",
type=int,
default=DEFAULT_MAX_GENERATE_LENGTH,
help="Maximum generation schedule length fixed for the app runtime",
)
parser.add_argument(
"--default-prompt-name",
default=DEFAULT_PROMPT_NAME,
help="Default built-in voice preset name",
)
parser.add_argument(
"--default-precision",
default=DEFAULT_PRECISION,
choices=["bfloat16", "float32", "float16"],
help="Default precision selected in the UI",
)
parser.add_argument(
"--default-num-steps",
type=int,
default=DEFAULT_NUM_STEPS,
help="Default Num Steps selected in the UI",
)
parser.add_argument(
"--default-guidance-scale",
type=float,
default=DEFAULT_GUIDANCE_SCALE,
help="Default Guidance Scale selected in the UI",
)
parser.add_argument(
"--default-speaker-scale",
type=float,
default=DEFAULT_SPEAKER_SCALE,
help="Default Speaker Scale selected in the UI",
)
parser.add_argument(
"--default-max-generate-length",
type=int,
default=DEFAULT_MAX_GENERATE_LENGTH,
help="Default Max Generate Length selected in the UI",
)
parser.add_argument(
"--skip-warmup",
action="store_true",
help="Start the Gradio server without running an initial synthesis warmup.",
)
return parser.parse_args(argv)
def build_startup_config_panel(gr, app_config) -> None:
with gr.Accordion("启动固定参数", open=False):
gr.Markdown("只读。修改这部分需要重启服务并传入新的启动参数。")
gr.Textbox(
label="Model",
value=app_config.default_model_name_or_path,
interactive=False,
)
with gr.Row():
gr.Textbox(
label="Execution Mode",
value=app_config.execution_mode,
interactive=False,
)
gr.Textbox(
label="Precision",
value=app_config.precision,
interactive=False,
)
with gr.Row():
gr.Number(
label="Max Generate Length",
value=app_config.max_generate_length,
precision=0,
interactive=False,
)
gr.Checkbox(
label="Optimize",
value=app_config.optimize,
interactive=False,
)
def build_demo(gr, app_config, app_service) -> "gr.Blocks":
from apps.gradio.service import (
GRADIO_SYNTHESIS_MODE_CHOICES,
SynthesisRequest,
build_prompt_choice_items,
resolve_prompt_selection,
)
def select_prompt_preset(prompt_name: str):
audio_path, prompt_text = resolve_prompt_selection(
prompt_name,
app_config.prompt_presets,
)
return audio_path, prompt_text
def run_synthesis(
text: str,
synthesis_mode: str,
prompt_audio_path: str | None,
prompt_text: str,
ode_method: str,
num_steps: float,
guidance_scale: float,
speaker_scale: float,
normalize_text: bool,
seed: float,
):
resolved_synthesis_mode = synthesis_mode if DEBUG_GRADIO_ENABLED else "tts"
request = SynthesisRequest(
model_name_or_path=app_config.default_model_name_or_path,
text=text,
prompt_audio_path=prompt_audio_path,
prompt_text=prompt_text,
execution_mode=app_config.execution_mode,
template_name=resolved_synthesis_mode,
ode_method=ode_method,
num_steps=int(num_steps),
guidance_scale=float(guidance_scale),
speaker_scale=float(speaker_scale),
normalize_text=normalize_text,
seed=int(seed),
)
result = app_service.generate(request)
return result.audio_path, result.metrics
show_prompt_preset = bool(app_config.prompt_presets)
with gr.Blocks(title="dots.tts") as demo:
gr.HTML(
"<style>\n"
+ PLAYGROUND_CSS
+ "\n</style>\n"
+ """
<div id="playground-banner">
<h1>dots.tts</h1>
<p class="subtitle">Fully-continuous Autoregressive TTS · 48 kHz · Voice Cloning</p>
</div>
""",
)
gr.HTML(
"""
<div class="info-card">
<span class="card-title">使用说明 · Instructions</span>
<ol>
<li>上传参考音频并填写对应转写文本 · Upload prompt audio and fill in its transcript.</li>
<li>在文本框中输入要合成的内容 · Enter the text to synthesize.</li>
<li>点击 <b>Generate</b> 合成声音 · Click <b>Generate</b> to synthesize speech.</li>
</ol>
</div>
""",
)
with gr.Row(equal_height=True, elem_classes="main-workspace"):
with gr.Column(scale=1, min_width=480, elem_classes="prompt-column"):
prompt_preset = gr.Dropdown(
label="音色 · Voice Preset",
choices=build_prompt_choice_items(app_config.prompt_presets),
value=app_config.default_prompt_name,
info="内置音色clone样本;选择后自动填入参考音频与转写。",
elem_id="voice-preset-dropdown",
elem_classes="strong-label",
visible=show_prompt_preset,
)
prompt_audio_path = gr.Audio(
label="参考音频 · Prompt Audio",
sources=["upload"],
type="filepath",
value=app_config.default_prompt_audio_path,
elem_classes="strong-label",
)
prompt_text = gr.Textbox(
label="参考音频转写 · Prompt Text",
lines=5,
value=app_config.default_prompt_text,
placeholder="Prompt audio 对应的文本转写(continuation cloning 必填)",
elem_classes="strong-label",
)
with gr.Column(scale=1, min_width=480, elem_classes="synthesis-column"):
text = gr.Textbox(
label="待合成文本 · Text",
lines=5,
max_lines=8,
value=DEFAULT_INPUT_TEXT,
placeholder="输入待合成的文本",
elem_classes="strong-label",
)
with gr.Accordion("⚙️ Settings", open=False, elem_classes="settings-card"):
with gr.Row(elem_classes="settings-slider-row"):
num_steps = gr.Slider(
label="Num Steps",
minimum=1,
maximum=32,
step=1,
value=app_config.default_num_steps,
)
with gr.Row(elem_classes="settings-slider-row"):
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1.0,
maximum=3.0,
step=0.1,
value=app_config.default_guidance_scale,
)
with gr.Row(elem_classes="control-row"):
seed = gr.Number(
label="Seed",
value=DEFAULT_SEED,
precision=0,
scale=1,
min_width=180,
)
normalize_text = gr.Checkbox(
label="Normalize Text",
value=False,
scale=1,
min_width=180,
)
generate = gr.Button(
"Generate",
variant="primary",
size="lg",
elem_classes="generate-button",
)
audio_out = gr.Audio(
label="生成音频 · Output",
type="filepath",
elem_classes="output-audio",
)
if DEBUG_GRADIO_ENABLED:
with gr.Accordion("Debug", open=False):
synthesis_mode = gr.Dropdown(
label="SynthesisMode",
choices=list(GRADIO_SYNTHESIS_MODE_CHOICES),
value="tts",
info="选择合成模式;界面显示名会自动映射到 runtime 对应模板。",
)
ode_method = gr.Textbox(
label="ODE Method",
value=DEFAULT_ODE_METHOD,
lines=1,
)
speaker_scale = gr.Slider(
label="Speaker Scale",
minimum=0.0,
maximum=3.0,
step=0.1,
value=app_config.default_speaker_scale,
info="说话人 x-vector 强度",
)
metrics = gr.JSON(label="Metrics", value=app_service.metadata())
build_startup_config_panel(gr, app_config)
else:
synthesis_mode = gr.State(value="tts")
ode_method = gr.State(value=DEFAULT_ODE_METHOD)
speaker_scale = gr.State(value=app_config.default_speaker_scale)
metrics = gr.State(value={})
generate.click(
fn=run_synthesis,
inputs=[
text,
synthesis_mode,
prompt_audio_path,
prompt_text,
ode_method,
num_steps,
guidance_scale,
speaker_scale,
normalize_text,
seed,
],
outputs=[audio_out, metrics],
concurrency_limit=1,
)
prompt_preset.change(
fn=select_prompt_preset,
inputs=[prompt_preset],
outputs=[prompt_audio_path, prompt_text],
concurrency_limit=1,
)
return demo.queue(default_concurrency_limit=1, max_size=8)
def main() -> None:
args = parse_args()
import gradio as gr
from loguru import logger
from apps.gradio.service import GradioAppService, build_gradio_app_config
from dots_tts.utils.logging import configure_logging
configure_logging(log_file=args.log_file)
logger.info(
"Gradio app starting: host={} port={} model_name_or_path={} output_dir={} "
"log_file={} output_retention_count={} max_generate_length={} execution_mode={} precision={} optimize={} "
"default_prompt_name={} skip_warmup={}",
args.host,
args.port,
args.model_name_or_path,
args.output_dir,
args.log_file,
args.output_retention_count,
args.max_generate_length,
args.execution_mode,
args.precision,
args.optimize,
args.default_prompt_name,
args.skip_warmup,
)
app_config = build_gradio_app_config(
host=args.host,
port=args.port,
execution_mode=args.execution_mode,
precision=args.precision,
optimize=args.optimize,
model_name_or_path=args.model_name_or_path,
output_dir=Path(args.output_dir),
output_retention_count=args.output_retention_count,
max_generate_length=args.max_generate_length,
default_prompt_name=args.default_prompt_name,
default_precision=args.default_precision,
default_num_steps=args.default_num_steps,
default_guidance_scale=args.default_guidance_scale,
default_speaker_scale=args.default_speaker_scale,
default_max_generate_length=args.default_max_generate_length,
)
app_service = GradioAppService(app_config)
if args.skip_warmup:
logger.info("Gradio app warmup skipped by --skip-warmup.")
else:
warmup_metrics = app_service.warmup()
logger.info("Gradio app warmup metrics: {}", warmup_metrics)
demo = build_demo(gr, app_config, app_service)
logger.info(
"Gradio app ready: host={} port={} execution_mode={} precision={} optimize={} default_model_name_or_path={}",
app_config.host,
app_config.port,
app_config.execution_mode,
app_config.precision,
app_config.optimize,
app_config.default_model_name_or_path,
)
demo.launch(
server_name=app_config.host,
server_port=app_config.port,
theme=build_playground_theme(gr),
css=PLAYGROUND_CSS,
)
if __name__ == "__main__":
main()
+26
View File
@@ -0,0 +1,26 @@
from __future__ import annotations
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 7860
DEFAULT_OUTPUT_DIR = REPO_ROOT / "apps" / "gradio" / "outputs"
DEFAULT_LOG_FILE = REPO_ROOT / "apps" / "gradio" / "gradio.log"
DEFAULT_PROMPTS_DIR = REPO_ROOT / "apps" / "gradio" / "default_prompts"
DEFAULT_PROMPT_SOURCE_DIR = DEFAULT_PROMPTS_DIR
DEFAULT_PROMPT_MAPPING_FILE = DEFAULT_PROMPTS_DIR / "prompt_text"
DEFAULT_OUTPUT_RETENTION = 20
DEFAULT_EXECUTION_MODE = "generate_stream"
DEFAULT_PRECISION = "bfloat16"
DEFAULT_ODE_METHOD = "euler"
DEFAULT_NUM_STEPS = 10
DEFAULT_GUIDANCE_SCALE = 1.2
DEFAULT_SPEAKER_SCALE = 1.5
DEFAULT_MAX_GENERATE_LENGTH = 500
DEFAULT_SEED = 42
DEFAULT_INPUT_TEXT = ""
DEFAULT_WARMUP_TEXT = "dots.tts is a 2B-parameter fully continuous, end-to-end autoregressive (AR) text-to-speech system. The backbone pairs a semantic encoder, an LLM, and an autoregressive flow-matching acoustic head over a 48 kHz AudioVAE"
DEFAULT_PROMPT_NAME = "male_zh"
DEFAULT_PROMPT_NONE = "__none__"
PROMPT_AUDIO_SUFFIXES = (".wav", ".mp3", ".flac", ".m4a", ".ogg")
+115
View File
@@ -0,0 +1,115 @@
from __future__ import annotations
SUPPORTED_LANGUAGE_CODE_BY_NAME = {
"普通话": "ZH",
"粤语": "口音:粤语",
"北京话": "口音:北京官话",
"东北话": "口音:东北话",
"四川话": "口音:四川话",
"闽南话": "口音:闽南话",
"吴语": "口音:吴语",
"英语": "EN",
"西班牙语": "ES",
"印地语": "HI",
"阿拉伯语": "AR",
"孟加拉语": "BN",
"葡萄牙语": "PT",
"俄语": "RU",
"日语": "JA",
"法语": "FR",
"德语": "DE",
"韩语": "KO",
"意大利语": "IT",
"土耳其语": "TR",
"越南语": "VI",
"印尼语": "ID",
"乌尔都语": "UR",
"波斯语": "FA",
"泰米尔语": "TA",
"泰卢固语": "TE",
"菲律宾语": "FIL",
"马来语": "MS",
"旁遮普语": "PA",
"马拉地语": "MR",
"古吉拉特语": "GU",
"马拉雅拉姆语": "ML",
"卡纳达语": "KN",
"波兰语": "PL",
"乌克兰语": "UK",
"荷兰语": "NL",
"泰语": "TH",
"罗马尼亚语": "RO",
"斯瓦希里语": "SW",
"希伯来语": "HE",
"捷克语": "CS",
"希腊语": "EL",
"匈牙利语": "HU",
"瑞典语": "SV",
"丹麦语": "DA",
"芬兰语": "FI",
"书面挪威语": "NB",
"斯洛伐克语": "SK",
"斯洛文尼亚语": "SL",
"塞尔维亚语": "SR",
"波斯尼亚语": "BS",
"克罗地亚语": "HR",
"保加利亚语": "BG",
"马其顿语": "MK",
"立陶宛语": "LT",
"拉脱维亚语": "LV",
"爱沙尼亚语": "ET",
"冰岛语": "IS",
"爱尔兰语": "GA",
"威尔士语": "CY",
"加泰罗尼亚语": "CA",
"加利西亚语": "GL",
"奥克语": "OC",
"阿斯图里亚斯语": "AST",
"尼泊尔语": "NE",
"信德语": "SD",
"奥里亚语": "OR",
"阿萨姆语": "AS",
"普什图语": "PS",
"缅甸语": "MY",
"高棉语": "KM",
"老挝语": "LO",
"哈萨克语": "KK",
"乌兹别克语": "UZ",
"吉尔吉斯语": "KY",
"塔吉克语": "TG",
"阿塞拜疆语": "AZ",
"格鲁吉亚语": "KA",
"亚美尼亚语": "HY",
"白俄罗斯语": "BE",
"卢森堡语": "LB",
"马耳他语": "MT",
"毛利语": "MI",
"南非荷兰语": "AF",
"祖鲁语": "ZU",
"科萨语": "XH",
"约鲁巴语": "YO",
"豪萨语": "HA",
"伊博语": "IG",
"阿姆哈拉语": "AM",
"奥罗莫语": "OM",
"北索托语": "NSO",
"尼扬贾语": "NY",
"修纳语": "SN",
"索马里语": "SO",
"卢干达语": "LG",
"林加拉语": "LN",
"卢奥语": "LUO",
"坎巴语": "KAM",
"翁本杜语": "UMB",
"富拉语": "FF",
"沃洛夫语": "WO",
"中库尔德语": "CKB",
"宿务语": "CEB",
"佛得角克里奥尔语": "KEA",
"蒙古语": "MN",
"爪哇语": "JV",
}
def build_language_choice_items() -> list[tuple[str, str]]:
return [("不指定", ""), *[(name, code) for name, code in SUPPORTED_LANGUAGE_CODE_BY_NAME.items()]]
+773
View File
@@ -0,0 +1,773 @@
from __future__ import annotations
import shutil
import sys
import threading
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal
REPO_ROOT = Path(__file__).resolve().parents[2]
SRC_ROOT = REPO_ROOT / "src"
for import_root in (REPO_ROOT, SRC_ROOT):
import_root_str = str(import_root)
if import_root_str not in sys.path:
sys.path.insert(0, import_root_str)
import soundfile as sf # noqa: E402
import torch # noqa: E402
from loguru import logger # noqa: E402
from apps.gradio.constants import ( # noqa: E402
DEFAULT_EXECUTION_MODE,
DEFAULT_GUIDANCE_SCALE,
DEFAULT_HOST,
DEFAULT_MAX_GENERATE_LENGTH,
DEFAULT_NUM_STEPS,
DEFAULT_ODE_METHOD,
DEFAULT_OUTPUT_DIR,
DEFAULT_OUTPUT_RETENTION,
DEFAULT_PORT,
DEFAULT_PRECISION,
DEFAULT_PROMPT_MAPPING_FILE,
DEFAULT_PROMPT_NAME,
DEFAULT_PROMPT_NONE,
DEFAULT_PROMPT_SOURCE_DIR,
DEFAULT_PROMPTS_DIR,
DEFAULT_SEED,
DEFAULT_SPEAKER_SCALE,
DEFAULT_WARMUP_TEXT,
PROMPT_AUDIO_SUFFIXES,
)
from apps.gradio.languages import ( # noqa: E402
SUPPORTED_LANGUAGE_CODE_BY_NAME,
build_language_choice_items,
)
from dots_tts.runtime import DotsTtsRuntime # noqa: E402
from dots_tts.utils.util import seed_everything # noqa: E402
ExecutionMode = Literal["generate", "generate_stream"]
GRADIO_SYNTHESIS_MODE_CHOICES = (
("tts", "tts"),
("instruct_tts", "instruction_tts"),
("instruct_tts_general", "text_to_audio"),
)
GRADIO_SYNTHESIS_MODE_TEMPLATE_NAMES = tuple(
value for _, value in GRADIO_SYNTHESIS_MODE_CHOICES
)
@dataclass(frozen=True)
class PromptPreset:
name: str
audio_path: str
prompt_text: str
def _is_prompt_asset(path: Path) -> bool:
return path.is_file() and (
path.name == "prompt_text" or path.suffix.lower() in PROMPT_AUDIO_SUFFIXES
)
def sync_default_prompt_library(
source_dir: Path = DEFAULT_PROMPT_SOURCE_DIR,
target_dir: Path = DEFAULT_PROMPTS_DIR,
) -> None:
source_dir = Path(source_dir)
if not source_dir.is_dir():
logger.info(
"Prompt library sync skipped: source_dir={} does not exist.",
source_dir,
)
return
target_dir = Path(target_dir)
target_dir.mkdir(parents=True, exist_ok=True)
logger.info(
"Prompt library sync started: source_dir={} target_dir={}",
source_dir,
target_dir,
)
source_assets = {
asset.name: asset for asset in sorted(source_dir.iterdir()) if _is_prompt_asset(asset)
}
copied_count = 0
for asset_name, source_asset in source_assets.items():
target_asset = target_dir / asset_name
if (
not target_asset.exists()
or target_asset.stat().st_size != source_asset.stat().st_size
or target_asset.stat().st_mtime_ns != source_asset.stat().st_mtime_ns
):
shutil.copy2(source_asset, target_asset)
copied_count += 1
removed_count = 0
for target_asset in sorted(target_dir.iterdir()):
if _is_prompt_asset(target_asset) and target_asset.name not in source_assets:
target_asset.unlink(missing_ok=True)
removed_count += 1
logger.info(
"Prompt library sync completed: copied_assets={} removed_assets={} "
"available_assets={}",
copied_count,
removed_count,
len(source_assets),
)
def _load_prompt_text_map(mapping_file: Path) -> dict[str, str]:
if not mapping_file.is_file():
return {}
prompt_text_map: dict[str, str] = {}
with mapping_file.open(encoding="utf-8") as file_obj:
for raw_line in file_obj:
line = raw_line.strip()
if not line or line.startswith("#") or "|" not in line:
continue
name, text = line.split("|", 1)
prompt_text_map[name.strip()] = text.strip()
return prompt_text_map
def discover_prompt_presets(
prompts_dir: Path = DEFAULT_PROMPTS_DIR,
mapping_file: Path = DEFAULT_PROMPT_MAPPING_FILE,
) -> tuple[PromptPreset, ...]:
prompts_dir = Path(prompts_dir)
if not prompts_dir.is_dir():
return ()
prompt_text_map = _load_prompt_text_map(Path(mapping_file))
prompt_audio_paths = [
audio_path
for audio_path in sorted(prompts_dir.iterdir(), key=lambda path: (path.stem == "child", path.stem))
if audio_path.is_file() and audio_path.suffix.lower() in PROMPT_AUDIO_SUFFIXES
]
return tuple(
PromptPreset(
name=audio_path.stem,
audio_path=str(audio_path.resolve()),
prompt_text=prompt_text_map.get(audio_path.stem, ""),
)
for audio_path in prompt_audio_paths
)
def build_prompt_choice_items(
prompt_presets: tuple[PromptPreset, ...],
) -> list[tuple[str, str]]:
return [("No Preset", DEFAULT_PROMPT_NONE), *[(preset.name, preset.name) for preset in prompt_presets]]
def resolve_default_prompt_selection(
prompt_presets: tuple[PromptPreset, ...],
default_prompt_name: str = DEFAULT_PROMPT_NAME,
) -> tuple[str, str | None, str]:
if not prompt_presets:
return DEFAULT_PROMPT_NONE, None, ""
preset_by_name = {preset.name: preset for preset in prompt_presets}
selected_name = default_prompt_name if default_prompt_name in preset_by_name else prompt_presets[0].name
selected_preset = preset_by_name[selected_name]
return selected_name, selected_preset.audio_path, selected_preset.prompt_text
def resolve_prompt_selection(
prompt_name: str,
prompt_presets: tuple[PromptPreset, ...],
) -> tuple[str | None, str]:
if prompt_name == DEFAULT_PROMPT_NONE:
return None, ""
for preset in prompt_presets:
if preset.name == prompt_name:
return preset.audio_path, preset.prompt_text
return None, ""
def discover_local_model_choices(repo_root: Path = REPO_ROOT) -> list[str]:
model_root = Path(repo_root) / "pretrained_models"
if not model_root.is_dir():
return []
return sorted(
path.relative_to(repo_root).as_posix()
for path in model_root.glob("**/model")
if path.is_dir()
)
def resolve_model_name_or_path(model_name_or_path: str, repo_root: Path = REPO_ROOT) -> str:
normalized = model_name_or_path.strip()
if not normalized:
raise ValueError("model_name_or_path 不能为空。")
direct_path = Path(normalized).expanduser()
if direct_path.exists():
return str(direct_path.resolve())
repo_relative_path = Path(repo_root) / normalized
if repo_relative_path.exists():
return str(repo_relative_path.resolve())
return normalized
def default_model_name_or_path(repo_root: Path = REPO_ROOT) -> str:
discovered = discover_local_model_choices(repo_root=repo_root)
if not discovered:
return ""
return discovered[0]
@dataclass(frozen=True)
class GradioAppConfig:
host: str
port: int
execution_mode: ExecutionMode
precision: str
optimize: bool
output_dir: Path
prompts_dir: Path
output_retention_count: int
max_generate_length: int
default_model_name_or_path: str
prompt_presets: tuple[PromptPreset, ...]
default_prompt_name: str
default_prompt_audio_path: str | None
default_prompt_text: str
default_precision: str
default_num_steps: int
default_guidance_scale: float
default_speaker_scale: float
default_max_generate_length: int
local_model_choices: tuple[str, ...]
repo_root: Path = REPO_ROOT
def build_gradio_app_config(
*,
host: str = DEFAULT_HOST,
port: int = DEFAULT_PORT,
execution_mode: ExecutionMode = DEFAULT_EXECUTION_MODE,
precision: str = DEFAULT_PRECISION,
optimize: bool = False,
output_dir: Path = DEFAULT_OUTPUT_DIR,
output_retention_count: int = DEFAULT_OUTPUT_RETENTION,
max_generate_length: int = DEFAULT_MAX_GENERATE_LENGTH,
model_name_or_path: str | None = None,
default_prompt_name: str = DEFAULT_PROMPT_NAME,
default_precision: str = DEFAULT_PRECISION,
default_num_steps: int = DEFAULT_NUM_STEPS,
default_guidance_scale: float = DEFAULT_GUIDANCE_SCALE,
default_speaker_scale: float = DEFAULT_SPEAKER_SCALE,
default_max_generate_length: int = DEFAULT_MAX_GENERATE_LENGTH,
repo_root: Path = REPO_ROOT,
prompts_dir: Path = DEFAULT_PROMPTS_DIR,
prompt_source_dir: Path = DEFAULT_PROMPT_SOURCE_DIR,
) -> GradioAppConfig:
sync_default_prompt_library(
source_dir=prompt_source_dir,
target_dir=prompts_dir,
)
discovered_models = discover_local_model_choices(repo_root=repo_root)
prompt_presets = discover_prompt_presets(
prompts_dir=prompts_dir,
mapping_file=prompts_dir / "prompt_text",
)
resolved_default_prompt_name, default_prompt_audio_path, default_prompt_text = (
resolve_default_prompt_selection(
prompt_presets,
default_prompt_name=default_prompt_name,
)
)
selected_model_name_or_path = (
model_name_or_path.strip()
if model_name_or_path is not None
else default_model_name_or_path(repo_root=repo_root)
)
if not selected_model_name_or_path:
raise ValueError("No default model found. Please pass --model-name-or-path.")
if execution_mode not in ("generate", "generate_stream"):
raise ValueError(f"Unsupported execution_mode: {execution_mode}")
resolved_max_generate_length = int(max_generate_length)
if resolved_max_generate_length <= 0:
raise ValueError("max_generate_length must be positive.")
resolved_precision = precision.strip() or DEFAULT_PRECISION
logger.info(
"Gradio app config prepared: host={} port={} output_dir={} "
"output_retention_count={} max_generate_length={} execution_mode={} precision={} optimize={} "
"default_model_name_or_path={} prompt_preset_count={} language_count={} local_model_choice_count={}",
host,
port,
output_dir,
output_retention_count,
resolved_max_generate_length,
execution_mode,
resolved_precision,
bool(optimize),
selected_model_name_or_path,
len(prompt_presets),
len(SUPPORTED_LANGUAGE_CODE_BY_NAME),
len(discovered_models),
)
return GradioAppConfig(
host=host,
port=int(port),
execution_mode=execution_mode,
precision=resolved_precision,
optimize=bool(optimize),
output_dir=Path(output_dir),
prompts_dir=Path(prompts_dir),
output_retention_count=int(output_retention_count),
max_generate_length=resolved_max_generate_length,
default_model_name_or_path=selected_model_name_or_path,
prompt_presets=prompt_presets,
default_prompt_name=resolved_default_prompt_name,
default_prompt_audio_path=default_prompt_audio_path,
default_prompt_text=default_prompt_text,
default_precision=default_precision,
default_num_steps=int(default_num_steps),
default_guidance_scale=float(default_guidance_scale),
default_speaker_scale=float(default_speaker_scale),
default_max_generate_length=int(default_max_generate_length),
local_model_choices=tuple(discovered_models),
repo_root=repo_root,
)
@dataclass(frozen=True)
class SynthesisRequest:
model_name_or_path: str
text: str
prompt_audio_path: str | None = None
prompt_text: str | None = None
execution_mode: ExecutionMode = DEFAULT_EXECUTION_MODE
template_name: str = "tts"
language: str | None = None
ode_method: str = DEFAULT_ODE_METHOD
num_steps: int = DEFAULT_NUM_STEPS
guidance_scale: float = DEFAULT_GUIDANCE_SCALE
speaker_scale: float = DEFAULT_SPEAKER_SCALE
normalize_text: bool = False
seed: int = DEFAULT_SEED
@dataclass(frozen=True)
class SynthesisResult:
audio_path: str
metrics: dict[str, Any]
status: str
class GradioAppService:
def __init__(self, config: GradioAppConfig):
self.config = config
self.config.output_dir.mkdir(parents=True, exist_ok=True)
self._lock = threading.Lock()
self._runtime: DotsTtsRuntime | None = None
self._runtime_model_name_or_path: str | None = None
logger.info(
"Gradio service initialized: output_dir={} default_model_name_or_path={} "
"output_retention_count={} max_generate_length={} execution_mode={} precision={} optimize={}",
self.config.output_dir,
self.config.default_model_name_or_path,
self.config.output_retention_count,
self.config.max_generate_length,
self.config.execution_mode,
self.config.precision,
self.config.optimize,
)
def metadata(self) -> dict[str, Any]:
return {
"repo_root": str(self.config.repo_root),
"default_model_name_or_path": self.config.default_model_name_or_path,
"local_model_choices": list(self.config.local_model_choices),
"prompts_dir": str(self.config.prompts_dir),
"prompt_preset_names": [preset.name for preset in self.config.prompt_presets],
"default_prompt_name": self.config.default_prompt_name,
"output_dir": str(self.config.output_dir),
"output_retention_count": self.config.output_retention_count,
"configured_max_generate_length": self.config.max_generate_length,
"configured_execution_mode": self.config.execution_mode,
"configured_precision": self.config.precision,
"optimize": self.config.optimize,
"loaded_model_name_or_path": self._runtime_model_name_or_path,
"loaded_max_generate_length": (
self.config.max_generate_length if self._runtime is not None else None
),
"loaded_precision": (
self.config.precision if self._runtime is not None else None
),
"model_loaded": self._runtime is not None,
"host": self.config.host,
"port": self.config.port,
"default_precision": self.config.default_precision,
"default_num_steps": self.config.default_num_steps,
"default_guidance_scale": self.config.default_guidance_scale,
"default_speaker_scale": self.config.default_speaker_scale,
"default_max_generate_length": self.config.default_max_generate_length,
"supported_languages": build_language_choice_items()[1:],
"supported_template_names": list(GRADIO_SYNTHESIS_MODE_TEMPLATE_NAMES),
}
def _get_runtime(
self,
model_name_or_path: str,
) -> tuple[DotsTtsRuntime, str]:
resolved_model_name_or_path = resolve_model_name_or_path(
model_name_or_path,
repo_root=self.config.repo_root,
)
if (
self._runtime is None
or self._runtime_model_name_or_path != resolved_model_name_or_path
):
logger.info(
"Gradio runtime cache miss: requested_model={} resolved_model={} "
"max_generate_length={} execution_mode={} precision={} optimize={}",
model_name_or_path,
resolved_model_name_or_path,
self.config.max_generate_length,
self.config.execution_mode,
self.config.precision,
self.config.optimize,
)
self._runtime = DotsTtsRuntime.from_pretrained(
resolved_model_name_or_path,
precision=self.config.precision,
optimize=self.config.optimize,
max_generate_length=self.config.max_generate_length,
)
self._runtime_model_name_or_path = resolved_model_name_or_path
else:
logger.info(
"Gradio runtime cache hit: requested_model={} resolved_model={} "
"max_generate_length={} execution_mode={} precision={} optimize={}",
model_name_or_path,
resolved_model_name_or_path,
self.config.max_generate_length,
self.config.execution_mode,
self.config.precision,
self.config.optimize,
)
return self._runtime, resolved_model_name_or_path
def _build_stream_request_id(
self,
runtime: DotsTtsRuntime,
request: SynthesisRequest,
) -> str:
normalized_text, normalized_language = runtime._process_text( # noqa: SLF001
request.text,
language=request.language,
normalize=request.normalize_text,
)
normalized_prompt_text = runtime._process_prompt_text( # noqa: SLF001
request.prompt_text,
language=normalized_language,
)
if normalized_language is not None and not normalized_prompt_text:
from dots_tts.utils.text import attach_language_tag # noqa: PLC0415
normalized_text = attach_language_tag(
normalized_text,
normalized_language,
)
request_id_kwargs = {
"text": normalized_text,
"prompt_audio_path": request.prompt_audio_path,
"prompt_text": normalized_prompt_text,
"template_name": request.template_name,
}
if normalized_language is not None:
request_id_kwargs["language"] = normalized_language
return runtime._build_request_id( # noqa: SLF001
**request_id_kwargs,
)
@staticmethod
def _build_runtime_generate_kwargs(request: SynthesisRequest) -> dict[str, Any]:
runtime_kwargs: dict[str, Any] = {
"text": request.text,
"prompt_audio_path": request.prompt_audio_path,
"prompt_text": request.prompt_text,
"template_name": request.template_name,
"ode_method": request.ode_method,
"num_steps": request.num_steps,
"guidance_scale": request.guidance_scale,
"speaker_scale": request.speaker_scale,
"normalize_text": request.normalize_text,
}
if request.language is not None:
runtime_kwargs["language"] = request.language
return runtime_kwargs
def _run_stream_generation(
self,
runtime: DotsTtsRuntime,
request: SynthesisRequest,
) -> dict[str, Any]:
start_time = time.time()
chunks = [
chunk.detach().float().cpu()
for chunk in runtime.generate_stream(
**self._build_runtime_generate_kwargs(request)
)
]
if not chunks:
raise ValueError("流式生成未返回任何音频块。")
audio = torch.cat(chunks, dim=-1)
elapsed_seconds = time.time() - start_time
audio_seconds = audio.shape[-1] / runtime.sample_rate
rtf = elapsed_seconds / audio_seconds if audio_seconds > 0 else float("inf")
return {
"fid": self._build_stream_request_id(runtime, request),
"audio": audio,
"sample_rate": runtime.sample_rate,
"time_used": elapsed_seconds,
"rtf": rtf,
"chunk_count": len(chunks),
}
def warmup(self, text: str | None = None) -> dict[str, Any]:
warmup_text = (text or "").strip() or DEFAULT_WARMUP_TEXT.strip()
if not warmup_text:
raise ValueError("DEFAULT_WARMUP_TEXT 不能为空。")
with self._lock:
logger.info(
"Gradio warmup requested: default_model_name_or_path={} execution_mode={} precision={} optimize={} seed={}",
self.config.default_model_name_or_path,
self.config.execution_mode,
self.config.precision,
self.config.optimize,
DEFAULT_SEED,
)
try:
seed_everything(DEFAULT_SEED)
runtime, resolved_model_name_or_path = self._get_runtime(
self.config.default_model_name_or_path,
)
warmup_request = SynthesisRequest(
model_name_or_path=self.config.default_model_name_or_path,
text=warmup_text,
execution_mode=self.config.execution_mode,
template_name="tts",
ode_method=DEFAULT_ODE_METHOD,
num_steps=self.config.default_num_steps,
guidance_scale=self.config.default_guidance_scale,
speaker_scale=self.config.default_speaker_scale,
normalize_text=False,
seed=DEFAULT_SEED,
)
request_id = self._build_stream_request_id(runtime, warmup_request)
if self.config.execution_mode == "generate_stream":
result = self._run_stream_generation(runtime, warmup_request)
else:
start_time = time.time()
result = runtime.generate(**self._build_runtime_generate_kwargs(warmup_request))
result["time_used"] = time.time() - start_time
result["chunk_count"] = 1
audio_samples = int(result["audio"].shape[-1])
except Exception:
logger.exception(
"Gradio warmup failed: default_model_name_or_path={}",
self.config.default_model_name_or_path,
)
raise
audio_seconds = audio_samples / runtime.sample_rate
metrics = {
"request_id": request_id,
"execution_mode": self.config.execution_mode,
"chunk_count": int(result["chunk_count"]),
"resolved_model_name_or_path": resolved_model_name_or_path,
"sample_rate": runtime.sample_rate,
"elapsed_seconds": round(float(result["time_used"]), 3),
"audio_seconds": round(float(audio_seconds), 3),
"rtf": round(float(result["rtf"]), 4),
"seed": DEFAULT_SEED,
"text": warmup_text,
}
logger.info(
"Gradio warmup ready: request_id={} execution_mode={} resolved_model_name_or_path={}",
metrics["request_id"],
metrics["execution_mode"],
metrics["resolved_model_name_or_path"],
)
return metrics
def _normalize_request(self, request: SynthesisRequest) -> SynthesisRequest:
normalized_text = request.text.strip()
if not normalized_text:
raise ValueError("text 不能为空。")
normalized_prompt_audio_path = request.prompt_audio_path or None
normalized_prompt_text = (request.prompt_text or "").strip() or None
if normalized_prompt_text and not normalized_prompt_audio_path:
raise ValueError("prompt_text requires prompt_audio_path.")
normalized_template_name = request.template_name.strip() or "tts"
if normalized_template_name not in GRADIO_SYNTHESIS_MODE_TEMPLATE_NAMES:
raise ValueError(
f"Unsupported template_name={normalized_template_name!r}. "
f"Expected one of {list(GRADIO_SYNTHESIS_MODE_TEMPLATE_NAMES)}."
)
normalized_language = (request.language or "").strip() or None
supported_language_codes = set(SUPPORTED_LANGUAGE_CODE_BY_NAME.values())
if (
normalized_language is not None
and normalized_language not in supported_language_codes
):
raise ValueError(
f"Unsupported language={normalized_language!r}. "
f"Expected one of {sorted(supported_language_codes)}."
)
resolved_seed = int(request.seed)
return SynthesisRequest(
model_name_or_path=request.model_name_or_path.strip(),
text=normalized_text,
prompt_audio_path=normalized_prompt_audio_path,
prompt_text=normalized_prompt_text,
execution_mode=request.execution_mode,
template_name=normalized_template_name,
language=normalized_language,
ode_method=request.ode_method.strip() or DEFAULT_ODE_METHOD,
num_steps=int(request.num_steps),
guidance_scale=float(request.guidance_scale),
speaker_scale=float(request.speaker_scale),
normalize_text=bool(request.normalize_text),
seed=resolved_seed,
)
def _build_output_path(self) -> Path:
output_name = f"{time.strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:8]}.wav"
return self.config.output_dir / output_name
def _cleanup_outputs(self) -> None:
if self.config.output_retention_count <= 0:
return
wav_files = sorted(
self.config.output_dir.glob("*.wav"),
key=lambda path: path.stat().st_mtime,
reverse=True,
)
removed_count = 0
for stale_file in wav_files[self.config.output_retention_count :]:
stale_file.unlink(missing_ok=True)
removed_count += 1
if removed_count > 0:
logger.info(
"Gradio output cleanup completed: removed_files={} retention_limit={}",
removed_count,
self.config.output_retention_count,
)
@staticmethod
def _waveform_to_numpy(audio: torch.Tensor):
waveform = audio.detach().float().cpu().squeeze()
if waveform.ndim == 0:
raise ValueError("生成音频为空。")
return waveform.numpy()
def _write_audio(self, audio: torch.Tensor, sample_rate: int) -> str:
output_path = self._build_output_path()
logger.info(
"Writing synthesized audio: output_path={} sample_rate={} samples={}",
output_path,
sample_rate,
audio.shape[-1],
)
sf.write(output_path, self._waveform_to_numpy(audio), sample_rate)
self._cleanup_outputs()
logger.info("Synthesized audio written: output_path={}", output_path)
return str(output_path)
def generate(self, request: SynthesisRequest) -> SynthesisResult:
normalized_request = self._normalize_request(request)
with self._lock:
try:
seed_everything(normalized_request.seed)
runtime, resolved_model_name_or_path = self._get_runtime(
normalized_request.model_name_or_path,
)
logger.info(
"Gradio request accepted: resolved_model_name_or_path={} execution_mode={} seed={}",
resolved_model_name_or_path,
normalized_request.execution_mode,
normalized_request.seed,
)
if normalized_request.execution_mode == "generate_stream":
result = self._run_stream_generation(runtime, normalized_request)
else:
result = runtime.generate(
**self._build_runtime_generate_kwargs(normalized_request)
)
result["chunk_count"] = 1
audio_path = self._write_audio(result["audio"], result["sample_rate"])
except Exception:
logger.exception(
"Gradio request failed: model_name_or_path={} execution_mode={} text_len={} has_prompt_audio={} has_prompt_text={} template_name={} language={} "
"precision={} ode_method={} num_steps={} guidance_scale={} speaker_scale={} max_generate_length={} "
"normalize_text={} seed={}",
normalized_request.model_name_or_path,
normalized_request.execution_mode,
len(normalized_request.text),
bool(normalized_request.prompt_audio_path),
bool(normalized_request.prompt_text),
normalized_request.template_name,
normalized_request.language,
self.config.precision,
normalized_request.ode_method,
normalized_request.num_steps,
normalized_request.guidance_scale,
normalized_request.speaker_scale,
self.config.max_generate_length,
normalized_request.normalize_text,
normalized_request.seed,
)
raise
audio_seconds = result["audio"].shape[-1] / result["sample_rate"]
metrics = {
"request_id": result["fid"],
"execution_mode": normalized_request.execution_mode,
"chunk_count": int(result["chunk_count"]),
"template_name": normalized_request.template_name,
"language": normalized_request.language,
"resolved_model_name_or_path": resolved_model_name_or_path,
"sample_rate": result["sample_rate"],
"elapsed_seconds": round(float(result["time_used"]), 3),
"audio_seconds": round(float(audio_seconds), 3),
"rtf": round(float(result["rtf"]), 4),
"seed": normalized_request.seed,
"output_path": audio_path,
}
logger.info(
"Gradio request output ready: request_id={} execution_mode={} resolved_model_name_or_path={} output_path={}",
metrics["request_id"],
metrics["execution_mode"],
metrics["resolved_model_name_or_path"],
metrics["output_path"],
)
status = (
f"完成:{Path(audio_path).name} | "
f"模式 {metrics['execution_mode']} | "
f"耗时 {metrics['elapsed_seconds']}s | "
f"音频 {metrics['audio_seconds']}s | "
f"RTF {metrics['rtf']}"
)
return SynthesisResult(
audio_path=audio_path,
metrics=metrics,
status=status,
)
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 185 KiB

+76
View File
@@ -0,0 +1,76 @@
train_data:
train_audio_sample_rate: 48000
audio_samples_per_llm_token: 7680
sources:
- name: ljspeech_basic
weight: 1.0
pipeline: basic
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: downloaded_data/ljspeech_48khz_manifest_train.jsonl
shuffle: true
- name: ljspeech_interleave
weight: 1.0
pipeline: interleave
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: downloaded_data/ljspeech_48khz_manifest_train.jsonl
shuffle: true
# append other sources here if need
num_tokens_per_epoch: 2000000
num_workers: 20
pin_memory: true
max_audio_seconds_in_batch: 30.0
max_text_tokens_in_batch: 2048
max_samples_per_batch: null
bucketing_pool_size: 100
val_data:
train_audio_sample_rate: 48000
audio_samples_per_llm_token: 7680
sources:
- name: ljspeech_valid_basic
weight: 1.0
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: downloaded_data/ljspeech_48khz_manifest_valid.jsonl
shuffle: false
pipeline: basic
- name: ljspeech_valid_interleave
weight: 1.0
pipeline: interleave
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: downloaded_data/ljspeech_48khz_manifest_valid.jsonl
shuffle: false
pipeline: interleave
# append other sources here if need
num_workers: 4
pin_memory: true
max_audio_seconds_in_batch: 30.0
max_text_tokens_in_batch: 2048
max_samples_per_batch: null
bucketing_pool_size: 64
train:
pretrained_model_path: pretrained_models/pretrain_cpt_decay/latest/model/
output_dir: debug_train/run_003
seed: 42
learning_rate: 1.0e-05
weight_decay: 0.01
warmup_steps: 50
max_train_steps: 500
gradient_accumulation_steps: 2
grad_clip_norm: 1
save_interval: 500
max_checkpoints_to_keep: 40
log_interval: 10
eval_interval: 100
max_eval_batches: null
run_eval_on_start: false
loss:
ce_weight: 1.0
fm_weight: 1.0
eos_weight: 1.0
+61
View File
@@ -0,0 +1,61 @@
train_data:
train_audio_sample_rate: 48000
audio_samples_per_llm_token: 7680
sources:
- name: ljspeech_meanflow_basic
weight: 1.0
pipeline: basic
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: downloaded_data/ljspeech_48khz_manifest_train.jsonl
shuffle: true
# append other sources here if need
num_tokens_per_epoch: 1000000
num_workers: 4
pin_memory: true
max_audio_seconds_in_batch: 10.0
max_text_tokens_in_batch: 1024
max_samples_per_batch: 1
bucketing_pool_size: 64
val_data:
train_audio_sample_rate: 48000
audio_samples_per_llm_token: 7680
sources:
- name: ljspeech_meanflow_valid_basic
weight: 1.0
pipeline: basic
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: downloaded_data/ljspeech_48khz_manifest_valid.jsonl
shuffle: false
# append other sources here if need
num_workers: 2
pin_memory: true
max_audio_seconds_in_batch: 10.0
max_text_tokens_in_batch: 1024
max_samples_per_batch: 1
bucketing_pool_size: 32
train:
pretrained_model_path: pretrained_models/dots.tts-soar
output_dir: debug_train/meanflow_run_001
seed: 42
learning_rate: 1.0e-05
cfg_droprate: 0.0
xvec_drop_rate: 0.5
weight_decay: 0.01
warmup_steps: 50
max_train_steps: 500
gradient_accumulation_steps: 1
grad_clip_norm: 1
save_interval: 500
max_checkpoints_to_keep: 10
log_interval: 10
eval_interval: 100
max_eval_batches: 10
run_eval_on_start: false
loss:
ce_weight: 1.0
fm_weight: 1.0
eos_weight: 1.0
+18
View File
@@ -0,0 +1,18 @@
# Recommended versions for reproducible installs.
# pyproject.toml keeps compatibility ranges; install with `pip -c` to use these
# versions by default while still allowing users to omit this file for newer
# compatible releases.
torch==2.8.0
torchaudio==2.8.0
transformers==4.57.0
librosa==0.11.0
soundfile==0.13.1
numpy==2.2.6
pydantic==2.12.5
PyYAML==6.0.3
safetensors==0.8.0rc0
# Optional `full` extra.
accelerate==1.12.0
tensorboard==2.20.0
ruff==0.15.12
+108
View File
@@ -0,0 +1,108 @@
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "dots.tts"
version = "0.1.0"
description = "dots.tts: a fully continuous autoregressive TTS system with self-corrective alignment and CFG-aware MeanFlow distillation."
readme = "README.md"
requires-python = ">=3.10"
license = { text = "Apache-2.0" }
authors = [{ name = "dots.tts Team" }]
keywords = [
"text-to-speech",
"tts",
"speech-synthesis",
"autoregressive",
"flow-matching",
"meanflow",
"diffusion",
"voice-cloning",
]
classifiers = [
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: Apache Software License",
"Operating System :: POSIX :: Linux",
"Operating System :: MacOS",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Topic :: Multimedia :: Sound/Audio :: Speech",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"torch>=2.8.0",
"torchaudio>=2.8.0",
"transformers>=4.57.0",
"huggingface-hub",
"loguru",
"langcodes[data]",
"gradio",
"einops",
"librosa>=0.11.0",
"soundfile>=0.13.1",
"numpy>=2.2.6",
"pydantic>=2.12.5,<3",
"PyYAML>=6.0.3",
"safetensors>=0.8.0rc0",
"torchdiffeq",
"tqdm",
"lingua-language-detector",
# WeTextProcessing disabled: requires pynini (needs MSVC build tools on Windows)
# "WeTextProcessing",
]
[project.optional-dependencies]
full = [
"accelerate>=1.12.0",
"tensorboard>=2.20.0",
"ruff>=0.15.12",
]
[project.scripts]
"dots.tts" = "dots_tts.cli:main"
[project.urls]
# TODO: fill in before public launch
Homepage = "https://github.com/<to-be-filled>/dots.tts"
Repository = "https://github.com/<to-be-filled>/dots.tts"
Issues = "https://github.com/<to-be-filled>/dots.tts/issues"
[tool.setuptools]
include-package-data = false
[tool.setuptools.package-dir]
"" = "src"
[tool.setuptools.packages.find]
where = ["src"]
include = ["dots_tts*"]
[tool.ruff]
target-version = "py310"
line-length = 88
src = ["src"]
extend-exclude = ["__pycache__"]
[tool.ruff.lint]
select = [
"E",
"F",
"I",
]
ignore = [
"E501",
]
[tool.ruff.lint.isort]
known-first-party = ["dots_tts"]
[tool.ruff.format]
quote-style = "double"
indent-style = "space"
skip-magic-trailing-comma = false
line-ending = "lf"
+166
View File
@@ -0,0 +1,166 @@
from __future__ import annotations
import argparse
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
SRC_ROOT = REPO_ROOT / "src"
for import_root in (REPO_ROOT, SRC_ROOT):
import_root_str = str(import_root)
if import_root_str not in sys.path:
sys.path.insert(0, import_root_str)
import soundfile as sf # noqa: E402
import torch # noqa: E402
from loguru import logger # noqa: E402
from dots_tts.utils.logging import configure_logging # noqa: E402
from dots_tts.runtime_double_streaming import ( # noqa: E402
DotsTtsRuntimeDoubleStreaming,
)
from dots_tts.utils.text import normalize_text # noqa: E402
from dots_tts.utils.util import seed_everything # noqa: E402
def parse_args(argv=None):
parser = argparse.ArgumentParser(
description="Temporary example for dots.tts double streaming session API."
)
parser.add_argument(
"--model-name-or-path",
required=True,
help="Local pretrained directory or Hugging Face repo id",
)
parser.add_argument("--text", required=True, help="Input text")
parser.add_argument("--output", default="double_streaming.wav", help="Output wav path")
parser.add_argument(
"--prompt-audio",
default=None,
help="Optional reference audio for ref_audio_only speaker conditioning",
)
parser.add_argument("--revision", default=None, help="Optional Hugging Face revision")
parser.add_argument("--cache-dir", default=None, help="Optional Hugging Face cache dir")
parser.add_argument("--precision", default="bfloat16", help="Inference precision")
parser.add_argument(
"--optimize",
action="store_true",
help="Enable inference optimization and warmup",
)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Random seed.",
)
parser.add_argument("--ode-method", default="euler", help="ODE solver method")
parser.add_argument("--num-steps", type=int, default=10, help="Diffusion sampling steps")
parser.add_argument(
"--guidance-scale",
type=float,
default=1.2,
help="Classifier-free guidance scale",
)
parser.add_argument(
"--eos-threshold",
type=float,
default=0.8,
help="EOS stop threshold for finish_text() tail decode",
)
parser.add_argument(
"--max-generate-length",
type=int,
default=500,
help="Maximum number of decoded audio patches in double streaming",
)
parser.add_argument(
"--normalize-text",
action="store_true",
help="Normalize text before tokenizer encode",
)
return parser.parse_args(argv)
def _prepare_text(text: str, *, normalize: bool) -> str:
prepared = text.strip()
if normalize:
prepared = normalize_text(prepared)
if not prepared:
raise ValueError("Input text is empty after preprocessing.")
return prepared
def main(argv=None):
configure_logging()
args = parse_args(argv)
seed_everything(args.seed)
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
runtime = DotsTtsRuntimeDoubleStreaming.from_pretrained(
args.model_name_or_path,
revision=args.revision,
cache_dir=args.cache_dir,
precision=args.precision,
optimize=args.optimize,
max_generate_length=args.max_generate_length,
)
prepared_text = _prepare_text(args.text, normalize=args.normalize_text)
text_token_ids = runtime.model.tokenizer.encode(
prepared_text,
add_special_tokens=False,
)
if not text_token_ids:
raise ValueError("Tokenizer produced no text tokens.")
logger.info(
"Double streaming example started: text_len={} text_token_count={} output={}",
len(prepared_text),
len(text_token_ids),
output_path,
)
session = runtime.start_double_streaming(
prompt_audio_path=args.prompt_audio,
ode_method=args.ode_method,
num_steps=args.num_steps,
guidance_scale=args.guidance_scale,
eos_threshold=args.eos_threshold,
)
chunks: list[torch.Tensor] = []
for index, token_id in enumerate(text_token_ids, start=1):
chunk = session.push_text_token(token_id)
logger.info(
"Double streaming step: token_index={} token_id={} emitted_audio={}",
index,
token_id,
chunk is not None,
)
if chunk is not None:
chunks.append(chunk.detach().cpu())
for chunk in session.finish_text():
chunks.append(chunk.detach().cpu())
if not chunks:
raise RuntimeError("Double streaming produced no audio chunks.")
audio = torch.cat(chunks, dim=-1)
sf.write(
output_path,
audio.float().squeeze().numpy(),
runtime.sample_rate,
)
logger.info(
"Double streaming example completed: output={} chunk_count={} samples={}",
output_path,
len(chunks),
audio.shape[-1],
)
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,134 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import csv
import json
import subprocess
from pathlib import Path
from huggingface_hub import hf_hub_download
REPO_ROOT = Path(__file__).resolve().parents[1]
REPO_ID = "alibabasglab/LJSpeech-1.1-48kHz"
ARCHIVE_NAME = "LJSpeech-1.1-48kHz.tar.bz2"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument(
"--cache-dir",
type=Path,
default=REPO_ROOT / "downloaded_data" / "hf_cache",
)
parser.add_argument(
"--extract-dir",
type=Path,
default=REPO_ROOT / "downloaded_data" / "hf_cache",
)
parser.add_argument(
"--output-dir",
type=Path,
default=REPO_ROOT / "downloaded_data",
)
parser.add_argument("--valid-size", type=int, default=100)
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.valid_size < 0:
raise ValueError("valid_size must be >= 0")
cache_dir = args.cache_dir.resolve()
extract_dir = args.extract_dir.resolve()
output_dir = args.output_dir.resolve()
train_manifest_path = output_dir / "ljspeech_48khz_manifest_train.jsonl"
valid_manifest_path = output_dir / "ljspeech_48khz_manifest_valid.jsonl"
cache_dir.mkdir(parents=True, exist_ok=True)
extract_dir.mkdir(parents=True, exist_ok=True)
output_dir.mkdir(parents=True, exist_ok=True)
archive_path = Path(
hf_hub_download(
repo_id=REPO_ID,
repo_type="dataset",
filename=ARCHIVE_NAME,
local_dir=str(cache_dir),
)
)
dataset_root = extract_dir / "LJSpeech-1.1-48kHz"
if not dataset_root.exists():
print("extracting archive...")
subprocess.run(
[
"tar",
"-xjf",
str(archive_path),
"-C",
str(extract_dir),
"--checkpoint=2000",
"--checkpoint-action=echo=extracting...",
],
check=True,
)
metadata_path = dataset_root / "metadata.csv"
audio_dir = dataset_root / "wavs" / "MossFormer2_SR_48K"
if not metadata_path.is_file():
raise FileNotFoundError(f"metadata.csv not found: {metadata_path}")
if not audio_dir.is_dir():
raise FileNotFoundError(f"audio dir not found: {audio_dir}")
train_count = 0
valid_count = 0
with (
metadata_path.open("r", encoding="utf-8", newline="") as fin,
train_manifest_path.open("w", encoding="utf-8") as train_fout,
valid_manifest_path.open("w", encoding="utf-8") as valid_fout,
):
reader = csv.reader(fin, delimiter="|")
for row in reader:
if not row:
continue
fid = row[0].strip()
text = (
row[2].strip() if len(row) >= 3 and row[2].strip() else row[1].strip()
)
audio_path = (audio_dir / f"{fid}.wav").resolve()
if not audio_path.is_file():
raise FileNotFoundError(f"audio not found: {audio_path}")
record = json.dumps(
{
"fid": fid,
"audio": str(audio_path),
"text": text,
},
ensure_ascii=False,
)
if valid_count < args.valid_size:
valid_fout.write(record)
valid_fout.write("\n")
valid_count += 1
else:
train_fout.write(record)
train_fout.write("\n")
train_count += 1
print(f"archive: {archive_path}")
print(f"dataset_root: {dataset_root}")
print(f"train_manifest: {train_manifest_path}")
print(f"valid_manifest: {valid_manifest_path}")
print(f"train_records: {train_count}")
print(f"valid_records: {valid_count}")
if __name__ == "__main__":
main()
+773
View File
@@ -0,0 +1,773 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import math
import time
from dataclasses import dataclass
from pathlib import Path
import torch
import yaml
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration
from torch.optim import AdamW
from transformers import get_cosine_schedule_with_warmup
from dots_tts.config import app as app_config
from dots_tts.data import builders as data_module
from dots_tts.models.dots_tts import model as dots_tts_model
from dots_tts.training import checkpoint as train_checkpoint
from dots_tts.training import losses as loss_ops
from dots_tts.training import utils as train_utils
from dots_tts.utils import util as util_module
_EMPTY_EPOCH_TOLERANCE = 32
_DEBUG_BATCH_LIMIT = 3
_DEBUG_GRAD_EARLY_STEP_LIMIT = 3
# region Training Step State
@dataclass(slots=True)
class _PreparedTrainingStep:
micro_batches: list[dict]
consumed_counts: list[int]
global_denominators: dict[str, float]
@dataclass(slots=True)
class _AccumulatedTrainingStep:
loss_totals: dict[str, float]
loss_denominators: dict[str, float]
source_loss_totals: dict[str, dict[str, float]]
source_loss_denominators: dict[str, dict[str, float]]
completed_optimizer_step: bool
grad_norm: torch.Tensor | None
@dataclass(slots=True)
class _CompletedTrainingStep:
reduced_metrics: dict[str, float]
learning_rate: float
grad_norm_value: float
# endregion Training Step State
class DotsTtsTrainingRun:
# region Lifecycle
def __init__(self, cfg: app_config.AppConfig, *, debug_enabled: bool = False):
self.cfg = cfg
self.progress = train_utils.TrainProgress()
self.max_train_steps = int(cfg.train.max_train_steps)
self.grad_accumulation_steps = int(cfg.train.gradient_accumulation_steps)
self.last_validation_step: int | None = None
self.consecutive_empty_epochs = 0
self.saved_latest_checkpoint = False
self._last_log_step = 0
self._last_log_time = 0.0
self._debug_enabled = bool(debug_enabled)
self._debug_batch_count = 0
self._debug_audio_sample_rate = int(self.cfg.train_data.train_audio_sample_rate)
project_config = ProjectConfiguration(
project_dir=self.cfg.train.output_dir,
total_limit=self.cfg.train.max_checkpoints_to_keep,
)
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
self.accelerator = Accelerator(
kwargs_handlers=[ddp_kwargs],
gradient_accumulation_steps=self.grad_accumulation_steps,
log_with="tensorboard",
project_config=project_config,
step_scheduler_with_optimizer=False,
)
util_module.seed_everything(self.cfg.train.seed)
model = dots_tts_model.DotsTtsModel.from_pretrained(
self.cfg.train.pretrained_model_path
)
# model.set_cfg_droprate(
# cfg_droprate=self.cfg.train.cfg_droprate,
# xvec_drop_rate=self.cfg.train.xvec_drop_rate,
# )
optimizer = AdamW(
(param for param in model.parameters() if param.requires_grad),
lr=self.cfg.train.learning_rate,
weight_decay=self.cfg.train.weight_decay,
)
scheduler = get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps=self.cfg.train.warmup_steps,
num_training_steps=self.max_train_steps,
)
self.model, self.optimizer, self.scheduler = self.accelerator.prepare(
model,
optimizer,
scheduler,
)
self.unwrapped_model = self.accelerator.unwrap_model(self.model)
expected_sample_rate = int(self.unwrapped_model.config.vocoder.sample_rate)
expected_audio_samples_per_llm_token = (
int(self.unwrapped_model.hop_size) * int(self.unwrapped_model.config.patch_size)
)
if int(self.cfg.train_data.train_audio_sample_rate) != expected_sample_rate:
raise ValueError(
f"train_data.train_audio_sample_rate={int(self.cfg.train_data.train_audio_sample_rate)} "
f"does not match the pretrained model sample rate {expected_sample_rate}."
)
if (
int(self.cfg.train_data.audio_samples_per_llm_token)
!= expected_audio_samples_per_llm_token
):
raise ValueError(
"train_data.audio_samples_per_llm_token="
f"{int(self.cfg.train_data.audio_samples_per_llm_token)} "
"does not match the pretrained model audio token contract "
f"{expected_audio_samples_per_llm_token}."
)
if self.cfg.val_data is not None:
if int(self.cfg.val_data.train_audio_sample_rate) != expected_sample_rate:
raise ValueError(
f"val_data.train_audio_sample_rate={int(self.cfg.val_data.train_audio_sample_rate)} "
f"does not match the pretrained model sample rate {expected_sample_rate}."
)
if (
int(self.cfg.val_data.audio_samples_per_llm_token)
!= expected_audio_samples_per_llm_token
):
raise ValueError(
"val_data.audio_samples_per_llm_token="
f"{int(self.cfg.val_data.audio_samples_per_llm_token)} "
"does not match the pretrained model audio token contract "
f"{expected_audio_samples_per_llm_token}."
)
if self.accelerator.is_main_process:
total_params = sum(param.numel() for param in self.unwrapped_model.parameters())
trainable_params = sum(
param.numel()
for param in self.unwrapped_model.parameters()
if param.requires_grad
)
self.accelerator.print(f"Total parameters: {total_params:,}")
self.accelerator.print(f"Trainable parameters: {trainable_params:,}")
self.accelerator.print(
f"Distributed type: {self.accelerator.distributed_type}"
)
tokenizer = self.unwrapped_model.tokenizer
self.tokenizer = tokenizer
train_dataset = data_module.build_training_dataset(
self.cfg.train_data,
tokenizer=tokenizer,
seed=int(self.cfg.train.seed),
accelerator=self.accelerator,
)
self.train_loader = data_module.build_training_dataloader(
train_dataset,
self.cfg.train_data,
tokenizer=tokenizer,
)
self.val_loader = None
if (
self.cfg.train.eval_interval is not None
or self.cfg.train.run_eval_on_start
):
if self.cfg.val_data is None:
raise ValueError(
"Validation requires val_data when eval_interval or "
"run_eval_on_start is enabled."
)
validation_data_cfg = self.cfg.val_data.model_copy(deep=True)
validation_data_cfg.num_tokens_per_epoch = None
val_dataset = data_module.build_validation_dataset(
validation_data_cfg,
tokenizer=tokenizer,
seed=int(self.cfg.train.seed),
accelerator=self.accelerator,
)
self.val_loader = data_module.build_validation_dataloader(
val_dataset,
validation_data_cfg,
tokenizer=tokenizer,
)
self._resume_if_available()
self.train_loader.set_epoch(self.progress.epoch)
def run(self) -> int:
self.accelerator.init_trackers("dots_tts")
self._write_run_config()
self.optimizer.zero_grad(set_to_none=True)
try:
if self.cfg.train.run_eval_on_start:
self._run_validation()
self.last_validation_step = self.progress.global_step
self._last_log_step = self.progress.global_step
self._last_log_time = time.perf_counter()
while self.progress.global_step < self.max_train_steps:
self._run_training_step()
if (
self.cfg.train.eval_interval is not None
and self.val_loader is not None
and self.progress.global_step > 0
and self.last_validation_step != self.progress.global_step
):
self._run_validation()
if not self.saved_latest_checkpoint:
self._save_checkpoint(float(self.optimizer.param_groups[0]["lr"]))
return 0
finally:
try:
self._close_data_streams()
finally:
self.accelerator.end_training()
def _write_run_config(self) -> None:
if not bool(getattr(self.accelerator, "is_main_process", True)):
return
output_dir = Path(self.cfg.train.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
config_path = output_dir / "config.yml"
with config_path.open("w", encoding="utf-8") as fout:
yaml.safe_dump(
self.cfg.to_dict(),
fout,
sort_keys=False,
allow_unicode=True,
)
def _close_data_streams(self) -> None:
for loader_name in ("train_loader", "val_loader"):
loader = getattr(self, loader_name, None)
close = getattr(loader, "close", None)
if callable(close):
close()
setattr(self, loader_name, None)
def _resume_if_available(self) -> None:
try:
resume_dir = train_checkpoint.resolve_latest_train_checkpoint(
self.cfg.train.output_dir
)
except FileNotFoundError:
return
resume_state = train_checkpoint.load_train_checkpoint(
self.accelerator,
self.model,
self.optimizer,
self.progress,
resume_dir,
self.scheduler,
)
saved_max_train_steps = int(resume_state["scheduler_state"]["max_train_steps"])
if saved_max_train_steps != self.max_train_steps:
self.accelerator.print(
"Warning: resumed scheduler was saved with "
f"max_train_steps={saved_max_train_steps}, but current run uses "
f"{self.max_train_steps}."
)
self.train_loader.load_state_dict(resume_state["data_state"])
self.accelerator.print(
"Resumed training from "
f"{resume_dir} at step {self.progress.global_step}. "
"Restored committed data state. "
"In-memory prefetch and batching state is rebuilt on restart, so only "
"committed sample progress is resumed."
)
# endregion Lifecycle
# region Training Step Pipeline
def _run_training_step(self) -> None:
try:
self.model.train()
# Stage 1: collect one synchronized accumulation window and its
# normalization factors before touching model state.
prepared_step = self._prepare_training_step()
# Stage 2: run forward/backward over the prepared micro-batches and
# accumulate overall/source statistics for the completed optimizer step.
accumulated_step = self._accumulate_training_step(prepared_step)
# Stage 3: advance counters, reduce metrics, then trigger side effects
# (logging, validation, checkpointing) only after a real optimizer step.
self._apply_consumed_counts(prepared_step.consumed_counts)
if not accumulated_step.completed_optimizer_step:
return
completed_step = self._finalize_completed_training_step(accumulated_step)
if train_utils.should_log_training_step(
self.progress.global_step,
int(self.cfg.train.log_interval),
):
reduced_by_source = train_utils.reduce_source_metrics(
accumulated_step.source_loss_totals,
accumulated_step.source_loss_denominators,
device=self.accelerator.device,
loss_config=self.cfg.loss,
)
current_time = time.perf_counter()
report = train_utils.build_train_step_report(
completed_step.reduced_metrics,
learning_rate=completed_step.learning_rate,
grad_norm=completed_step.grad_norm_value,
current_time=current_time,
last_log_step=self._last_log_step,
last_log_time=self._last_log_time,
progress=self.progress,
max_train_steps=self.max_train_steps,
reduced_by_source=reduced_by_source,
)
self.accelerator.log(
report.log_values,
step=self.progress.global_step,
)
self.accelerator.print(report.console_line)
self._last_log_step = self.progress.global_step
self._last_log_time = current_time
if (
self.cfg.train.eval_interval is not None
and self.progress.global_step % self.cfg.train.eval_interval == 0
):
self._run_validation()
self.last_validation_step = self.progress.global_step
if self.progress.global_step % self.cfg.train.save_interval == 0:
self._save_checkpoint(completed_step.learning_rate)
self.saved_latest_checkpoint = True
except BaseException as exc:
train_utils.abort_on_out_of_memory(
exc,
stage="train",
batch=None,
progress=self.progress,
device=self.accelerator.device,
process_index=int(getattr(self.accelerator, "process_index", 0)),
num_processes=int(getattr(self.accelerator, "num_processes", 1)),
)
raise
def _prepare_training_step(self) -> _PreparedTrainingStep:
micro_batches: list[dict] = []
local_denominators: dict[str, float] = {}
while len(micro_batches) < self.grad_accumulation_steps:
batch, has_batch = self.train_loader.peek_batch()
if train_utils.any_rank_true(not has_batch, device=self.accelerator.device):
self._advance_epoch_after_empty_batch(has_local_batch=has_batch)
continue
self.consecutive_empty_epochs = 0
self.train_loader.commit_batch()
prepared_batch = self.unwrapped_model.prepare_training_batch(batch)
self._maybe_debug_training_batch(prepared_batch)
batch_denominators = loss_ops.to_host_named_scalars(
loss_ops.collapse_loss_masks(prepared_batch["loss_masks"])
)
if not local_denominators:
local_denominators = {name: 0.0 for name in batch_denominators}
loss_ops.accumulate_named_scalars_(local_denominators, batch_denominators)
micro_batches.append(prepared_batch)
consumed_counts = train_utils.sum_integer_counters_across_ranks(
[
sum(int(batch["input_ids_lengths"].sum().item()) for batch in micro_batches),
sum(int(batch["num_audio_tokens"].sum().item()) for batch in micro_batches),
sum(int(batch["num_text_tokens"].sum().item()) for batch in micro_batches),
],
device=self.accelerator.device,
)
global_denominators = loss_ops.sum_named_scalars_across_ranks(
local_denominators,
device=self.accelerator.device,
)
return _PreparedTrainingStep(
micro_batches=micro_batches,
consumed_counts=consumed_counts,
global_denominators=global_denominators,
)
def _advance_epoch_after_empty_batch(self, *, has_local_batch: bool) -> None:
if has_local_batch:
self.train_loader.discard_batch()
self.progress.epoch += 1
self.train_loader.set_epoch(self.progress.epoch)
self.consecutive_empty_epochs += 1
if self.consecutive_empty_epochs > _EMPTY_EPOCH_TOLERANCE:
raise RuntimeError(
"Unable to obtain a synchronized training batch across ranks. "
"Check shard assignment, dataset size, and filtering constraints."
)
def _accumulate_training_step(
self,
prepared_step: _PreparedTrainingStep,
) -> _AccumulatedTrainingStep:
accumulated_loss_totals: dict[str, float] = {}
accumulated_loss_denominators: dict[str, float] = {}
accumulated_source_loss_totals: dict[str, dict[str, float]] = {}
accumulated_source_loss_denominators: dict[str, dict[str, float]] = {}
completed_optimizer_step = False
grad_norm = None
for batch in prepared_step.micro_batches:
batch = train_utils.move_to_device(batch, self.accelerator.device)
with self.accelerator.accumulate(self.model):
with self.accelerator.autocast():
loss_terms = self.model(batch)
loss = loss_ops.compute_gradient_loss(
loss_terms,
global_normalizers=prepared_step.global_denominators,
loss_config=self.cfg.loss,
ddp_world_size=int(self.accelerator.num_processes),
gradient_accumulation_steps=self.grad_accumulation_steps,
)
batch_loss_totals, batch_loss_denominators = (
loss_ops.collapse_loss_terms(loss_terms)
)
batch_loss_totals = loss_ops.to_host_named_scalars(batch_loss_totals)
batch_loss_denominators = loss_ops.to_host_named_scalars(
batch_loss_denominators
)
if not accumulated_loss_totals:
accumulated_loss_totals = {name: 0.0 for name in batch_loss_totals}
accumulated_loss_denominators = {
name: 0.0 for name in batch_loss_denominators
}
loss_ops.accumulate_named_scalars_(
accumulated_loss_totals,
batch_loss_totals,
)
loss_ops.accumulate_named_scalars_(
accumulated_loss_denominators,
batch_loss_denominators,
)
batch_source_totals, batch_source_denominators = (
loss_ops.collapse_loss_terms_by_source(
loss_terms,
source_names=batch["source_names"],
)
)
loss_ops.accumulate_grouped_named_scalars_(
accumulated_source_loss_totals,
batch_source_totals,
)
loss_ops.accumulate_grouped_named_scalars_(
accumulated_source_loss_denominators,
batch_source_denominators,
)
self.accelerator.backward(loss)
if self.accelerator.sync_gradients:
grad_norm = self.accelerator.clip_grad_norm_(
self.model.parameters(),
self.cfg.train.grad_clip_norm,
)
self._maybe_print_gradient_debug(grad_norm)
self.optimizer.step()
completed_optimizer_step = (
not self.accelerator.optimizer_step_was_skipped
)
if completed_optimizer_step:
self.scheduler.step()
self.optimizer.zero_grad(set_to_none=True)
batch.clear()
return _AccumulatedTrainingStep(
loss_totals=accumulated_loss_totals,
loss_denominators=accumulated_loss_denominators,
source_loss_totals=accumulated_source_loss_totals,
source_loss_denominators=accumulated_source_loss_denominators,
completed_optimizer_step=completed_optimizer_step,
grad_norm=grad_norm,
)
def _apply_consumed_counts(self, consumed_counts: list[int]) -> None:
self.progress.total_tokens += consumed_counts[0]
self.progress.audio_tokens += consumed_counts[1]
self.progress.text_tokens += consumed_counts[2]
def _finalize_completed_training_step(
self,
accumulated_step: _AccumulatedTrainingStep,
) -> _CompletedTrainingStep:
if not accumulated_step.loss_totals or not accumulated_step.loss_denominators:
raise RuntimeError("Training step produced no accumulated loss totals.")
if all(
float(value) == 0.0 for value in accumulated_step.loss_denominators.values()
):
raise RuntimeError("Accumulated training step produced no loss statistics.")
self.progress.global_step += 1
self.saved_latest_checkpoint = False
reduced_totals = loss_ops.sum_named_scalars_across_ranks(
accumulated_step.loss_totals,
device=self.accelerator.device,
)
reduced_denominators = loss_ops.sum_named_scalars_across_ranks(
accumulated_step.loss_denominators,
device=self.accelerator.device,
)
reduced_metrics = loss_ops.reduce_loss_statistics(
reduced_totals,
reduced_denominators,
loss_config=self.cfg.loss,
)
learning_rate = float(self.optimizer.param_groups[0]["lr"])
grad_norm_value = (
math.nan
if accumulated_step.grad_norm is None
else float(accumulated_step.grad_norm.detach().float().item())
)
return _CompletedTrainingStep(
reduced_metrics=reduced_metrics,
learning_rate=learning_rate,
grad_norm_value=grad_norm_value,
)
# endregion Training Step Pipeline
# region Validation
def _run_validation(self) -> None:
try:
if self.val_loader is None:
raise ValueError(
"Validation requested, but validation loader was not initialized."
)
self.val_loader.set_epoch(0)
was_training = bool(self.model.training)
self.model.eval()
overall_loss_totals = None
overall_loss_denominators = None
source_loss_totals: dict[str, dict[str, float]] = {}
source_loss_denominators: dict[str, dict[str, float]] = {}
processed_batches = 0
# Collect rank-local partial sums using the same batch preparation and
# loss aggregation path as training.
with torch.no_grad():
for batch_idx, batch in enumerate(self.val_loader):
if (
self.cfg.train.max_eval_batches is not None
and batch_idx >= self.cfg.train.max_eval_batches
):
break
batch = self.unwrapped_model.prepare_training_batch(batch)
batch = train_utils.move_to_device(batch, self.accelerator.device)
with self.accelerator.autocast():
loss_terms = self.model(batch)
batch_loss_totals, batch_loss_denominators = (
loss_ops.collapse_loss_terms(loss_terms)
)
batch_loss_totals = loss_ops.to_host_named_scalars(batch_loss_totals)
batch_loss_denominators = loss_ops.to_host_named_scalars(
batch_loss_denominators
)
if overall_loss_totals is None:
overall_loss_totals = {name: 0.0 for name in batch_loss_totals}
overall_loss_denominators = {
name: 0.0 for name in batch_loss_denominators
}
loss_ops.accumulate_named_scalars_(
overall_loss_totals,
batch_loss_totals,
)
loss_ops.accumulate_named_scalars_(
overall_loss_denominators,
batch_loss_denominators,
)
batch_source_totals, batch_source_denominators = (
loss_ops.collapse_loss_terms_by_source(
loss_terms,
source_names=batch["source_names"],
)
)
loss_ops.accumulate_grouped_named_scalars_(
source_loss_totals,
batch_source_totals,
)
loss_ops.accumulate_grouped_named_scalars_(
source_loss_denominators,
batch_source_denominators,
)
processed_batches += 1
# Merge rank-local partial sums with tensor reductions only. Validation
# runs close to the training memory ceiling, so object collectives are
# not acceptable here because NCCL materializes pickled payloads on GPU.
processed_batches = train_utils.sum_integer_counters_across_ranks(
[processed_batches],
device=self.accelerator.device,
)[0]
overall_loss_totals = loss_ops.sum_named_scalars_across_ranks(
overall_loss_totals or {},
device=self.accelerator.device,
)
overall_loss_denominators = loss_ops.sum_named_scalars_across_ranks(
overall_loss_denominators or {},
device=self.accelerator.device,
)
source_loss_totals = loss_ops.sum_grouped_named_scalars_across_ranks(
source_loss_totals,
device=self.accelerator.device,
)
source_loss_denominators = (
loss_ops.sum_grouped_named_scalars_across_ranks(
source_loss_denominators,
device=self.accelerator.device,
)
)
if processed_batches <= 0:
raise RuntimeError(
"Validation produced no batches. Check validation data configuration."
)
if not overall_loss_totals or not overall_loss_denominators:
raise RuntimeError("Validation produced no aggregate loss totals.")
reduced_metrics = loss_ops.reduce_loss_statistics(
overall_loss_totals,
overall_loss_denominators,
loss_config=self.cfg.loss,
)
reduced_by_source = loss_ops.reduce_loss_statistics_by_source(
source_loss_totals,
source_loss_denominators,
loss_config=self.cfg.loss,
)
if was_training:
self.model.train()
self.accelerator.log(
train_utils.build_validation_log_dict(
reduced_metrics,
reduced_by_source=reduced_by_source,
),
step=self.progress.global_step,
)
self.accelerator.print(
train_utils.format_validation_line(
reduced_metrics,
global_step=self.progress.global_step,
reduced_by_source=reduced_by_source,
)
)
except BaseException as exc:
train_utils.abort_on_out_of_memory(
exc,
stage="validation",
batch=None,
progress=self.progress,
device=self.accelerator.device,
process_index=int(getattr(self.accelerator, "process_index", 0)),
num_processes=int(getattr(self.accelerator, "num_processes", 1)),
)
raise
# endregion Validation
# region Checkpointing
def _save_checkpoint(self, learning_rate: float) -> None:
train_checkpoint.save_train_checkpoint(
self.accelerator,
self.model,
self.optimizer,
self.progress,
self.cfg.train.output_dir,
self.cfg.train.max_checkpoints_to_keep,
self.train_loader.state_dict(),
{
"type": "transformers_cosine_with_warmup",
"global_step": int(self.progress.global_step),
"base_lr": float(self.cfg.train.learning_rate),
"current_lr": float(learning_rate),
"warmup_steps": int(self.cfg.train.warmup_steps),
"max_train_steps": int(self.max_train_steps),
"state_dict": self.scheduler.state_dict(),
},
)
# endregion Checkpointing
# region Debug Logging
def _maybe_debug_training_batch(self, batch: dict[str, object]) -> None:
if not bool(getattr(self, "_debug_enabled", False)):
return
if not bool(getattr(self.accelerator, "is_main_process", True)):
return
if self._debug_batch_count >= _DEBUG_BATCH_LIMIT:
return
batch_index = self._debug_batch_count
self._debug_batch_count += 1
for line in train_utils.build_data_debug_lines(
batch,
batch_index=batch_index,
tokenizer=self.tokenizer,
sample_rate=self._debug_audio_sample_rate,
):
self.accelerator.print(line)
def _maybe_print_gradient_debug(self, grad_norm: torch.Tensor | None) -> None:
if grad_norm is None:
return
if not train_utils.should_print_gradient_debug(
debug_enabled=bool(getattr(self, "_debug_enabled", False)),
is_main_process=bool(getattr(self.accelerator, "is_main_process", True)),
next_global_step=self.progress.global_step + 1,
log_interval=int(self.cfg.train.log_interval),
early_step_limit=_DEBUG_GRAD_EARLY_STEP_LIMIT,
):
return
for line in train_utils.build_gradient_debug_lines(
self.unwrapped_model,
global_step=self.progress.global_step + 1,
grad_norm=float(grad_norm.detach().float().item()),
grad_clip_norm=float(self.cfg.train.grad_clip_norm),
):
self.accelerator.print(line)
# endregion Debug Logging
# region CLI
def parse_args(argv=None):
parser = argparse.ArgumentParser(
description="Accelerate training entrypoint for dots.tts."
)
parser.add_argument("--config", default=app_config.DEFAULT_CONFIG_PATH)
parser.add_argument(
"--debug",
action="store_true",
help="Print training debug information.",
)
return parser.parse_args(argv)
def main(argv=None):
args = parse_args(argv)
return DotsTtsTrainingRun(
app_config.load_config(args.config),
debug_enabled=args.debug,
).run()
if __name__ == "__main__":
raise SystemExit(main())
# endregion CLI
+956
View File
@@ -0,0 +1,956 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import torch
import torch.nn as nn
import yaml
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration
from einops import rearrange
from torch.optim import AdamW
from train_dots_tts import DotsTtsTrainingRun
from transformers import get_cosine_schedule_with_warmup
from dots_tts.config import app as app_config
from dots_tts.data import builders as data_module
from dots_tts.models.dots_tts import model as dots_tts_model
from dots_tts.models.dots_tts.config import MeanFlowConfig
from dots_tts.models.dots_tts.core import DotsTtsForwardOutput
from dots_tts.modules.backbone.dit import DiT
from dots_tts.training import checkpoint as train_checkpoint
from dots_tts.training import utils as train_utils
from dots_tts.utils import util as util_module
_ALLOWED_TEACHER_SOLVERS = ("euler", "midpoint", "rk4")
_ALLOWED_CFG_DISTILL_MODES = ("natural", "fused")
_ALLOWED_ANCHOR_TARGETS = ("formula", "teacher")
@dataclass(frozen=True, slots=True)
class MeanFlowSettings:
teacher_model_path: str | None
teacher_steps: int = 8
teacher_solver: str = "euler"
cfg_distill_mode: str = "fused"
distill_cfg_scale: float = 1.2
anchor_prob: float = 0.5
anchor_target: str = "formula"
time_sampling_mean: float = -0.4
time_sampling_std: float = 1.0
train_all_parameters: bool = False
def __post_init__(self) -> None:
if int(self.teacher_steps) <= 0:
raise ValueError("teacher_steps must be positive.")
if self.teacher_solver not in _ALLOWED_TEACHER_SOLVERS:
raise ValueError(
f"teacher_solver must be one of {_ALLOWED_TEACHER_SOLVERS}, "
f"got {self.teacher_solver!r}."
)
if self.cfg_distill_mode not in _ALLOWED_CFG_DISTILL_MODES:
raise ValueError(
"cfg_distill_mode must be one of "
f"{_ALLOWED_CFG_DISTILL_MODES}, got {self.cfg_distill_mode!r}."
)
if self.anchor_target not in _ALLOWED_ANCHOR_TARGETS:
raise ValueError(
f"anchor_target must be one of {_ALLOWED_ANCHOR_TARGETS}, "
f"got {self.anchor_target!r}."
)
if not 0.0 <= float(self.anchor_prob) <= 1.0:
raise ValueError("anchor_prob must be in [0, 1].")
def to_dict(self) -> dict[str, Any]:
return {
"teacher_model_path": self.teacher_model_path,
"teacher_steps": int(self.teacher_steps),
"teacher_solver": self.teacher_solver,
"cfg_distill_mode": self.cfg_distill_mode,
"distill_cfg_scale": float(self.distill_cfg_scale),
"anchor_prob": float(self.anchor_prob),
"anchor_target": self.anchor_target,
"time_sampling_mean": float(self.time_sampling_mean),
"time_sampling_std": float(self.time_sampling_std),
"train_all_parameters": bool(self.train_all_parameters),
}
def enable_meanflow_student(model: dots_tts_model.DotsTtsModel) -> None:
meanflow_config = MeanFlowConfig(enabled=True, use_duration_embedding=True)
model.config.meanflow = meanflow_config
model.core.meanflow_config = meanflow_config
model.core.mode = "meanflow"
old_dit = model.core.velocity_field_predictor
if getattr(old_dit, "duration_embedder", None) is not None:
return
new_dit = DiT(
in_dim=model.core.fm_hidden_size,
out_dim=model.core.latent_dim,
transformer_config=model.core.config.DiT,
mode="meanflow",
)
missing_keys, unexpected_keys = new_dit.load_state_dict(
old_dit.state_dict(),
strict=False,
)
missing_keys = [
key for key in missing_keys if not key.startswith("duration_embedder.")
]
if missing_keys or unexpected_keys:
raise RuntimeError(
"Failed to initialize MeanFlow DiT from the pretrained flow-matching "
f"DiT: missing={missing_keys[:5]} unexpected={unexpected_keys[:5]}"
)
duration_output = new_dit.duration_embedder.mlp[-1]
nn.init.zeros_(duration_output.weight)
nn.init.zeros_(duration_output.bias)
model.core.velocity_field_predictor = new_dit
class MeanFlowDotsTtsModel(nn.Module):
def __init__(
self,
student: dots_tts_model.DotsTtsModel,
settings: MeanFlowSettings,
):
super().__init__()
self.student = student
self.settings = settings
self._teacher_holder: dict[str, dots_tts_model.DotsTtsModel] = {}
@property
def config(self):
return self.student.config
@property
def tokenizer(self):
return self.student.tokenizer
@property
def teacher(self) -> dots_tts_model.DotsTtsModel:
teacher = self._teacher_holder.get("model")
if teacher is None:
raise RuntimeError("MeanFlow teacher model has not been initialized.")
return teacher
def set_teacher(self, teacher: dots_tts_model.DotsTtsModel) -> None:
for param in teacher.parameters():
param.requires_grad_(False)
teacher.eval()
self._teacher_holder["model"] = teacher
def to(self, *args, **kwargs):
super().to(*args, **kwargs)
teacher = self._teacher_holder.get("model")
if teacher is not None:
self._teacher_holder["model"] = teacher.to(*args, **kwargs)
self._teacher_holder["model"].eval()
return self
def cuda(self, device=None):
super().cuda(device)
teacher = self._teacher_holder.get("model")
if teacher is not None:
self._teacher_holder["model"] = teacher.cuda(device).eval()
return self
def train(self, mode: bool = True):
super().train(mode)
teacher = self._teacher_holder.get("model")
if teacher is not None:
teacher.eval()
return self
def prepare_training_batch(self, data: dict[str, Any]) -> dict[str, Any]:
return self.student.prepare_training_batch(data)
def save_pretrained(self, save_directory: str | Path) -> Path:
return self.student.save_pretrained(save_directory)
def load_pretrained_weights(
self, pretrained_model_name_or_path: str | Path
) -> None:
self.student.load_pretrained_weights(pretrained_model_name_or_path)
def set_cfg_droprate(
self,
cfg_droprate: float | None = None,
xvec_drop_rate: float | None = None,
) -> None:
self.student.set_cfg_droprate(
cfg_droprate=cfg_droprate,
xvec_drop_rate=xvec_drop_rate,
)
@torch.no_grad()
def compute_teacher_meanflow_target(
self,
*,
xt: torch.Tensor,
t: torch.Tensor,
delta_t: torch.Tensor,
prefix_data: dict[str, Any],
g_cond: torch.Tensor | None,
cfg_distill: bool,
uncond_prefix_data: dict[str, Any] | None,
uncond_g_cond: torch.Tensor | None,
) -> torch.Tensor:
teacher_core = self.teacher.core
teacher_dit = teacher_core.velocity_field_predictor
io_helper = teacher_core.io_helper
noisy_proj = teacher_core.coordinate_proj
n_steps = int(self.settings.teacher_steps)
solver = self.settings.teacher_solver
cfg_scale = float(self.settings.distill_cfg_scale)
if solver not in _ALLOWED_TEACHER_SOLVERS:
raise ValueError(f"Unsupported teacher solver: {solver!r}.")
device = xt.device
batch_size = xt.size(0)
latent_lens = prefix_data["latent_lens"]
latent_patch_size = int(prefix_data["latent_patch_size"])
anchor_mask = delta_t.float() == 0
autocast_device = "cuda" if device.type == "cuda" else "cpu"
with torch.autocast(device_type=autocast_device, enabled=False):
z = xt.float()
cur_t = t.float()
safe_dt = delta_t.float().clamp(min=1e-6)
step_dt = safe_dt / n_steps
def evaluate(z_in: torch.Tensor, t_val: torch.Tensor) -> torch.Tensor:
fm_seq = io_helper.replace_noise_latents_in_fm_seq(
prefix_data,
z_in.to(xt.dtype),
noisy_proj,
).float()
vt = teacher_dit(
x=fm_seq,
timesteps=t_val,
pos_ids=prefix_data["fm_pos_ids"],
mask=prefix_data["fm_seq_mask"],
attn_mask=prefix_data["fm_attn_mask"],
g_cond=None if g_cond is None else g_cond.float(),
)
pred = io_helper.get_dit_outputs(
pred_v=vt,
fm_prefix_lengths=prefix_data["fm_prefix_lengths"],
fm_gen_lengths=prefix_data["fm_gen_lengths"],
fm_gen_patch_size=prefix_data["fm_gen_patch_size"],
latent_patch_size=prefix_data["latent_patch_size"],
)
if cfg_distill:
if uncond_prefix_data is None:
raise RuntimeError(
"CFG distillation requires an uncond prefix."
)
fm_seq_u = io_helper.replace_noise_latents_in_fm_seq(
uncond_prefix_data,
z_in.to(xt.dtype),
noisy_proj,
).float()
vt_u = teacher_dit(
x=fm_seq_u,
timesteps=t_val,
pos_ids=uncond_prefix_data["fm_pos_ids"],
mask=uncond_prefix_data["fm_seq_mask"],
attn_mask=uncond_prefix_data["fm_attn_mask"],
g_cond=None if uncond_g_cond is None else uncond_g_cond.float(),
)
pred_u = io_helper.get_dit_outputs(
pred_v=vt_u,
fm_prefix_lengths=uncond_prefix_data["fm_prefix_lengths"],
fm_gen_lengths=uncond_prefix_data["fm_gen_lengths"],
fm_gen_patch_size=uncond_prefix_data["fm_gen_patch_size"],
latent_patch_size=uncond_prefix_data["latent_patch_size"],
)
pred = pred + cfg_scale * (pred - pred_u)
return rearrange(pred, "n p d -> (n p) d")
v_init_flat = evaluate(z, cur_t)
def apply_velocity(
z_cur: torch.Tensor,
v_flat: torch.Tensor,
*,
dt_factor: float,
) -> torch.Tensor:
new_z = z_cur.clone()
offset = 0
for batch_idx in range(batch_size):
length = int(latent_lens[batch_idx].item())
if length <= 0:
continue
if not bool(anchor_mask[batch_idx].item()):
new_z[batch_idx, :length, :] = z_cur[
batch_idx, :length, :
] + v_flat[offset : offset + length, :] * (
step_dt[batch_idx] * float(dt_factor)
)
offset += length
return new_z
if solver == "euler":
v_flat = v_init_flat
for step in range(n_steps):
if step > 0:
v_flat = evaluate(z, cur_t)
z = apply_velocity(z, v_flat, dt_factor=1.0)
cur_t = cur_t + step_dt
elif solver == "midpoint":
for step in range(n_steps):
k1 = v_init_flat if step == 0 else evaluate(z, cur_t)
z_mid = apply_velocity(z, k1, dt_factor=0.5)
k2 = evaluate(z_mid, cur_t + 0.5 * step_dt)
z = apply_velocity(z, k2, dt_factor=1.0)
cur_t = cur_t + step_dt
else:
for step in range(n_steps):
k1 = v_init_flat if step == 0 else evaluate(z, cur_t)
z1 = apply_velocity(z, k1, dt_factor=0.5)
k2 = evaluate(z1, cur_t + 0.5 * step_dt)
z2 = apply_velocity(z, k2, dt_factor=0.5)
k3 = evaluate(z2, cur_t + 0.5 * step_dt)
z3 = apply_velocity(z, k3, dt_factor=1.0)
k4 = evaluate(z3, cur_t + step_dt)
z = apply_velocity(
z,
(k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0,
dt_factor=1.0,
)
cur_t = cur_t + step_dt
mean_velocity = (z - xt.float()) / safe_dt.view(-1, 1, 1)
target_chunks = []
offset = 0
for batch_idx in range(batch_size):
length = int(latent_lens[batch_idx].item())
if length <= 0:
continue
if bool(anchor_mask[batch_idx].item()):
target_b = v_init_flat[offset : offset + length, :]
else:
target_b = mean_velocity[batch_idx, :length, :]
target_chunks.append(
rearrange(target_b, "(n p) d -> n p d", p=latent_patch_size)
)
offset += length
if not target_chunks:
raise RuntimeError("Teacher rollout produced no MeanFlow target.")
return torch.cat(target_chunks, dim=0).to(xt.dtype)
def forward(self, data: dict[str, Any]):
loss_masks = data["loss_masks"]
processed = self.student.prepare_training_inputs(data)
processed["input_span_mask"] = data["input_span_mask"]
processed["output_span_mask"] = data["output_span_mask"]
outputs = self.meanflow_forward(processed)
return self.student._compute_loss_terms(
outputs,
labels=processed["labels"],
loss_masks=loss_masks,
)
def meanflow_forward(self, data: dict[str, Any]) -> DotsTtsForwardOutput:
core = self.student.core
input_ids: torch.Tensor = data["input_ids"]
input_ids_lengths: torch.Tensor = data["input_ids_lengths"]
input_span_mask: torch.Tensor = data["input_span_mask"]
output_span_mask: torch.Tensor = data["output_span_mask"]
batch_size = input_ids.size(0)
device = input_ids.device
latents: torch.Tensor | None = data.get("latents")
latents_sampled: torch.Tensor | None = data.get("latents_sampled")
latent_lengths: torch.Tensor | None = data.get("latent_lengths")
has_latents = latents is not None or latents_sampled is not None
if has_latents:
if latents_sampled is None:
latents_sampled = core.io_helper.sample_from_latent(latents)
patch_embeddings = core.patch_encoder(
latents_sampled, x_lens=latent_lengths
)
valid_patch_counts = latent_lengths // core.latent_patch_size
latents_sampled = core.io_helper.normalize(latents_sampled)
else:
latents_sampled = None
patch_embeddings = None
valid_patch_counts = torch.zeros(
batch_size,
dtype=torch.long,
device=device,
)
input_span_counts = input_span_mask.sum(dim=1)
if input_span_counts.sum() > 0 and patch_embeddings is None:
raise RuntimeError(
"Found audio span tokens but no latents provided to compute patch embeddings."
)
inputs_embeds = core.llm.get_input_embeddings()(input_ids)
if patch_embeddings is not None:
inputs_embeds = inputs_embeds.clone()
patch_embeddings = patch_embeddings.to(inputs_embeds.dtype)
for batch_idx in range(batch_size):
span_num = int(input_span_counts[batch_idx].item())
if span_num == 0:
continue
expected = int(valid_patch_counts[batch_idx].item())
if expected != span_num:
raise RuntimeError(
f"Mismatch between span tokens ({span_num}) and latent patches "
f"({expected}) for sample {batch_idx}."
)
indices = input_span_mask[batch_idx].nonzero(as_tuple=False).squeeze(-1)
inputs_embeds[batch_idx, indices, :] = patch_embeddings[
batch_idx,
:span_num,
:,
]
_llm_attn_mask, llm_seq_mask, _ = core.causal_helper.create_causal_mask_and_pos(
seq_lens=input_ids_lengths,
max_len=input_ids.size(1),
)
llm_outputs = core.llm(
inputs_embeds=inputs_embeds,
attention_mask=llm_seq_mask.long(),
use_cache=False,
output_hidden_states=True,
return_dict=True,
)
llm_logits = llm_outputs.logits
llm_hidden = llm_outputs.hidden_states[-1]
eos = core.eos_proj(llm_hidden.detach())
total_patches = int(output_span_mask.sum().item())
if total_patches > 0 and latents_sampled is None:
raise RuntimeError("MeanFlow training requested but latents are missing.")
if total_patches > 0:
pred, target = self.meanflow_fm_segment(
data,
llm_hidden=llm_hidden,
inputs_embeds=inputs_embeds,
output_span_mask=output_span_mask,
latents_sampled=latents_sampled,
latent_lengths=latent_lengths,
)
else:
pred, target = self.dummy_fm_forward(core, llm_hidden, device)
return DotsTtsForwardOutput(
llm_logits=llm_logits,
pred=pred,
target=target,
eos_out=eos,
)
def meanflow_fm_segment(
self,
data: dict[str, Any],
*,
llm_hidden: torch.Tensor,
inputs_embeds: torch.Tensor,
output_span_mask: torch.Tensor,
latents_sampled: torch.Tensor,
latent_lengths: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
core = self.student.core
teacher_core = self.teacher.core
settings = self.settings
batch_size = latents_sampled.size(0)
device = latents_sampled.device
latent_dtype = latents_sampled.dtype
first_t = torch.randn(batch_size, device=device, dtype=latent_dtype)
second_t = torch.randn(batch_size, device=device, dtype=latent_dtype)
first_t = torch.sigmoid(
first_t * float(settings.time_sampling_std)
+ float(settings.time_sampling_mean)
)
second_t = torch.sigmoid(
second_t * float(settings.time_sampling_std)
+ float(settings.time_sampling_mean)
)
t_vec = torch.minimum(first_t, second_t)
delta_t = (first_t - second_t).abs()
anchor_mask = torch.rand(batch_size, device=device, dtype=latent_dtype) < float(
settings.anchor_prob
)
delta_t = torch.where(anchor_mask, torch.zeros_like(delta_t), delta_t)
z0 = torch.randn_like(latents_sampled)
xt = core.fm_helper.sample_x_t(
z0,
latents_sampled,
t_vec.view(-1, 1, 1).to(latent_dtype),
)
fused_cfg = settings.cfg_distill_mode == "fused"
if fused_cfg:
cfg_mask = torch.zeros(batch_size, device=device, dtype=torch.bool)
xvec_drop_mask = torch.zeros(batch_size, device=device, dtype=torch.bool)
else:
cfg_mask = torch.empty(
batch_size, device=device, dtype=torch.float32
).uniform_(0, 1) < float(core.cfg_droprate)
xvec_drop_mask = torch.empty(
batch_size, device=device, dtype=torch.float32
).uniform_(0, 1) < float(core.xvec_drop_rate)
xvec_cond = core.xvec_proj(data["xvector"])
vocal_mask = data.get("vocal_mask")
if vocal_mask is None:
vocal_mask = torch.ones(batch_size, device=device, dtype=torch.bool)
xvec_cond = util_module.mask_data(xvec_cond, xvec_drop_mask & vocal_mask)
hiddens_for_fm = torch.where(
output_span_mask.unsqueeze(-1),
llm_hidden,
inputs_embeds,
)
prefix_data = core.io_helper.prepare_meanflow_inputs_for_dit(
hiddens=hiddens_for_fm,
latents=latents_sampled,
latent_lens=latent_lengths,
hidden_proj=core.hidden_proj,
latent_proj=core.latent_proj,
noisy_proj=core.coordinate_proj,
span_mask=output_span_mask,
hidden_patch_size=core.hidden_patch_size,
latent_patch_size=core.latent_patch_size,
cfg_mask=cfg_mask,
noise_latents=xt,
)
uncond_prefix_data = None
uncond_g_cond = None
with torch.no_grad():
teacher_xvec_cond = teacher_core.xvec_proj(data["xvector"])
teacher_xvec_cond = util_module.mask_data(
teacher_xvec_cond,
xvec_drop_mask & vocal_mask,
)
teacher_prefix_data = (
teacher_core.io_helper.prepare_meanflow_inputs_for_dit(
hiddens=hiddens_for_fm,
latents=latents_sampled,
latent_lens=latent_lengths,
hidden_proj=teacher_core.hidden_proj,
latent_proj=teacher_core.latent_proj,
noisy_proj=teacher_core.coordinate_proj,
span_mask=output_span_mask,
hidden_patch_size=teacher_core.hidden_patch_size,
latent_patch_size=teacher_core.latent_patch_size,
cfg_mask=cfg_mask,
noise_latents=xt,
)
)
if fused_cfg:
uncond_prefix_data = (
teacher_core.io_helper.prepare_meanflow_inputs_for_dit(
hiddens=hiddens_for_fm,
latents=latents_sampled,
latent_lens=latent_lengths,
hidden_proj=teacher_core.hidden_proj,
latent_proj=teacher_core.latent_proj,
noisy_proj=teacher_core.coordinate_proj,
span_mask=output_span_mask,
hidden_patch_size=teacher_core.hidden_patch_size,
latent_patch_size=teacher_core.latent_patch_size,
cfg_mask=torch.ones(
batch_size, device=device, dtype=torch.bool
),
noise_latents=xt,
)
)
uncond_g_cond = torch.zeros_like(teacher_xvec_cond)
teacher_target = self.compute_teacher_meanflow_target(
xt=xt,
t=t_vec,
delta_t=delta_t,
prefix_data=teacher_prefix_data,
g_cond=teacher_xvec_cond,
cfg_distill=fused_cfg,
uncond_prefix_data=uncond_prefix_data,
uncond_g_cond=uncond_g_cond,
)
if anchor_mask.any() and settings.anchor_target == "formula":
target = self.replace_anchor_targets_with_formula(
teacher_target,
z0=z0,
latents_sampled=latents_sampled,
latent_lengths=latent_lengths,
anchor_mask=anchor_mask,
)
else:
target = teacher_target
student_vt = core.velocity_field_predictor(
x=prefix_data["fm_seq"],
timesteps=t_vec,
duration=delta_t,
pos_ids=prefix_data["fm_pos_ids"],
mask=prefix_data["fm_seq_mask"],
attn_mask=prefix_data["fm_attn_mask"],
g_cond=xvec_cond,
)
pred = core.io_helper.get_dit_outputs(
pred_v=student_vt,
fm_prefix_lengths=prefix_data["fm_prefix_lengths"],
fm_gen_lengths=prefix_data["fm_gen_lengths"],
fm_gen_patch_size=prefix_data["fm_gen_patch_size"],
latent_patch_size=prefix_data["latent_patch_size"],
)
return pred, target
def replace_anchor_targets_with_formula(
self,
teacher_target: torch.Tensor,
*,
z0: torch.Tensor,
latents_sampled: torch.Tensor,
latent_lengths: torch.Tensor,
anchor_mask: torch.Tensor,
) -> torch.Tensor:
core = self.student.core
formula_target = core.fm_helper.compute_u_t(z0, latents_sampled)
chunks = []
offset = 0
for batch_idx in range(latents_sampled.size(0)):
length = int(latent_lengths[batch_idx].item())
if length <= 0:
continue
patch_count = length // core.latent_patch_size
if bool(anchor_mask[batch_idx].item()):
chunks.append(
rearrange(
formula_target[batch_idx, :length, :],
"(n p) d -> n p d",
p=core.latent_patch_size,
)
)
else:
chunks.append(teacher_target[offset : offset + patch_count])
offset += patch_count
if not chunks:
raise RuntimeError("Anchor target replacement produced no target.")
return torch.cat(chunks, dim=0)
def dummy_fm_forward(
self,
core,
llm_hidden: torch.Tensor,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
dummy_length = core.latent_patch_size
dummy_seq_h = llm_hidden.new_zeros((1, dummy_length, core.llm_hidden_size))
dummy_seq_h = core.hidden_proj(dummy_seq_h) * 0.0
dummy_seq_l = llm_hidden.new_zeros((1, dummy_length, core.latent_dim))
dummy_seq_l = core.latent_proj(dummy_seq_l) * 0.0
dummy_seq_c = llm_hidden.new_zeros((1, dummy_length, core.latent_dim))
dummy_seq_c = core.coordinate_proj(dummy_seq_c) * 0.0
dummy_seq = dummy_seq_h + dummy_seq_l + dummy_seq_c
dummy_times = torch.zeros((1,), device=device, dtype=torch.float32)
dummy_duration = torch.zeros((1,), device=device, dtype=torch.float32)
dummy_attn_mask = torch.ones(
(1, dummy_length, dummy_length),
device=device,
dtype=torch.bool,
)
dummy_out = core.velocity_field_predictor(
x=dummy_seq,
timesteps=dummy_times,
duration=dummy_duration,
attn_mask=dummy_attn_mask,
)
pred = dummy_out[:, -core.latent_patch_size :, :]
return pred, pred.detach()
class DotsTtsMeanFlowTrainingRun(DotsTtsTrainingRun):
def __init__(
self,
cfg: app_config.AppConfig,
*,
meanflow_settings: MeanFlowSettings,
debug_enabled: bool = False,
):
self.cfg = cfg
self.meanflow_settings = meanflow_settings
self.progress = train_utils.TrainProgress()
self.max_train_steps = int(cfg.train.max_train_steps)
self.grad_accumulation_steps = int(cfg.train.gradient_accumulation_steps)
self.last_validation_step: int | None = None
self.consecutive_empty_epochs = 0
self.saved_latest_checkpoint = False
self._last_log_step = 0
self._last_log_time = 0.0
self._debug_enabled = bool(debug_enabled)
self._debug_batch_count = 0
self._debug_audio_sample_rate = int(self.cfg.train_data.train_audio_sample_rate)
project_config = ProjectConfiguration(
project_dir=self.cfg.train.output_dir,
total_limit=self.cfg.train.max_checkpoints_to_keep,
)
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
self.accelerator = Accelerator(
kwargs_handlers=[ddp_kwargs],
gradient_accumulation_steps=self.grad_accumulation_steps,
log_with="tensorboard",
project_config=project_config,
step_scheduler_with_optimizer=False,
)
util_module.seed_everything(self.cfg.train.seed)
student = dots_tts_model.DotsTtsModel.from_pretrained(
self.cfg.train.pretrained_model_path
)
student.set_cfg_droprate(
cfg_droprate=self.cfg.train.cfg_droprate,
xvec_drop_rate=self.cfg.train.xvec_drop_rate,
)
enable_meanflow_student(student)
if not bool(meanflow_settings.train_all_parameters):
for param in student.parameters():
param.requires_grad_(False)
for param in student.core.velocity_field_predictor.parameters():
param.requires_grad_(True)
model = MeanFlowDotsTtsModel(student, meanflow_settings)
teacher_path = (
meanflow_settings.teacher_model_path or self.cfg.train.pretrained_model_path
)
teacher = dots_tts_model.DotsTtsModel.from_pretrained(teacher_path)
model.set_teacher(teacher)
optimizer = AdamW(
(param for param in model.parameters() if param.requires_grad),
lr=self.cfg.train.learning_rate,
weight_decay=self.cfg.train.weight_decay,
)
scheduler = get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps=self.cfg.train.warmup_steps,
num_training_steps=self.max_train_steps,
)
self.model, self.optimizer, self.scheduler = self.accelerator.prepare(
model,
optimizer,
scheduler,
)
self.unwrapped_model = self.accelerator.unwrap_model(self.model)
self.unwrapped_model.to(self.accelerator.device)
expected_sample_rate = int(self.unwrapped_model.config.vocoder.sample_rate)
expected_audio_samples_per_llm_token = int(
self.unwrapped_model.student.hop_size
) * int(self.unwrapped_model.config.patch_size)
if int(self.cfg.train_data.train_audio_sample_rate) != expected_sample_rate:
raise ValueError(
f"train_data.train_audio_sample_rate={int(self.cfg.train_data.train_audio_sample_rate)} "
f"does not match the pretrained model sample rate {expected_sample_rate}."
)
if (
int(self.cfg.train_data.audio_samples_per_llm_token)
!= expected_audio_samples_per_llm_token
):
raise ValueError(
"train_data.audio_samples_per_llm_token="
f"{int(self.cfg.train_data.audio_samples_per_llm_token)} "
"does not match the pretrained model audio token contract "
f"{expected_audio_samples_per_llm_token}."
)
if self.cfg.val_data is not None:
if int(self.cfg.val_data.train_audio_sample_rate) != expected_sample_rate:
raise ValueError(
f"val_data.train_audio_sample_rate={int(self.cfg.val_data.train_audio_sample_rate)} "
f"does not match the pretrained model sample rate {expected_sample_rate}."
)
if (
int(self.cfg.val_data.audio_samples_per_llm_token)
!= expected_audio_samples_per_llm_token
):
raise ValueError(
"val_data.audio_samples_per_llm_token="
f"{int(self.cfg.val_data.audio_samples_per_llm_token)} "
"does not match the pretrained model audio token contract "
f"{expected_audio_samples_per_llm_token}."
)
if self.accelerator.is_main_process:
total_params = sum(
param.numel() for param in self.unwrapped_model.parameters()
)
trainable_params = sum(
param.numel()
for param in self.unwrapped_model.parameters()
if param.requires_grad
)
self.accelerator.print(f"Total parameters: {total_params:,}")
self.accelerator.print(f"Trainable parameters: {trainable_params:,}")
self.accelerator.print(
f"MeanFlow teacher path: {Path(teacher_path).expanduser()}"
)
self.accelerator.print(
f"Distributed type: {self.accelerator.distributed_type}"
)
tokenizer = self.unwrapped_model.tokenizer
self.tokenizer = tokenizer
train_dataset = data_module.build_training_dataset(
self.cfg.train_data,
tokenizer=tokenizer,
seed=int(self.cfg.train.seed),
accelerator=self.accelerator,
)
self.train_loader = data_module.build_training_dataloader(
train_dataset,
self.cfg.train_data,
tokenizer=tokenizer,
)
self.val_loader = None
if self.cfg.train.eval_interval is not None or self.cfg.train.run_eval_on_start:
if self.cfg.val_data is None:
raise ValueError(
"Validation requires val_data when eval_interval or "
"run_eval_on_start is enabled."
)
validation_data_cfg = self.cfg.val_data.model_copy(deep=True)
validation_data_cfg.num_tokens_per_epoch = None
val_dataset = data_module.build_validation_dataset(
validation_data_cfg,
tokenizer=tokenizer,
seed=int(self.cfg.train.seed),
accelerator=self.accelerator,
)
self.val_loader = data_module.build_validation_dataloader(
val_dataset,
validation_data_cfg,
tokenizer=tokenizer,
)
self._resume_if_available()
self.train_loader.set_epoch(self.progress.epoch)
def _write_run_config(self) -> None:
if not bool(getattr(self.accelerator, "is_main_process", True)):
return
output_dir = Path(self.cfg.train.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
config_path = output_dir / "config.yml"
payload = self.cfg.to_dict()
payload["meanflow_train"] = self.meanflow_settings.to_dict()
with config_path.open("w", encoding="utf-8") as fout:
yaml.safe_dump(
payload,
fout,
sort_keys=False,
allow_unicode=True,
)
def _save_checkpoint(self, learning_rate: float) -> None:
train_checkpoint.save_train_checkpoint(
self.accelerator,
self.model,
self.optimizer,
self.progress,
self.cfg.train.output_dir,
self.cfg.train.max_checkpoints_to_keep,
self.train_loader.state_dict(),
{
"type": "transformers_cosine_with_warmup_meanflow",
"global_step": int(self.progress.global_step),
"base_lr": float(self.cfg.train.learning_rate),
"current_lr": float(learning_rate),
"warmup_steps": int(self.cfg.train.warmup_steps),
"max_train_steps": int(self.max_train_steps),
"meanflow": self.meanflow_settings.to_dict(),
"state_dict": self.scheduler.state_dict(),
},
)
def parse_args(argv=None):
parser = argparse.ArgumentParser(
description="Accelerate MeanFlow distillation entrypoint for dots.tts."
)
parser.add_argument("--config", default=app_config.DEFAULT_CONFIG_PATH)
parser.add_argument(
"--debug",
action="store_true",
help="Print training debug information.",
)
parser.add_argument(
"--teacher-model-path",
default=None,
help=(
"Frozen flow-matching teacher model path. Defaults to "
"train.pretrained_model_path."
),
)
parser.add_argument("--teacher-steps", type=int, default=8)
parser.add_argument(
"--teacher-solver",
choices=_ALLOWED_TEACHER_SOLVERS,
default="euler",
)
parser.add_argument(
"--cfg-distill-mode",
choices=_ALLOWED_CFG_DISTILL_MODES,
default="fused",
)
parser.add_argument("--distill-cfg-scale", type=float, default=1.2)
parser.add_argument("--anchor-prob", type=float, default=0.5)
parser.add_argument(
"--anchor-target",
choices=_ALLOWED_ANCHOR_TARGETS,
default="formula",
)
parser.add_argument("--time-sampling-mean", type=float, default=-0.4)
parser.add_argument("--time-sampling-std", type=float, default=1.0)
parser.add_argument(
"--train-all-parameters",
action="store_true",
help="Train all regular dots.tts parameters instead of only the DiT.",
)
return parser.parse_args(argv)
def main(argv=None):
args = parse_args(argv)
settings = MeanFlowSettings(
teacher_model_path=args.teacher_model_path,
teacher_steps=args.teacher_steps,
teacher_solver=args.teacher_solver,
cfg_distill_mode=args.cfg_distill_mode,
distill_cfg_scale=args.distill_cfg_scale,
anchor_prob=args.anchor_prob,
anchor_target=args.anchor_target,
time_sampling_mean=args.time_sampling_mean,
time_sampling_std=args.time_sampling_std,
train_all_parameters=args.train_all_parameters,
)
return DotsTtsMeanFlowTrainingRun(
app_config.load_config(args.config),
meanflow_settings=settings,
debug_enabled=args.debug,
).run()
if __name__ == "__main__":
raise SystemExit(main())
+1
View File
@@ -0,0 +1 @@
"""dots.tts package."""
+152
View File
@@ -0,0 +1,152 @@
from __future__ import annotations
import argparse
from pathlib import Path
def parse_args(argv=None):
parser = argparse.ArgumentParser(description="dots.tts inference CLI.")
template_choices = ("tts", "instruction_tts", "text_to_audio", "tts_interleave")
parser.add_argument(
"--model-name-or-path",
required=True,
help="Local pretrained directory or Hugging Face repo id",
)
parser.add_argument(
"--revision", default=None, help="Optional Hugging Face revision"
)
parser.add_argument(
"--cache-dir", default=None, help="Optional Hugging Face cache dir"
)
parser.add_argument("--text", type=str, required=True, help="Input text")
parser.add_argument("--output", default="output.wav", help="Output wav file path")
parser.add_argument(
"--precision", type=str, default="bfloat16", help="Inference precision"
)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Random seed for inference.",
)
parser.add_argument(
"--prompt-audio", type=str, default=None, help="Path to prompt audio"
)
parser.add_argument(
"--prompt-text", type=str, default=None, help="Transcript of prompt audio"
)
parser.add_argument(
"--language",
type=str,
default=None,
help="Language tag mode. Default: none. Supported values: none, auto_detect, or a language code/name such as EN/en/english/chinese.",
)
parser.add_argument(
"--template-name",
choices=template_choices,
default=None,
help="Named template preset for generation.",
)
parser.add_argument(
"--ode-method", type=str, default="euler", help="ODE solver method"
)
parser.add_argument(
"--num-steps", type=int, default=10, help="Diffusion sampling steps"
)
parser.add_argument(
"--guidance-scale",
type=float,
default=1.2,
help="Classifier-free guidance scale",
)
parser.add_argument(
"--speaker-scale",
type=float,
default=1.5,
help="Scale applied to the reference speaker embedding",
)
parser.add_argument(
"--max-generate-length",
type=int,
default=500,
help="Maximum total audio patch count (prompt + generated)",
)
parser.add_argument(
"--normalize-text",
action="store_true",
help="Whether to normalize text before inference",
)
parser.add_argument(
"--profile-inference",
action="store_true",
help="Collect per-module inference timing statistics",
)
return parser.parse_args(argv)
def main(argv=None):
args = parse_args(argv)
import soundfile as sf
from loguru import logger
from dots_tts.runtime import DotsTtsRuntime
from dots_tts.utils.logging import configure_logging
from dots_tts.utils.util import seed_everything
configure_logging()
seed_everything(args.seed)
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
logger.info(
"CLI command started: model={} output={} seed={}",
args.model_name_or_path,
output_path,
args.seed,
)
try:
runtime = DotsTtsRuntime.from_pretrained(
args.model_name_or_path,
revision=args.revision,
cache_dir=args.cache_dir,
precision=args.precision,
max_generate_length=args.max_generate_length,
)
result = runtime.generate(
text=args.text,
prompt_audio_path=args.prompt_audio,
prompt_text=args.prompt_text,
language=args.language,
template_name=args.template_name,
ode_method=args.ode_method,
num_steps=args.num_steps,
guidance_scale=args.guidance_scale,
speaker_scale=args.speaker_scale,
normalize_text=args.normalize_text,
profile_inference=args.profile_inference,
)
sf.write(
output_path,
result["audio"].float().cpu().squeeze().numpy(),
result["sample_rate"],
)
except Exception:
logger.exception(
"CLI inference failed: model={} output={}",
args.model_name_or_path,
output_path,
)
raise
logger.info(
"CLI output written: request_id={} output={} sample_rate={} samples={}",
result["fid"],
output_path,
result["sample_rate"],
int(result["audio"].shape[-1]),
)
if __name__ == "__main__":
raise SystemExit(main())
+1
View File
@@ -0,0 +1 @@
"""Configuration package."""
+32
View File
@@ -0,0 +1,32 @@
from __future__ import annotations
from pathlib import Path
import yaml
from dots_tts.config.base import StrictConfigBase
from dots_tts.config.data import DataConfig
from dots_tts.config.train import TrainConfig
from dots_tts.models.dots_tts.config import LossConfig
DEFAULT_CONFIG_PATH = "configs/dots_tts.yaml"
class AppConfig(StrictConfigBase):
train_data: DataConfig
val_data: DataConfig | None = None
loss: LossConfig
train: TrainConfig
@classmethod
def from_yaml(cls, config_path: str = DEFAULT_CONFIG_PATH) -> AppConfig:
with Path(config_path).open(encoding="utf-8") as fin:
raw_config = yaml.safe_load(fin)
return cls.model_validate(raw_config)
def load_config(config_path: str = DEFAULT_CONFIG_PATH) -> AppConfig:
return AppConfig.from_yaml(config_path)
__all__ = ["AppConfig", "DEFAULT_CONFIG_PATH", "load_config"]
+64
View File
@@ -0,0 +1,64 @@
from __future__ import annotations
from typing import Any
from pydantic import BaseModel, ConfigDict
class ConfigBase(BaseModel):
model_config = ConfigDict(
extra="allow",
validate_assignment=True,
arbitrary_types_allowed=True,
)
def get(self, key: str, default=None):
value = getattr(self, key, default)
if value is default:
return value
fields_set = self.model_fields_set
if value is None and key not in fields_set:
return default
return value
def to_dict(self) -> dict[str, Any]:
return self.model_dump(exclude_none=True)
@classmethod
def _declared_field_names(cls) -> list[str]:
return [name for name in cls.model_fields if name != "model_config"]
@classmethod
def _serialize_declared_value(cls, value):
if isinstance(value, ConfigBase):
return value.to_declared_dict()
if isinstance(value, list):
return [cls._serialize_declared_value(item) for item in value]
if isinstance(value, tuple):
return [cls._serialize_declared_value(item) for item in value]
if isinstance(value, dict):
return {
key: cls._serialize_declared_value(item) for key, item in value.items()
}
return value
def to_declared_dict(self) -> dict[str, Any]:
data = {}
for name in self._declared_field_names():
value = getattr(self, name, None)
if value is None:
continue
data[name] = self._serialize_declared_value(value)
return data
class StrictConfigBase(ConfigBase):
model_config = ConfigDict(
extra="forbid",
validate_assignment=True,
arbitrary_types_allowed=True,
)
__all__ = ["ConfigBase", "StrictConfigBase"]
+63
View File
@@ -0,0 +1,63 @@
from __future__ import annotations
from typing import Any, Literal
from pydantic import Field, model_validator
from dots_tts.config.base import StrictConfigBase
DEFAULT_SOURCE_ADAPTER_CLASS_NAME = "JsonlManifestSourceAdapter"
class SourceAdapterConfig(StrictConfigBase):
class_name: Literal["JsonlManifestSourceAdapter"] = (
DEFAULT_SOURCE_ADAPTER_CLASS_NAME
)
params: dict[str, Any] = Field(default_factory=dict)
class DataSourceConfig(StrictConfigBase):
name: str
weight: float = Field(default=1.0, gt=0.0)
pipeline: Literal["basic", "interleave"] = "basic"
adapter: SourceAdapterConfig = Field(default_factory=SourceAdapterConfig)
class DataConfig(StrictConfigBase):
sources: list[DataSourceConfig]
train_audio_sample_rate: int = Field(ge=1)
audio_samples_per_llm_token: int = Field(ge=1)
num_tokens_per_epoch: int | None = Field(
default=None,
ge=1,
description="Global token budget across all ranks for one training epoch.",
)
num_workers: int = Field(default=0, ge=0)
pin_memory: bool = False
prefetch_factor: int = Field(
default=2,
ge=1,
description="Samples prefetched by each DataLoader worker.",
)
max_audio_seconds_in_batch: float = Field(gt=0.0)
max_text_tokens_in_batch: int = Field(ge=1)
max_samples_per_batch: int | None = Field(default=None, ge=1)
bucketing_pool_size: int = Field(default=64, ge=1)
@model_validator(mode="after")
def _validate_unique_source_names(self) -> "DataConfig":
counts: dict[str, int] = {}
for source in self.sources:
counts[source.name] = counts.get(source.name, 0) + 1
duplicated = [name for name, count in counts.items() if count > 1]
if duplicated:
raise ValueError(f"Source names must be unique: {duplicated}")
return self
__all__ = [
"DEFAULT_SOURCE_ADAPTER_CLASS_NAME",
"DataConfig",
"DataSourceConfig",
"SourceAdapterConfig",
]
+28
View File
@@ -0,0 +1,28 @@
from __future__ import annotations
from pydantic import Field
from dots_tts.config.base import StrictConfigBase
class TrainConfig(StrictConfigBase):
pretrained_model_path: str
output_dir: str
seed: int = 42
learning_rate: float
cfg_droprate: float = 0.0
xvec_drop_rate: float = 0.5
weight_decay: float = 0.01
warmup_steps: int = 0
max_train_steps: int
gradient_accumulation_steps: int = Field(default=1, ge=1)
grad_clip_norm: float = 1.0
save_interval: int = Field(default=1000, ge=1)
max_checkpoints_to_keep: int = 10
log_interval: int = Field(default=10, ge=1)
eval_interval: int | None = Field(default=None, ge=1)
max_eval_batches: int | None = None
run_eval_on_start: bool = False
__all__ = ["TrainConfig"]
+124
View File
@@ -0,0 +1,124 @@
# Data Source Extension Guide
This document answers exactly one question: how to plug a new training data source into the current `dots_tts` data pipeline.
If you only need to swap in a different JSONL manifest, no code changes are required. To support a new raw data format, you usually only need to add:
- one **source adapter**
- optionally one **sample pipeline**
## Data flow
1. An **adapter** reads from the raw data source and yields raw samples.
2. A **pipeline** turns each raw sample into a training sample (1:1).
3. A **multi-source wrapper** handles mixing across sources and resume state.
4. `StreamingSampleDataset` / `DataLoader` pulls samples.
5. `OnlineBatcher` assembles batches and `PadCollator` performs padding.
## What an adapter must implement
Subclass `BaseSourceAdapter`:
```python
class BaseSourceAdapter(ABC):
@abstractmethod
def initial_state(self) -> dict[str, Any]:
...
@abstractmethod
def iter_samples(
self,
context: SourceContext,
*,
state: dict[str, Any] | None = None,
) -> Iterable[dict[str, Any]]:
...
@abstractmethod
def is_cycle_start_state(self, state: dict[str, Any] | None) -> bool:
...
# Optional — only required when used under WeightedMultiSourceAdapter,
# which cycles each finite child source independently. The default
# implementation raises if your adapter never gets re-cycled.
def advance_cycle(self, state: dict[str, Any] | None) -> dict[str, Any]:
...
```
Each emitted sample **must** carry these fields:
- `fid`
- `text`
- `audio`
- `_adapter_state`
Key constraints:
- `_adapter_state` must describe **where to resume next**, not the position of the current item.
- The state must be plain Python data — serializable and recoverable after a restart.
- If your source needs to be split across workers, use `context.global_worker_id` and `context.global_worker_count` (or subclass `ShardableSourceAdapter` and use its `is_assigned_index` / `shard_items` helpers).
- If the source will participate in weighted cyclic sampling, you must implement `advance_cycle` and make `is_cycle_start_state` correct — otherwise `WeightedMultiSourceAdapter` cannot detect an empty cycle and will raise.
After implementing the adapter, register the class in `dots_tts/data/builders.py::_SOURCE_ADAPTER_CLASSES` so that the YAML config can resolve it by `class_name`.
## What a pipeline must implement
Pipelines must subclass `BaseSamplePipeline` and perform a strict **1:1** sample transform.
Minimum implementation:
```python
class MyPipeline(BaseSamplePipeline):
def process_sample(self, sample: dict) -> dict:
sample["text"] = str(sample["text"]).strip()
return sample
```
Do **not**:
- filter samples out
- expand a single sample into multiple samples
- assemble batches inside the pipeline
`BaseSamplePipeline.__call__` automatically merges the original raw sample (including `_adapter_state` and any extra fields the adapter attached) with whatever your `process_sample` returns. You do not need to copy these fields manually — just return the fields you produced or want to overwrite.
To wire a new pipeline into config, also extend `dots_tts/data/builders.py::_build_source_pipeline` so it can be selected by name in YAML.
## How multi-source wrappers affect you
There are two wrappers in the current codebase:
- `SequentialMultiSourceAdapter` — used for validation. Reads sources in the configured order, exhaustively, once.
- `WeightedMultiSourceAdapter` — used for training. Draws sources by weight, cycles each child source independently when exhausted.
Both wrappers **replace** the `_adapter_state` produced by your child adapter with their own resume state before yielding to the dataset. Even so, the child adapter must still emit its own `_adapter_state` — the wrapper reads it to track where each sub-source has read to.
## Config
Each source is configured independently:
```yaml
train_data:
sources:
- name: train_a
weight: 1.0
pipeline: basic
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: train_a.jsonl
- name: train_b
weight: 2.0
pipeline: interleave
adapter:
class_name: JsonlManifestSourceAdapter
params:
manifest_path: train_b.jsonl
```
Constraints:
- `sources[].name` must be unique within the same `train_data` / `val_data` block (it is used as a dict key for resume state).
- `sources[].pipeline` is a per-source setting, not shared across the dataset.
- All sources must ultimately produce the same training-sample structure, since they feed into the same batcher and collator.
- `class_name` must match a key registered in `_SOURCE_ADAPTER_CLASSES`; `params` is forwarded verbatim as kwargs to the adapter constructor.
+1
View File
@@ -0,0 +1 @@
"""Data package."""
+188
View File
@@ -0,0 +1,188 @@
from __future__ import annotations
import warnings
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from dots_tts.utils.profiling import ensure_data_profiler
@dataclass(slots=True)
class BatchDecision:
dropped_samples: list[dict]
batch_samples: list[dict]
@dataclass(slots=True)
class _PoolSample:
sample: dict
num_audio_tokens: int
num_text_tokens: int
arrival_step: int
class OnlineBatcher:
def __init__(
self,
*,
max_audio_tokens_in_batch: int,
max_text_tokens_in_batch: int,
max_batch_size: int | None,
sample_pool_size: int,
profiler=None,
):
self.max_audio_tokens_in_batch = max(1, int(max_audio_tokens_in_batch))
self.max_text_tokens_in_batch = max(1, int(max_text_tokens_in_batch))
self.max_batch_size = max_batch_size
self.sample_pool_size = max(1, int(sample_pool_size))
self.profiler = ensure_data_profiler(profiler)
@staticmethod
def _sort_pool(pool: list[_PoolSample]) -> None:
pool.sort(
key=lambda item: (
item.num_audio_tokens,
item.num_text_tokens,
-item.arrival_step,
),
reverse=True,
)
def _choose_anchor_index(
self,
pool: list[_PoolSample],
*,
decision_step: int,
) -> int:
oldest_waiting_index = -1
oldest_waiting_step = decision_step
for index, item in enumerate(pool):
waited_steps = decision_step - item.arrival_step
if waited_steps < self.sample_pool_size:
continue
if item.arrival_step <= oldest_waiting_step:
oldest_waiting_index = index
oldest_waiting_step = item.arrival_step
return 0 if oldest_waiting_index < 0 else oldest_waiting_index
def _build_next_decision(
self,
pool: list[_PoolSample],
*,
decision_step: int,
) -> BatchDecision:
dropped_samples: list[dict] = []
batch_samples: list[dict] = []
selected_indices: list[int] = []
anchor_index = self._choose_anchor_index(pool, decision_step=decision_step)
anchor = pool[anchor_index]
exceed_audio_budget = anchor.num_audio_tokens > self.max_audio_tokens_in_batch
exceed_text_budget = anchor.num_text_tokens > self.max_text_tokens_in_batch
exceed_batch_size = self.max_batch_size is not None and self.max_batch_size < 1
if exceed_audio_budget or exceed_text_budget or exceed_batch_size:
skipped = pool.pop(anchor_index).sample
dropped_samples.append(skipped)
warnings.warn(
"Skipping sample that exceeds batching limits on its own: "
f"fid={skipped.get('fid')!r}, "
f"num_audio_tokens={anchor.num_audio_tokens}, "
f"input_ids_length={anchor.num_text_tokens}, "
f"max_audio_tokens_in_batch={self.max_audio_tokens_in_batch}, "
f"max_text_tokens_in_batch={self.max_text_tokens_in_batch}, "
f"max_batch_size={self.max_batch_size}",
RuntimeWarning,
stacklevel=2,
)
return BatchDecision(
dropped_samples=dropped_samples,
batch_samples=batch_samples,
)
longest_audio_tokens = anchor.num_audio_tokens
longest_text_tokens = anchor.num_text_tokens
batch_samples.append(anchor.sample)
selected_indices.append(anchor_index)
for index, item in enumerate(pool):
if index == anchor_index:
continue
if (
self.max_batch_size is not None
and len(batch_samples) >= self.max_batch_size
):
break
proposed_batch_size = len(batch_samples) + 1
proposed_longest_audio_tokens = max(
longest_audio_tokens,
item.num_audio_tokens,
)
proposed_longest_text_tokens = max(
longest_text_tokens,
item.num_text_tokens,
)
if (
proposed_longest_audio_tokens * proposed_batch_size
> self.max_audio_tokens_in_batch
):
continue
if (
proposed_longest_text_tokens * proposed_batch_size
> self.max_text_tokens_in_batch
):
continue
batch_samples.append(item.sample)
selected_indices.append(index)
longest_audio_tokens = proposed_longest_audio_tokens
longest_text_tokens = proposed_longest_text_tokens
for index in sorted(set(selected_indices), reverse=True):
pool.pop(index)
return BatchDecision(
dropped_samples=dropped_samples,
batch_samples=batch_samples,
)
def build_decisions(self, sample_iter: Iterable[dict]) -> Iterator[BatchDecision]:
pool: list[_PoolSample] = []
source_exhausted = False
decision_step = 0
iterator = iter(sample_iter)
while not source_exhausted or pool:
while not source_exhausted and len(pool) < self.sample_pool_size:
try:
sample = next(iterator)
except StopIteration:
source_exhausted = True
break
pool.append(
_PoolSample(
sample=sample,
num_audio_tokens=int(sample.get("num_audio_tokens", 0)),
num_text_tokens=int(sample.get("input_ids_length", 0)),
arrival_step=decision_step,
)
)
if not pool:
break
profiler = self.profiler
with profiler.measure("main.sort_pool", count=len(pool)):
self._sort_pool(pool)
with profiler.measure("main.build_batch_decision"):
decision = self._build_next_decision(
pool,
decision_step=decision_step,
)
if decision.dropped_samples or decision.batch_samples:
decision_step += 1
yield decision
continue
raise RuntimeError("OnlineBatcher failed to make progress on a non-empty pool.")
+194
View File
@@ -0,0 +1,194 @@
from __future__ import annotations
from torch.utils.data import DataLoader
from dots_tts.config.data import DataConfig
from dots_tts.data.pipelines.base import BaseSamplePipeline
from dots_tts.data.pipelines.tts_pipeline import BasicTtsPipeline, InterleaveTtsPipeline
from dots_tts.data.source_adapters.jsonl_manifest_adapter import (
JsonlManifestSourceAdapter,
)
from dots_tts.data.source_adapters.multi_source_adapter import (
SequentialMultiSourceAdapter,
SourceSpec,
WeightedMultiSourceAdapter,
)
from dots_tts.data.streaming import (
BatchedDataStream,
StreamingSampleDataset,
identity_collate,
)
_SOURCE_ADAPTER_CLASSES = {
"JsonlManifestSourceAdapter": JsonlManifestSourceAdapter,
}
def _build_source_pipeline(
tokenizer, data_cfg, pipeline_name: str, *, profiler=None
) -> BaseSamplePipeline:
if pipeline_name == "basic":
return BasicTtsPipeline(tokenizer, data_cfg, profiler=profiler)
if pipeline_name == "interleave":
return InterleaveTtsPipeline(tokenizer, data_cfg, profiler=profiler)
raise ValueError(f"Unsupported data pipeline: {pipeline_name!r}")
def _build_source_specs(data_cfg, tokenizer, *, profiler=None) -> list[SourceSpec]:
specs = []
for source_cfg in data_cfg.sources:
adapter_cls = _SOURCE_ADAPTER_CLASSES[source_cfg.adapter.class_name]
adapter = adapter_cls(**source_cfg.adapter.params)
specs.append(
SourceSpec(
name=source_cfg.name,
weight=float(source_cfg.weight),
adapter=adapter,
pipeline=_build_source_pipeline(
tokenizer, data_cfg, source_cfg.pipeline, profiler=profiler
),
)
)
return specs
def _resolve_rank_info(accelerator=None) -> tuple[int, int]:
rank = (
int(getattr(accelerator, "process_index", 0)) if accelerator is not None else 0
)
world_size = (
int(getattr(accelerator, "num_processes", 1)) if accelerator is not None else 1
)
return rank, world_size
def _local_num_tokens_per_epoch(
global_num_tokens_per_epoch: int, *, rank: int, world_size: int
) -> int:
if world_size <= 0:
raise ValueError(f"world_size must be positive, but got {world_size}.")
if rank < 0 or rank >= world_size:
raise ValueError(
f"rank must be in [0, {world_size}), but got rank={rank}."
)
base, remainder = divmod(int(global_num_tokens_per_epoch), int(world_size))
return base + int(rank < remainder)
def _build_dataset(
data_cfg: DataConfig,
*,
tokenizer,
seed: int,
accelerator=None,
sequential: bool,
profiler=None,
):
rank, world_size = _resolve_rank_info(accelerator)
source_cls = SequentialMultiSourceAdapter if sequential else WeightedMultiSourceAdapter
source = source_cls(
sources=_build_source_specs(data_cfg, tokenizer, profiler=profiler)
)
return StreamingSampleDataset(
source=source,
rank=rank,
world_size=world_size,
seed=int(seed),
)
def build_training_dataset(
data_cfg: DataConfig,
tokenizer,
*,
seed: int,
accelerator=None,
profiler=None,
):
if data_cfg.num_tokens_per_epoch is None:
raise ValueError("Training data requires num_tokens_per_epoch.")
return _build_dataset(
data_cfg,
tokenizer=tokenizer,
seed=seed,
accelerator=accelerator,
sequential=False,
profiler=profiler,
)
def build_validation_dataset(
data_cfg: DataConfig,
tokenizer,
*,
seed: int,
accelerator=None,
profiler=None,
):
return _build_dataset(
data_cfg,
tokenizer=tokenizer,
seed=seed,
accelerator=accelerator,
sequential=True,
profiler=profiler,
)
def _build_sample_loader(dataset, data_cfg: DataConfig) -> DataLoader:
loader_kwargs = {
"dataset": dataset,
"batch_size": None,
"collate_fn": identity_collate,
"num_workers": data_cfg.num_workers,
"pin_memory": data_cfg.pin_memory,
"persistent_workers": data_cfg.num_workers > 0,
}
if data_cfg.num_workers > 0:
loader_kwargs["prefetch_factor"] = int(data_cfg.prefetch_factor)
sample_loader = DataLoader(**loader_kwargs)
return sample_loader
def build_training_dataloader(
dataset, data_cfg: DataConfig, tokenizer, *, profiler=None
):
local_num_tokens_per_epoch = _local_num_tokens_per_epoch(
int(data_cfg.num_tokens_per_epoch),
rank=int(dataset.rank),
world_size=int(dataset.world_size),
)
sample_loader = _build_sample_loader(dataset, data_cfg)
batched_stream = BatchedDataStream(
sample_dataset=dataset,
data_cfg=data_cfg,
tokenizer=tokenizer,
num_tokens_per_epoch=local_num_tokens_per_epoch,
profiler=profiler,
)
batched_stream.attach_loader(sample_loader)
return batched_stream
def build_validation_dataloader(
dataset, data_cfg: DataConfig, tokenizer, *, profiler=None
):
sample_loader = _build_sample_loader(dataset, data_cfg)
batched_stream = BatchedDataStream(
sample_dataset=dataset,
data_cfg=data_cfg,
tokenizer=tokenizer,
num_tokens_per_epoch=None,
profiler=profiler,
)
batched_stream.attach_loader(sample_loader)
return batched_stream
__all__ = [
"build_training_dataloader",
"build_training_dataset",
"build_validation_dataloader",
"build_validation_dataset",
]
+87
View File
@@ -0,0 +1,87 @@
from __future__ import annotations
from typing import Any
import torch
from torch.nn.utils.rnn import pad_sequence
class PadCollator:
def __init__(self, tokenizer):
self.tokenizer = tokenizer
self.pad_token_id = tokenizer.pad_token_id
if self.pad_token_id is None:
self.pad_token_id = tokenizer.eos_token_id or 0
def __call__(self, samples: list[dict[str, Any]]) -> dict[str, Any]:
if not samples:
raise ValueError("PadCollator received an empty sample list.")
order = sorted(
range(len(samples)),
key=lambda idx: samples[idx]["sample_length"],
reverse=True,
)
ordered = [samples[idx] for idx in order]
input_ids = [
torch.tensor(sample["input_ids"], dtype=torch.long) for sample in ordered
]
labels = [
torch.tensor(sample["labels"], dtype=torch.long) for sample in ordered
]
loss_masks = [
torch.tensor(sample["loss_mask"], dtype=torch.float32) for sample in ordered
]
waveforms = [sample["sample"].squeeze(0) for sample in ordered]
fbank = [sample["fbank"] for sample in ordered]
return {
"fids": [sample["fid"] for sample in ordered],
"source_names": [sample.get("source_name") for sample in ordered],
"input_ids": pad_sequence(
input_ids,
batch_first=True,
padding_value=self.pad_token_id,
),
"input_ids_lengths": torch.tensor(
[len(sample["input_ids"]) for sample in ordered],
dtype=torch.long,
),
"labels": pad_sequence(
labels,
batch_first=True,
padding_value=self.pad_token_id,
),
"loss_mask": pad_sequence(
loss_masks,
batch_first=True,
padding_value=0.0,
),
"sample": pad_sequence(
waveforms,
batch_first=True,
padding_value=0.0,
).unsqueeze(1),
"sample_lengths": torch.tensor(
[sample["sample_length"] for sample in ordered],
dtype=torch.long,
),
"num_text_tokens": torch.tensor(
[sample["num_text_tokens"] for sample in ordered],
dtype=torch.long,
),
"num_audio_tokens": torch.tensor(
[sample["num_audio_tokens"] for sample in ordered],
dtype=torch.long,
),
"fbank": pad_sequence(
fbank,
batch_first=True,
padding_value=0.0,
),
"fbank_lengths": torch.tensor(
[sample["fbank_length"] for sample in ordered],
dtype=torch.long,
),
}
@@ -0,0 +1 @@
"""Data pipelines package."""
@@ -0,0 +1,32 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Iterable, Iterator
class BaseSamplePipeline(ABC):
"""1:1 sample pipeline that preserves adapter resume metadata."""
@staticmethod
def _validate_input_sample(sample: dict) -> None:
if "_adapter_state" not in sample:
raise RuntimeError(
"Source sample is missing required '_adapter_state' for resume."
)
@abstractmethod
def process_sample(self, sample: dict) -> dict:
"""Transform one raw sample into one processed sample."""
def __call__(self, samples: Iterable[dict]) -> Iterator[dict]:
for raw_sample in samples:
self._validate_input_sample(raw_sample)
processed = self.process_sample(dict(raw_sample))
if not isinstance(processed, dict):
raise RuntimeError(
f"{self.__class__.__name__}.process_sample() must return a dict."
)
item = dict(raw_sample)
item.update(processed)
self._validate_input_sample(item)
yield item
@@ -0,0 +1,84 @@
from __future__ import annotations
import torch
import torch.nn.functional as F
DEFAULT_EDGE_SILENCE_MS = 250.0
DEFAULT_EDGE_SILENCE_TOP_DB = 30.0
def align_length(num_samples: int, multiple_of: int | None) -> int:
if multiple_of is None or multiple_of <= 0:
return int(num_samples)
if num_samples % multiple_of == 0:
return int(num_samples)
return int(((num_samples + multiple_of - 1) // multiple_of) * multiple_of)
def pad_waveform_align_only(
waveform: torch.Tensor,
*,
multiple_of: int | None,
) -> torch.Tensor:
if multiple_of is None or multiple_of <= 0:
return waveform
target_length = align_length(waveform.size(-1), multiple_of)
delta = target_length - waveform.size(-1)
if delta <= 0:
return waveform
return F.pad(waveform, (0, delta), "constant", 0.0)
def normalize_edge_silence_duration(
waveform: torch.Tensor,
*,
sample_rate: int,
target_silence_duration_ms: float = DEFAULT_EDGE_SILENCE_MS,
top_db: float = DEFAULT_EDGE_SILENCE_TOP_DB,
) -> torch.Tensor:
mono_waveform = waveform[0]
target_samples = int(round(float(sample_rate) * float(target_silence_duration_ms) / 1000.0))
amplitude = mono_waveform.abs()
peak = float(amplitude.max().item())
if peak <= 0.0:
waveform = waveform[..., :target_samples]
current_length = int(waveform.size(-1))
if current_length < target_samples:
waveform = F.pad(waveform, (0, target_samples - current_length), "constant", 0.0)
return waveform
threshold = peak * (10.0 ** (-float(top_db) / 20.0))
non_silent = torch.nonzero(amplitude > threshold, as_tuple=False).flatten()
first_non_silent = int(non_silent[0].item())
last_non_silent = int(non_silent[-1].item())
leading_silence_samples = first_non_silent
trailing_silence_samples = int(mono_waveform.numel()) - last_non_silent - 1
leading_delta = target_samples - leading_silence_samples
if leading_delta > 0:
waveform = F.pad(waveform, (leading_delta, 0), "constant", 0.0)
else:
trim_from_start = min(-leading_delta, int(waveform.size(-1)))
waveform = waveform[..., trim_from_start:]
trailing_delta = target_samples - trailing_silence_samples
if trailing_delta > 0:
return F.pad(waveform, (0, trailing_delta), "constant", 0.0)
trim_from_end = min(-trailing_delta, int(waveform.size(-1)))
if trim_from_end <= 0:
return waveform
return waveform[..., :-trim_from_end]
def compute_num_audio_tokens(
num_samples: int, *, audio_samples_per_llm_token: int
) -> int:
if num_samples % audio_samples_per_llm_token != 0:
raise ValueError(
f"Waveform length {num_samples} is not aligned to token hop {audio_samples_per_llm_token}."
)
return num_samples // audio_samples_per_llm_token
@@ -0,0 +1,339 @@
from __future__ import annotations
from dataclasses import dataclass
import re
from typing import Any
from loguru import logger
from dots_tts.utils.tokenizer import (
AUDIO_GEN_END_TOKEN,
AUDIO_GEN_SPAN_TOKEN,
AUDIO_GEN_START_TOKEN,
TEXT_COND_END_TOKEN,
require_token_id,
)
TEMPLATE_PATTERN = re.compile(r"\{text\}|\{audio\}|\{interleave\}|[^\{]+")
@dataclass(frozen=True)
class ParsedTemplate:
parts: tuple[str, ...]
has_audio_placeholder: bool
has_interleave_placeholder: bool
@dataclass(frozen=True)
class TokenizedTemplatePart:
kind: str
token_ids: tuple[int, ...] = ()
raw_text: str | None = None
def parse_template(template: str) -> ParsedTemplate:
parts = tuple(re.findall(TEMPLATE_PATTERN, template))
has_audio_placeholder = "{audio}" in parts
interleave_count = parts.count("{interleave}")
if has_audio_placeholder and interleave_count:
raise ValueError("Template cannot mix audio and interleave placeholders.")
if interleave_count > 1:
raise ValueError(
"Interleave generation template must contain exactly one interleave placeholder."
)
return ParsedTemplate(
parts=parts,
has_audio_placeholder=has_audio_placeholder,
has_interleave_placeholder=interleave_count == 1,
)
def _prepare_template_tokens(
*, text: str, tokenizer, template: str
) -> tuple[ParsedTemplate, list[int]]:
return parse_template(template), tokenizer.encode(text, add_special_tokens=False)
def _iter_tokenized_template_parts(
*,
parsed_template: ParsedTemplate,
tokenizer,
text_tokens: list[int],
):
for part in parsed_template.parts:
if part == "{text}":
yield TokenizedTemplatePart(kind="text", token_ids=tuple(text_tokens))
continue
if part == "{audio}":
yield TokenizedTemplatePart(kind="audio")
continue
if part == "{interleave}":
yield TokenizedTemplatePart(kind="interleave")
continue
yield TokenizedTemplatePart(
kind="literal",
token_ids=tuple(tokenizer.encode(part, add_special_tokens=False)),
raw_text=part,
)
def _extend_tokens_with_loss(
*, full_ids: list[int], loss_mask: list[float], token_ids: tuple[int, ...], loss: float
) -> None:
full_ids.extend(token_ids)
loss_mask.extend([loss] * len(token_ids))
def build_tokenized_example(
*, text: str, tokenizer, template: str, num_audio_tokens: int
) -> dict[str, Any]:
if tokenizer.eos_token_id is None:
raise ValueError("Tokenizer eos_token_id is required for generation targets.")
parsed_template, text_tokens = _prepare_template_tokens(
text=text,
tokenizer=tokenizer,
template=template,
)
full_ids: list[int] = []
loss_mask: list[float] = []
audio_tokens: list[int] | None = None
if parsed_template.has_audio_placeholder:
audio_gen_start_id = require_token_id(tokenizer, AUDIO_GEN_START_TOKEN)
audio_gen_span_id = require_token_id(tokenizer, AUDIO_GEN_SPAN_TOKEN)
audio_gen_end_id = require_token_id(tokenizer, AUDIO_GEN_END_TOKEN)
audio_tokens = (
[audio_gen_start_id]
+ [audio_gen_span_id] * num_audio_tokens
+ [audio_gen_end_id]
)
elif parsed_template.has_interleave_placeholder:
audio_gen_span_id = require_token_id(tokenizer, AUDIO_GEN_SPAN_TOKEN)
audio_gen_end_id = require_token_id(tokenizer, AUDIO_GEN_END_TOKEN)
text_cond_end_id = require_token_id(tokenizer, TEXT_COND_END_TOKEN)
for part in _iter_tokenized_template_parts(
parsed_template=parsed_template,
tokenizer=tokenizer,
text_tokens=text_tokens,
):
if part.kind == "text":
_extend_tokens_with_loss(
full_ids=full_ids,
loss_mask=loss_mask,
token_ids=part.token_ids,
loss=0.0,
)
continue
if part.kind == "audio":
if audio_tokens is None:
raise RuntimeError("Audio placeholder tokens were not initialized.")
full_ids.extend(audio_tokens)
loss_mask.extend([0.0])
loss_mask.extend([1.0] * max(0, len(audio_tokens) - 2))
loss_mask.append(0.0)
continue
if part.kind == "interleave":
_append_interleave_generation_tokens(
full_ids=full_ids,
loss_mask=loss_mask,
text_tokens=text_tokens,
num_audio_tokens=num_audio_tokens,
audio_span_id=audio_gen_span_id,
audio_end_id=audio_gen_end_id,
text_cond_end_id=text_cond_end_id,
)
continue
_extend_tokens_with_loss(
full_ids=full_ids,
loss_mask=loss_mask,
token_ids=part.token_ids,
loss=0.0,
)
full_ids.append(tokenizer.eos_token_id)
loss_mask.append(0.0)
return {
"input_ids": full_ids[:-1],
"labels": full_ids[1:],
"loss_mask": loss_mask[1:],
"text_token_count": len(text_tokens),
}
def build_generation_schedule(
*,
text: str,
tokenizer,
template: str,
max_audio_tokens: int,
) -> dict[str, Any]:
if max_audio_tokens <= 0:
raise ValueError("max_audio_tokens must be positive for generation.")
parsed_template, text_tokens = _prepare_template_tokens(
text=text,
tokenizer=tokenizer,
template=template,
)
schedule_ids: list[int] = []
audio_gen_start_id = require_token_id(tokenizer, AUDIO_GEN_START_TOKEN)
audio_gen_span_id = require_token_id(tokenizer, AUDIO_GEN_SPAN_TOKEN)
if parsed_template.has_audio_placeholder:
for part in _iter_tokenized_template_parts(
parsed_template=parsed_template,
tokenizer=tokenizer,
text_tokens=text_tokens,
):
if part.kind == "audio":
schedule_ids.append(audio_gen_start_id)
schedule_ids.extend([audio_gen_span_id] * max_audio_tokens)
continue
schedule_ids.extend(part.token_ids)
visible_schedule_ids = [
token_id for token_id in schedule_ids if token_id != audio_gen_span_id
]
decoded_schedule = (
tokenizer.decode(
visible_schedule_ids,
skip_special_tokens=False,
clean_up_tokenization_spaces=False,
)
if hasattr(tokenizer, "decode")
else repr(visible_schedule_ids)
)
logger.info(
"Built generation schedule: interleave={} max_audio_tokens={} sequence={!r}",
False,
int(max_audio_tokens),
decoded_schedule,
)
return {
"schedule_ids": schedule_ids,
"interleave": False,
}
if not parsed_template.has_interleave_placeholder:
raise ValueError(
"Generation template must contain either {audio} or {interleave}."
)
text_cond_end_id = require_token_id(tokenizer, TEXT_COND_END_TOKEN)
if max_audio_tokens < len(text_tokens):
raise ValueError(
"Interleave generation requires at least one audio span per text token: "
f"text_token_count={len(text_tokens)} "
f"max_audio_patch_count={max_audio_tokens}."
)
interleave_started = False
for part in _iter_tokenized_template_parts(
parsed_template=parsed_template,
tokenizer=tokenizer,
text_tokens=text_tokens,
):
if part.kind == "interleave":
_append_interleave_schedule_tokens(
schedule_ids=schedule_ids,
text_tokens=text_tokens,
max_audio_tokens=max_audio_tokens,
audio_span_id=audio_gen_span_id,
text_cond_end_id=text_cond_end_id,
)
interleave_started = True
continue
if part.kind == "text":
raise ValueError(
"Generation schedule does not support {text} inside an interleave template."
)
if part.kind == "audio":
raise ValueError(
"Generation schedule does not support {audio} inside an interleave template."
)
if interleave_started:
if (part.raw_text or "").strip():
raise ValueError(
"Generation schedule does not support non-empty suffix text after the interleave placeholder."
)
continue
schedule_ids.extend(part.token_ids)
visible_schedule_ids = [
token_id for token_id in schedule_ids if token_id != audio_gen_span_id
]
decoded_schedule = (
tokenizer.decode(
visible_schedule_ids,
skip_special_tokens=False,
clean_up_tokenization_spaces=False,
)
if hasattr(tokenizer, "decode")
else repr(visible_schedule_ids)
)
logger.info(
"Built generation schedule: interleave={} max_audio_tokens={} sequence={!r}",
True,
int(max_audio_tokens),
decoded_schedule,
)
return {
"schedule_ids": schedule_ids,
"interleave": True,
}
def _append_interleave_generation_tokens(
*,
full_ids: list[int],
loss_mask: list[float],
text_tokens: list[int],
num_audio_tokens: int,
audio_span_id: int,
audio_end_id: int,
text_cond_end_id: int,
) -> None:
audio_tokens = [audio_span_id] * num_audio_tokens + [audio_end_id]
text_index = 0
audio_index = 0
text_cond_end_added = False
while text_index < len(text_tokens) or audio_index < len(audio_tokens):
if text_index < len(text_tokens):
full_ids.append(text_tokens[text_index])
loss_mask.append(0.0)
text_index += 1
elif not text_cond_end_added:
full_ids.append(text_cond_end_id)
loss_mask.append(0.0)
text_cond_end_added = True
if audio_index < len(audio_tokens):
full_ids.append(audio_tokens[audio_index])
loss_mask.append(1.0 if audio_index < num_audio_tokens else 0.0)
audio_index += 1
if not text_cond_end_added:
full_ids.append(text_cond_end_id)
loss_mask.append(0.0)
def _append_interleave_schedule_tokens(
*,
schedule_ids: list[int],
text_tokens: list[int],
max_audio_tokens: int,
audio_span_id: int,
text_cond_end_id: int,
) -> None:
for token_id in text_tokens:
schedule_ids.append(token_id)
schedule_ids.append(audio_span_id)
schedule_ids.append(text_cond_end_id)
remaining_audio_tokens = max_audio_tokens - len(text_tokens)
if remaining_audio_tokens > 0:
schedule_ids.extend([audio_span_id] * remaining_audio_tokens)
@@ -0,0 +1,132 @@
from __future__ import annotations
import soundfile as sf
import torch
from dots_tts.utils.profiling import ensure_data_profiler
from dots_tts.data.pipelines.base import BaseSamplePipeline
from dots_tts.data.pipelines.preprocessing import (
compute_num_audio_tokens,
normalize_edge_silence_duration,
pad_waveform_align_only,
)
from dots_tts.data.pipelines.tokenizing import build_tokenized_example
from dots_tts.modules.speaker.fbank import extract_speaker_fbank
from dots_tts.utils.audio import high_quality_resample
TTS_TEXT_PREFIX = "[文本]"
TTS_AUDIO_PREFIX = "[文本对应语音]"
TTS_INSTRUCTION_TEXT_PREFIX = "[带指令文本]"
TTA_TEXT_PREFIX = "[声音描述]"
TTA_AUDIO_PREFIX = "[描述对应声音]"
TTS_INTERLEAVE_PREFIX = "[流式语音合成]"
DEFAULT_TRAIN_TEMPLATE = f"{TTS_TEXT_PREFIX}{{text}}{TTS_AUDIO_PREFIX}{{audio}}"
DEFAULT_INSTRUCTION_TTS_TEMPLATE = (
f"{TTS_INSTRUCTION_TEXT_PREFIX}{{text}}{TTS_AUDIO_PREFIX}{{audio}}"
)
DEFAULT_TEXT_TO_AUDIO_TEMPLATE = f"{TTA_TEXT_PREFIX}{{text}}{TTA_AUDIO_PREFIX}{{audio}}"
DEFAULT_INTERLEAVE_TRAIN_TEMPLATE = f"{TTS_INTERLEAVE_PREFIX}{{interleave}}"
class BasicTtsPipeline(BaseSamplePipeline):
"""Fixed internal training pipeline for adapter-emitted samples."""
template = DEFAULT_TRAIN_TEMPLATE
def __init__(self, tokenizer, data_cfg, *, profiler=None):
self.tokenizer = tokenizer
self.train_audio_sample_rate = int(data_cfg.train_audio_sample_rate)
self.audio_samples_per_llm_token = int(data_cfg.audio_samples_per_llm_token)
self.profiler = ensure_data_profiler(profiler)
@staticmethod
def _load_waveform(audio_path: str) -> tuple[torch.Tensor, int]:
if not isinstance(audio_path, str):
raise TypeError(
f"Training audio must be a filesystem path, got {type(audio_path)}."
)
audio_data, sample_rate = sf.read(
audio_path,
dtype="float32",
always_2d=True,
)
waveform = torch.from_numpy(audio_data.T)
if waveform.size(0) > 1:
waveform = waveform.mean(dim=0, keepdim=True)
return waveform.contiguous(), int(sample_rate)
@staticmethod
def _validate_source_sample(sample: dict) -> None:
missing = [field for field in ("fid", "text", "audio") if field not in sample]
if missing:
raise ValueError(
"Source adapter must emit fid/text/audio. "
f"Missing fields: {missing}. Sample keys: {sorted(sample.keys())}"
)
def process_sample(self, raw_sample: dict) -> dict:
sample = dict(raw_sample)
self._validate_source_sample(sample)
sample["fid"] = str(sample["fid"])
with self.profiler.measure("worker.process_sample_total"):
return self._process_sample_impl(sample)
def _process_sample_impl(self, sample: dict) -> dict:
profiler = self.profiler
with profiler.measure("worker.load_audio"):
waveform, sample_rate = self._load_waveform(sample["audio"])
with profiler.measure("worker.resample_audio"):
waveform = high_quality_resample(
waveform,
orig_sr=sample_rate,
target_sr=self.train_audio_sample_rate,
)
with profiler.measure("worker.normalize_edge_silence"):
waveform = normalize_edge_silence_duration(
waveform,
sample_rate=self.train_audio_sample_rate,
)
sample["sample"] = waveform
sample["sample_rate"] = self.train_audio_sample_rate
sample["unpadded_sample_length"] = int(waveform.size(-1))
with profiler.measure("worker.pad_audio"):
waveform = pad_waveform_align_only(
waveform,
multiple_of=self.audio_samples_per_llm_token,
)
sample["sample"] = waveform
sample["sample_length"] = int(waveform.size(-1))
num_audio_tokens = compute_num_audio_tokens(
sample["sample_length"],
audio_samples_per_llm_token=self.audio_samples_per_llm_token,
)
with profiler.measure("worker.tokenize"):
tokenized = build_tokenized_example(
text=sample["text"],
tokenizer=self.tokenizer,
template=self.template,
num_audio_tokens=num_audio_tokens,
)
sample["input_ids"] = tokenized["input_ids"]
sample["labels"] = tokenized["labels"]
sample["loss_mask"] = tokenized["loss_mask"]
sample["input_ids_length"] = len(tokenized["input_ids"])
sample["num_text_tokens"] = tokenized["text_token_count"]
sample["num_audio_tokens"] = num_audio_tokens
sample["num_total_tokens"] = sample["input_ids_length"]
with profiler.measure("worker.extract_fbank"):
fbank = extract_speaker_fbank(
sample["sample"],
sample_rate=sample["sample_rate"],
)
sample["fbank"] = fbank
sample["fbank_length"] = int(fbank.size(0))
return sample
class InterleaveTtsPipeline(BasicTtsPipeline):
template = DEFAULT_INTERLEAVE_TRAIN_TEMPLATE
@@ -0,0 +1 @@
"""Source adapter package."""
@@ -0,0 +1,91 @@
from __future__ import annotations
import random
from abc import ABC, abstractmethod
from collections.abc import Iterable, Sequence
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, TypeVar
@dataclass(frozen=True)
class SourceContext:
"""Execution context for a single adapter iterator."""
epoch: int
rank: int
world_size: int
worker_id: int
num_workers: int
seed: int
@property
def global_worker_count(self) -> int:
return max(1, self.world_size * self.num_workers)
@property
def global_worker_id(self) -> int:
return self.rank * self.num_workers + self.worker_id
class BaseSourceAdapter(ABC):
"""State-aware streaming source interface used by the training pipeline."""
@abstractmethod
def initial_state(self) -> dict[str, Any]:
"""Return the default iterator state for a new worker/epoch."""
@abstractmethod
def iter_samples(
self,
context: SourceContext,
*,
state: dict[str, Any] | None = None,
) -> Iterable[dict[str, Any]]:
"""Yield raw samples and attach the next adapter state to each item."""
@abstractmethod
def is_cycle_start_state(self, state: dict[str, Any] | None) -> bool:
"""Return whether ``state`` points at the beginning of a source cycle."""
def normalize_state(self, state: dict[str, Any] | None) -> dict[str, Any]:
merged = self.initial_state()
if state:
merged.update(deepcopy(state))
return merged
def clone_state(self, state: dict[str, Any] | None) -> dict[str, Any]:
return deepcopy(self.normalize_state(state))
def advance_cycle(self, state: dict[str, Any] | None) -> dict[str, Any]:
raise RuntimeError(
f"{self.__class__.__name__} does not support repeated cycling."
)
_T = TypeVar("_T")
class ShardableSourceAdapter(BaseSourceAdapter):
"""Helper mixin for deterministic rank/worker sharding."""
@staticmethod
def is_assigned_index(index: int, context: SourceContext) -> bool:
return index % context.global_worker_count == context.global_worker_id
@staticmethod
def shard_items(
items: Sequence[_T],
context: SourceContext,
*,
shuffle: bool = False,
seed_offset: int = 0,
) -> list[_T]:
assigned = list(items)
if shuffle:
random.Random(context.seed + context.epoch + seed_offset).shuffle(assigned)
return [
item
for index, item in enumerate(assigned)
if ShardableSourceAdapter.is_assigned_index(index, context)
]
@@ -0,0 +1,132 @@
from __future__ import annotations
import json
import random
from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import Any
from dots_tts.data.source_adapters.base_adapter import (
BaseSourceAdapter,
ShardableSourceAdapter,
SourceContext,
)
class JsonlManifestSourceAdapter(ShardableSourceAdapter, BaseSourceAdapter):
"""Finite adapter for line-delimited JSON manifests."""
def __init__(
self,
*,
manifest_path: str,
fid_key: str = "fid",
text_key: str = "text",
audio_key: str = "audio",
shuffle: bool = False,
encoding: str = "utf-8",
):
self.manifest_path = Path(manifest_path)
self.fid_key = fid_key
self.text_key = text_key
self.audio_key = audio_key
self.shuffle = shuffle
self.encoding = encoding
self._records: list[dict[str, Any]] | None = None
def initial_state(self) -> dict[str, Any]:
return {"cycle": 0, "cursor": 0}
def is_cycle_start_state(self, state: dict[str, Any] | None) -> bool:
normalized = self.normalize_state(state)
return int(normalized["cursor"]) == 0
def advance_cycle(self, state: dict[str, Any] | None) -> dict[str, Any]:
normalized = self.normalize_state(state)
return {"cycle": int(normalized["cycle"]) + 1, "cursor": 0}
def _iter_records(self) -> Iterator[dict[str, Any]]:
if not self.manifest_path.is_file():
raise FileNotFoundError(f"Manifest file not found: {self.manifest_path!s}")
with self.manifest_path.open("r", encoding=self.encoding) as fin:
for line_no, raw_line in enumerate(fin, start=1):
line = raw_line.strip()
if not line:
continue
try:
yield json.loads(line)
except json.JSONDecodeError as exc:
raise ValueError(
f"Invalid JSON at {self.manifest_path}:{line_no}"
) from exc
def _base_records(self) -> list[dict[str, Any]]:
if self._records is None:
self._records = list(self._iter_records())
return self._records
def _build_sample(self, record: dict[str, Any]) -> dict[str, Any]:
missing = [
key
for key in (self.fid_key, self.text_key, self.audio_key)
if key not in record
]
if missing:
raise KeyError(
f"Manifest record is missing required keys {missing}: {record}"
)
sample = {
"fid": str(record[self.fid_key]),
"text": record[self.text_key],
"audio": record[self.audio_key],
}
for key, value in record.items():
if key in {self.fid_key, self.text_key, self.audio_key}:
continue
sample[key] = value
return sample
def _indices_for_cycle(
self,
context: SourceContext,
*,
cycle: int,
) -> list[int]:
indices = list(range(len(self._base_records())))
if self.shuffle:
random.Random(context.seed + context.epoch + 1009 * int(cycle)).shuffle(
indices
)
indices = [
record_index
for shuffled_index, record_index in enumerate(indices)
if self.is_assigned_index(shuffled_index, context)
]
else:
indices = [
record_index
for record_index in indices
if self.is_assigned_index(record_index, context)
]
return indices
def iter_samples(
self,
context: SourceContext,
*,
state: dict[str, Any] | None = None,
) -> Iterable[dict[str, Any]]:
live_state = self.normalize_state(state)
cycle = int(live_state["cycle"])
cursor = int(live_state["cursor"])
records = self._base_records()
indices = self._indices_for_cycle(context, cycle=cycle)
for position in range(cursor, len(indices)):
sample = self._build_sample(records[indices[position]])
sample["_adapter_state"] = {
"cycle": cycle,
"cursor": position + 1,
}
yield sample
@@ -0,0 +1,222 @@
from __future__ import annotations
from collections.abc import Iterable
from copy import deepcopy
from dataclasses import dataclass
from dots_tts.data.pipelines.base import BaseSamplePipeline
from dots_tts.data.source_adapters.base_adapter import (
BaseSourceAdapter,
SourceContext,
)
@dataclass(frozen=True)
class SourceSpec:
name: str
weight: float
adapter: BaseSourceAdapter
pipeline: BaseSamplePipeline
_UINT64_MASK = 0xFFFFFFFFFFFFFFFF
def _mix_uint64(value: int) -> int:
value = (value ^ (value >> 30)) * 0xBF58476D1CE4E5B9
value &= _UINT64_MASK
value = (value ^ (value >> 27)) * 0x94D049BB133111EB
value &= _UINT64_MASK
return (value ^ (value >> 31)) & _UINT64_MASK
def _stable_seed(*parts: int) -> int:
value = 0x9E3779B97F4A7C15
for part in parts:
value = (value + int(part) + 0x9E3779B97F4A7C15) & _UINT64_MASK
value = _mix_uint64(value)
return value
class SequentialMultiSourceAdapter(BaseSourceAdapter):
"""Finite adapter that concatenates sources in the configured order."""
def __init__(self, *, sources: list[SourceSpec]):
if not sources:
raise ValueError(
"SequentialMultiSourceAdapter requires at least one source."
)
self.sources = list(sources)
def initial_state(self) -> dict:
return {
"source_index": 0,
"sources": {
source.name: source.adapter.initial_state() for source in self.sources
},
}
def is_cycle_start_state(self, state: dict | None) -> bool:
normalized = self.normalize_state(state)
if int(normalized["source_index"]) != 0:
return False
return all(
source.adapter.is_cycle_start_state(normalized["sources"][source.name])
for source in self.sources
)
def normalize_state(self, state: dict | None) -> dict:
normalized = super().normalize_state(state)
source_states = normalized.get("sources") or {}
normalized["sources"] = {
source.name: source.adapter.clone_state(source_states.get(source.name))
for source in self.sources
}
normalized["source_index"] = int(normalized.get("source_index", 0))
return normalized
def clone_state(self, state: dict | None) -> dict:
return deepcopy(self.normalize_state(state))
def iter_samples(
self,
context: SourceContext,
*,
state: dict | None = None,
) -> Iterable[dict]:
live_state = self.normalize_state(state)
start_index = int(live_state["source_index"])
for index in range(start_index, len(self.sources)):
source = self.sources[index]
child_state = live_state["sources"][source.name]
raw_iter = source.adapter.iter_samples(context, state=child_state)
for sample in source.pipeline(raw_iter):
item = dict(sample)
next_child_state = item.pop("_adapter_state", None)
if next_child_state is None:
raise RuntimeError(
f"{source.adapter.__class__.__name__} must attach '_adapter_state' to samples."
)
live_state["source_index"] = index
live_state["sources"][source.name] = source.adapter.clone_state(
next_child_state
)
item["source_name"] = source.name
item["_adapter_state"] = self.clone_state(live_state)
yield item
live_state["source_index"] = index + 1
class WeightedMultiSourceAdapter(BaseSourceAdapter):
"""Infinite weighted sampler that cycles each child source independently."""
def __init__(self, *, sources: list[SourceSpec]):
if not sources:
raise ValueError("WeightedMultiSourceAdapter requires at least one source.")
invalid = [source.name for source in sources if float(source.weight) <= 0.0]
if invalid:
raise ValueError(f"Source weights must be positive: {invalid}")
self.sources = list(sources)
self._cumulative_weights = []
total = 0.0
for source in self.sources:
total += float(source.weight)
self._cumulative_weights.append(total)
self._total_weight = total
def initial_state(self) -> dict:
return {
"draw_count": 0,
"sources": {
source.name: source.adapter.initial_state() for source in self.sources
},
}
def is_cycle_start_state(self, state: dict | None) -> bool:
normalized = self.normalize_state(state)
if int(normalized["draw_count"]) != 0:
return False
return all(
source.adapter.is_cycle_start_state(normalized["sources"][source.name])
for source in self.sources
)
def normalize_state(self, state: dict | None) -> dict:
normalized = super().normalize_state(state)
source_states = normalized.get("sources") or {}
normalized["sources"] = {
source.name: source.adapter.clone_state(source_states.get(source.name))
for source in self.sources
}
normalized["draw_count"] = int(normalized.get("draw_count", 0))
return normalized
def clone_state(self, state: dict | None) -> dict:
return deepcopy(self.normalize_state(state))
def _source_draw_value(self, context: SourceContext, draw_count: int) -> float:
raw = _stable_seed(
context.seed,
context.epoch,
context.rank,
context.worker_id,
draw_count,
)
return (raw / float(1 << 64)) * self._total_weight
def _pick_source(self, context: SourceContext, draw_count: int) -> SourceSpec:
draw_value = self._source_draw_value(context, draw_count)
for source, upper in zip(self.sources, self._cumulative_weights, strict=True):
if draw_value < upper:
return source
return self.sources[-1]
def iter_samples(
self,
context: SourceContext,
*,
state: dict | None = None,
) -> Iterable[dict]:
live_state = self.normalize_state(state)
iterators: dict[str, object] = {}
while True:
draw_count = int(live_state["draw_count"])
source = self._pick_source(context, draw_count)
while True:
child_state = live_state["sources"][source.name]
child_iter = iterators.get(source.name)
if child_iter is None:
raw_iter = source.adapter.iter_samples(context, state=child_state)
child_iter = iter(source.pipeline(raw_iter))
iterators[source.name] = child_iter
try:
sample = dict(next(child_iter))
except StopIteration:
if source.adapter.is_cycle_start_state(child_state):
raise RuntimeError(
"Weighted source yielded no samples for this worker. "
f"source={source.name!r}, worker={context.global_worker_id}, "
f"epoch={context.epoch}"
)
iterators.pop(source.name, None)
live_state["sources"][source.name] = source.adapter.advance_cycle(
child_state
)
continue
next_child_state = sample.pop("_adapter_state", None)
if next_child_state is None:
raise RuntimeError(
f"{source.adapter.__class__.__name__} must attach '_adapter_state' to samples."
)
live_state["sources"][source.name] = source.adapter.clone_state(
next_child_state
)
live_state["draw_count"] = draw_count + 1
sample["source_name"] = source.name
sample["_adapter_state"] = self.clone_state(live_state)
yield sample
break
+400
View File
@@ -0,0 +1,400 @@
from __future__ import annotations
import math
import multiprocessing as mp
from collections.abc import Iterable
from copy import deepcopy
from torch.utils.data import DataLoader, IterableDataset, get_worker_info
from dots_tts.data.batchers import OnlineBatcher
from dots_tts.utils.profiling import ensure_data_profiler
from dots_tts.data.source_adapters.base_adapter import BaseSourceAdapter, SourceContext
_TRACKING_KEY = "__tracking_state__"
_RESUME_TOPOLOGY_KEY = "resume_topology"
def identity_collate(sample):
return sample
class StreamingSampleDataset(IterableDataset):
def __init__(
self,
*,
source: BaseSourceAdapter,
rank: int,
world_size: int,
seed: int,
):
self.source = source
self.rank = int(rank)
self.world_size = int(world_size)
self.seed = int(seed)
self._epoch = mp.Value("q", 0)
self._pending_resume_state: dict | None = None
def load_state_dict(self, state: dict | None) -> None:
self._pending_resume_state = deepcopy(state) if state else None
def set_epoch(self, epoch: int) -> None:
with self._epoch.get_lock():
self._epoch.value = int(epoch)
def _current_epoch(self) -> int:
with self._epoch.get_lock():
return int(self._epoch.value)
def _take_resume_state(self, epoch: int) -> dict | None:
if (
self._pending_resume_state is None
or int(self._pending_resume_state.get("epoch", -1)) != int(epoch)
):
return None
state = deepcopy(self._pending_resume_state)
self._pending_resume_state = None
return state
@staticmethod
def _validate_resume_topology(
resume_state: dict,
*,
context: SourceContext,
loader_num_workers: int,
) -> None:
resume_topology = resume_state.get(_RESUME_TOPOLOGY_KEY)
if not isinstance(resume_topology, dict):
raise RuntimeError(
"Resume state is missing required worker topology metadata."
)
expected_world_size = int(resume_topology["world_size"])
expected_num_workers = int(resume_topology["loader_num_workers"])
expected_global_worker_count = int(resume_topology["global_worker_count"])
current_num_workers = int(loader_num_workers)
current_global_worker_count = int(context.global_worker_count)
if (
expected_world_size != int(context.world_size)
or expected_num_workers != current_num_workers
or expected_global_worker_count != current_global_worker_count
):
raise RuntimeError(
"Resume requires the same data worker topology as the saved state. "
f"saved(world_size={expected_world_size}, "
f"num_workers_per_rank={expected_num_workers}, "
f"global_worker_count={expected_global_worker_count}), "
f"current(world_size={context.world_size}, "
f"num_workers_per_rank={current_num_workers}, "
f"global_worker_count={current_global_worker_count})."
)
def __iter__(self) -> Iterable[dict]:
worker_info = get_worker_info()
if worker_info is None:
worker_id = 0
loader_num_workers = 0
effective_num_workers = 1
else:
worker_id = worker_info.id
loader_num_workers = worker_info.num_workers
effective_num_workers = worker_info.num_workers
epoch = self._current_epoch()
context = SourceContext(
epoch=epoch,
rank=self.rank,
world_size=self.world_size,
worker_id=worker_id,
num_workers=effective_num_workers,
seed=self.seed,
)
resume_state = self._take_resume_state(epoch)
if resume_state is not None:
self._validate_resume_topology(
resume_state,
context=context,
loader_num_workers=loader_num_workers,
)
worker_state = (
None
if resume_state is None
else (resume_state.get("workers") or {}).get(str(context.global_worker_id))
)
sample_iter = self.source.iter_samples(
context,
state=None if worker_state is None else worker_state.get("adapter_state"),
)
for sample in sample_iter:
sample["data_worker_id"] = context.worker_id
sample["data_global_worker_id"] = context.global_worker_id
yield sample
class _DataStateTracker:
def __init__(self, *, num_tokens_per_epoch: int | None):
self.num_tokens_per_epoch = (
None if num_tokens_per_epoch is None else int(num_tokens_per_epoch)
)
self._pending_state: dict | None = None
self._reset_for_epoch(epoch=0)
def _reset_for_epoch(self, *, epoch: int) -> None:
self.epoch = int(epoch)
self.samples_emitted = 0
self.num_text_tokens = 0
self.num_audio_tokens = 0
self.num_total_tokens = 0
self.workers: dict[str, dict] = {}
self._next_sample_order_by_worker: dict[str, int] = {}
def load_state_dict(self, state: dict | None) -> None:
self._pending_state = deepcopy(state) if state else None
def set_epoch(self, epoch: int) -> None:
if self._pending_state is not None and int(
self._pending_state.get("epoch", -1)
) == int(epoch):
state = deepcopy(self._pending_state)
self._pending_state = None
self.epoch = int(state.get("epoch", epoch))
self.samples_emitted = int(state.get("samples_emitted", 0))
self.num_text_tokens = int(state.get("num_text_tokens", 0))
self.num_audio_tokens = int(state.get("num_audio_tokens", 0))
self.num_total_tokens = int(state.get("num_total_tokens", 0))
self.workers = deepcopy(state.get("workers") or {})
self._next_sample_order_by_worker = {
worker_key: int((worker_state or {}).get("sample_order", -1)) + 1
for worker_key, worker_state in self.workers.items()
}
return
self._reset_for_epoch(epoch=int(epoch))
def should_stop(self) -> bool:
return (
self.num_tokens_per_epoch is not None
and self.num_total_tokens >= self.num_tokens_per_epoch
)
def stage_sample(self, sample: dict) -> dict:
item = dict(sample)
worker_key = str(item.pop("data_global_worker_id"))
item.pop("data_worker_id", None)
adapter_state = item.pop("_adapter_state", None)
sample_order = int(self._next_sample_order_by_worker.get(worker_key, 0))
self._next_sample_order_by_worker[worker_key] = sample_order + 1
item[_TRACKING_KEY] = {
"worker_key": worker_key,
"adapter_state": deepcopy(adapter_state),
"sample_order": sample_order,
"num_text_tokens": int(item["num_text_tokens"]),
"num_audio_tokens": int(item["num_audio_tokens"]),
"num_total_tokens": int(
item.get("num_total_tokens", item["input_ids_length"])
),
}
return item
def _pop_tracking(self, sample: dict) -> tuple[dict, dict]:
item = dict(sample)
tracking = item.pop(_TRACKING_KEY, None)
if not isinstance(tracking, dict):
raise RuntimeError("Tracked sample is missing internal resume metadata.")
return item, tracking
def _advance_worker(self, tracking: dict) -> None:
adapter_state = tracking.get("adapter_state")
if adapter_state is None:
return
worker_key = str(tracking["worker_key"])
sample_order = int(tracking.get("sample_order", -1))
current_state = self.workers.get(worker_key)
current_order = int((current_state or {}).get("sample_order", -1))
if current_order >= sample_order:
return
self.workers[worker_key] = {
"adapter_state": deepcopy(adapter_state),
"sample_order": sample_order,
}
def mark_samples_dropped(self, samples: list[dict]) -> None:
for sample in samples:
_, tracking = self._pop_tracking(sample)
self._advance_worker(tracking)
def commit_batch(self, samples: list[dict]) -> list[dict]:
committed: list[dict] = []
for sample in samples:
item, tracking = self._pop_tracking(sample)
self._advance_worker(tracking)
self.samples_emitted += 1
self.num_text_tokens += int(tracking["num_text_tokens"])
self.num_audio_tokens += int(tracking["num_audio_tokens"])
self.num_total_tokens += int(tracking["num_total_tokens"])
committed.append(item)
return committed
def state_dict(self) -> dict:
return {
"epoch": int(self.epoch),
"samples_emitted": int(self.samples_emitted),
"num_text_tokens": int(self.num_text_tokens),
"num_audio_tokens": int(self.num_audio_tokens),
"num_total_tokens": int(self.num_total_tokens),
"workers": deepcopy(self.workers),
"num_tokens_per_epoch": self.num_tokens_per_epoch,
}
class BatchedDataStream:
def __init__(
self,
*,
sample_dataset: StreamingSampleDataset,
data_cfg,
tokenizer,
num_tokens_per_epoch: int | None,
profiler=None,
):
from dots_tts.data.collator import PadCollator
self.sample_dataset = sample_dataset
self.profiler = ensure_data_profiler(profiler)
llm_token_rate = (
float(data_cfg.train_audio_sample_rate)
/ float(data_cfg.audio_samples_per_llm_token)
)
self.batcher = OnlineBatcher(
max_audio_tokens_in_batch=max(
1,
math.ceil(float(data_cfg.max_audio_seconds_in_batch) * llm_token_rate),
),
max_text_tokens_in_batch=data_cfg.max_text_tokens_in_batch,
max_batch_size=data_cfg.max_samples_per_batch,
sample_pool_size=data_cfg.bucketing_pool_size,
profiler=self.profiler,
)
self.sample_loader = None
self.collator = PadCollator(tokenizer)
self.data_state = _DataStateTracker(
num_tokens_per_epoch=num_tokens_per_epoch
)
self._decision_iterator = None
self._sample_iterator = None
self._pending_batch = None
self._pending_samples = None
def attach_loader(self, loader: DataLoader) -> None:
self.sample_loader = loader
def close(self) -> None:
self._reset_iteration_state()
self.sample_loader = None
def load_state_dict(self, state: dict | None) -> None:
self.data_state.load_state_dict(state)
self.sample_dataset.load_state_dict(state)
self._reset_iteration_state()
def state_dict(self) -> dict:
if self.sample_loader is None:
raise RuntimeError("BatchedDataStream has no attached sample loader.")
if self._pending_batch is not None or self._pending_samples is not None:
raise RuntimeError(
"Cannot serialize BatchedDataStream while a batch is pending commit."
)
loader_num_workers = int(getattr(self.sample_loader, "num_workers", 0))
effective_num_workers = max(1, loader_num_workers)
state = self.data_state.state_dict()
state[_RESUME_TOPOLOGY_KEY] = {
"world_size": int(self.sample_dataset.world_size),
"loader_num_workers": loader_num_workers,
"global_worker_count": int(self.sample_dataset.world_size)
* effective_num_workers,
}
return state
def set_epoch(self, epoch: int) -> None:
self.sample_dataset.set_epoch(epoch)
self.data_state.set_epoch(epoch)
self._reset_iteration_state()
def _reset_iteration_state(self) -> None:
close_iterator = getattr(self._decision_iterator, "close", None)
if callable(close_iterator):
close_iterator()
self._decision_iterator = None
self._sample_iterator = None
self._pending_batch = None
self._pending_samples = None
def _iter_staged_samples(self):
if self.sample_loader is None:
raise RuntimeError("BatchedDataStream has no attached sample loader.")
self._sample_iterator = iter(self.sample_loader)
profiler = self.profiler
try:
while True:
if self.data_state.should_stop():
return
try:
with profiler.measure("main.loader_wait_next_sample"):
sample = next(self._sample_iterator)
except StopIteration:
return
if sample is None:
continue
with profiler.measure("main.stage_sample"):
staged = self.data_state.stage_sample(sample)
yield staged
finally:
self._sample_iterator = None
def _decision_stream(self):
if self._decision_iterator is None:
self._decision_iterator = iter(
self.batcher.build_decisions(self._iter_staged_samples())
)
return self._decision_iterator
def peek_batch(self) -> tuple[dict | None, bool]:
if self._pending_batch is not None:
return self._pending_batch, True
for decision in self._decision_stream():
if decision.dropped_samples:
self.data_state.mark_samples_dropped(decision.dropped_samples)
if not decision.batch_samples:
continue
self._pending_samples = decision.batch_samples
with self.profiler.measure(
"main.collate_batch",
count=len(decision.batch_samples),
):
self._pending_batch = self.collator(decision.batch_samples)
return self._pending_batch, True
return None, False
def commit_batch(self) -> dict:
if self._pending_batch is None or self._pending_samples is None:
raise RuntimeError("BatchedDataStream has no pending batch to commit.")
pending_batch = self._pending_batch
self.data_state.commit_batch(self._pending_samples)
self._pending_batch = None
self._pending_samples = None
return pending_batch
def discard_batch(self) -> None:
if self._pending_batch is None or self._pending_samples is None:
raise RuntimeError("BatchedDataStream has no pending batch to discard.")
self._pending_batch = None
self._pending_samples = None
def __iter__(self):
while True:
batch, has_batch = self.peek_batch()
if not has_batch:
return
self.commit_batch()
yield batch
if self.data_state.should_stop():
return
+1
View File
@@ -0,0 +1 @@
"""Model families."""
@@ -0,0 +1 @@
"""dots_tts model package."""
@@ -0,0 +1,71 @@
from __future__ import annotations
from dots_tts.config.base import ConfigBase, StrictConfigBase
from dots_tts.modules.vocoder.config import AudioVAEConfig
class _EncoderConfig(ConfigBase):
num_layers: int = 6
num_heads: int = 16
hidden_size: int = 1024
ffn_hidden_size: int = 4096
modulation: bool = False
qkv_bias: bool = False
qk_norm: bool = False
attn_dropout: float = 0.0
dropout: float = 0.0
norm_layer: str = "LayerNorm"
alibi_bias: bool = False
rotary_bias: bool = False
rotary_theta: float | None = 10000
input_dim: int = 1024
causal: bool = True
class _DiTConfig(ConfigBase):
num_layers: int = 18
num_heads: int = 16
hidden_size: int = 1024
ffn_hidden_size: int = 4096
modulation: bool = True
qkv_bias: bool = False
qk_norm: bool = False
attn_dropout: float = 0.0
dropout: float = 0.0
norm_layer: str = "LayerNorm"
alibi_bias: bool = False
rotary_bias: bool = True
rotary_theta: float | None = 10000
class LossConfig(StrictConfigBase):
ce_weight: float = 1.0
fm_weight: float = 1.0
eos_weight: float = 1.0
class MeanFlowConfig(ConfigBase):
enabled: bool = False
use_duration_embedding: bool = True
class ModelConfig(ConfigBase):
model_type: str = "dots_tts"
latent_dim: int
patch_size: int
cfg_droprate: float = 0.2
PatchEncoder: _EncoderConfig
DiT: _DiTConfig
vocoder: AudioVAEConfig
fm_sigma: float = 0.0
xvec_drop_rate: float = 0.2
campplus_embedding_size: int | None = 512
xvec_max_audio_seconds: float = 10.0
meanflow: MeanFlowConfig | None = None
__all__ = [
"LossConfig",
"MeanFlowConfig",
"ModelConfig",
]

Some files were not shown because too many files have changed in this diff Show More