Update Live-streaming code (auto-daily features)

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

After

Width:  |  Height:  |  Size: 764 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 503 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.5 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.3 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 608 B

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

After

Width:  |  Height:  |  Size: 185 KiB

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

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