first commit
@@ -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/
|
||||
@@ -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`
|
||||
@@ -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 '-'}")
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
"""Core helpers for the live streaming app."""
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
@@ -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}
|
||||
@@ -0,0 +1,965 @@
|
||||
"""
|
||||
音乐播放器 -> 直播间 UI 音乐状态同步
|
||||
|
||||
通过 Windows 10/11 的 SMTC(System 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,
|
||||
不取 playing(playing 交给音量峰值判断)。"""
|
||||
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))
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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" %*
|
||||
@@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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>
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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]"
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
After Width: | Height: | Size: 764 B |
|
After Width: | Height: | Size: 503 B |
|
After Width: | Height: | Size: 1.5 KiB |
|
After Width: | Height: | Size: 2.3 KiB |
|
After Width: | Height: | Size: 1.4 KiB |
|
After Width: | Height: | Size: 608 B |
@@ -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");
|
||||
|
||||
})();
|
||||
@@ -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": []
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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."
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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.10–3.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.
|
||||
@@ -0,0 +1 @@
|
||||
"""Application entrypoints for dots.tts."""
|
||||
@@ -0,0 +1 @@
|
||||
"""Gradio application for dots.tts."""
|
||||
@@ -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()
|
||||
@@ -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")
|
||||
@@ -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()]]
|
||||
@@ -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,
|
||||
)
|
||||
|
After Width: | Height: | Size: 185 KiB |
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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())
|
||||
@@ -0,0 +1 @@
|
||||
"""dots.tts package."""
|
||||
@@ -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())
|
||||
@@ -0,0 +1 @@
|
||||
"""Configuration package."""
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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.
|
||||
@@ -0,0 +1 @@
|
||||
"""Data package."""
|
||||
@@ -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.")
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||