hmo
7de5ab3902
- 选择文件夹/单文件 → SenseVoice 转录 → LLM 修正 → 多路线翻译 → 双语 SRT - 路线：直译中文 / 英转中 / 仅转录 - 并发流水线：转录(可调) + LLM 并发(可调)，降噪拆锁并发 - 断点续跑：.subtitle-work/ 中间产物，三阶段独立续跑 + 翻译段级续跑 - 去重：history 跟随视频文件夹，防重复任务保护 - 实时进度：SSE 推送 + 耗时显示 + 子任务状态 + 重新生成按钮 - 时间戳调优：VAD silence_schedule + noisereduce 降噪 + 完整性校验
471 lines
20 KiB
Python
471 lines
20 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""任务管理器：队列 + 状态机 + SSE 广播"""
|
||
import os, json, threading, time, uuid, shutil
|
||
from dataclasses import dataclass, field, asdict
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
from .config import OUTPUT_DIR, MAX_CONCURRENT_TASKS, LLM_CONCURRENCY, WORK_DIRNAME
|
||
from .pipeline.transcribe import transcribe, extract_and_denoise, VIDEO_EXTS
|
||
from .pipeline.fix import fix
|
||
from .pipeline.translate import translate_direct, translate_via_en
|
||
from .pipeline.srt import write_srts
|
||
|
||
# 去重记录文件：放在被处理视频的所在文件夹下
|
||
HISTORY_FILENAME = ".subtitle-history.json"
|
||
|
||
|
||
def _history_path(video_path):
|
||
"""视频/文件夹的 history 文件路径（跟视频在同一文件夹）"""
|
||
if os.path.isdir(video_path):
|
||
folder = video_path
|
||
else:
|
||
folder = os.path.dirname(video_path)
|
||
return os.path.join(folder, HISTORY_FILENAME)
|
||
|
||
|
||
def _load_history(video_path):
|
||
"""从视频所在文件夹加载 history（按 视频|路线 记录）"""
|
||
hp = _history_path(video_path)
|
||
if os.path.exists(hp):
|
||
try:
|
||
with open(hp, encoding='utf-8') as f:
|
||
return json.load(f)
|
||
except Exception:
|
||
return {}
|
||
return {}
|
||
|
||
|
||
def _save_history(video_path, history):
|
||
"""保存 history 到视频所在文件夹"""
|
||
hp = _history_path(video_path)
|
||
try:
|
||
with open(hp, 'w', encoding='utf-8') as f:
|
||
json.dump(history, f, ensure_ascii=False, indent=1)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
@dataclass
|
||
class Task:
|
||
id: str
|
||
source: str # 文件或文件夹路径
|
||
route: str # direct / via_en / transcribe_only
|
||
language: str = "ja"
|
||
status: str = "pending" # pending/transcribing/fixing/translating/srt/done/error
|
||
progress: int = 0
|
||
message: str = ""
|
||
created: float = field(default_factory=time.time)
|
||
finished: float = None
|
||
error: str = ""
|
||
files: list = field(default_factory=list) # 生成的 SRT 文件
|
||
transcript_path: str = ""
|
||
video_errors: list = field(default_factory=list) # 单视频失败记录 [{video, error}]
|
||
processed_count: int = 0 # 成功处理视频数
|
||
failed_count: int = 0 # 失败视频数
|
||
total_videos: int = 0 # 总视频数（文件夹）
|
||
started_at: float = None # 任务开始时间（总体耗时）
|
||
subtasks: dict = field(default_factory=dict) # {video: {status, progress, message, started_at, elapsed}} 子任务状态
|
||
|
||
def to_dict(self):
|
||
d = asdict(self)
|
||
# 动态耗时：运行中 = now - started_at, 完成 = finished - created
|
||
if self.finished:
|
||
d["elapsed"] = round(self.finished - self.created, 1)
|
||
elif self.started_at:
|
||
d["elapsed"] = round(time.time() - self.started_at, 1)
|
||
else:
|
||
d["elapsed"] = 0
|
||
return d
|
||
|
||
|
||
class TaskManager:
|
||
def __init__(self):
|
||
self.tasks = {} # id -> Task
|
||
self.queue = [] # 等待队列（串行执行）
|
||
self.lock = threading.Lock()
|
||
self.worker = None # 当前执行线程
|
||
self.subscribers = [] # SSE 客户端队列
|
||
# 并发控制
|
||
self.transcribe_lock = threading.BoundedSemaphore(MAX_CONCURRENT_TASKS) # 转录并发（GPU 模型共享，默认 2）
|
||
self.llm_sem = threading.BoundedSemaphore(LLM_CONCURRENCY) # LLM 并发信号量
|
||
self.llm_concurrency = LLM_CONCURRENCY
|
||
self.transcribe_concurrency = MAX_CONCURRENT_TASKS
|
||
|
||
def set_llm_concurrency(self, n):
|
||
"""动态调整 LLM 并发数（重建信号量）"""
|
||
n = max(1, min(int(n), 20))
|
||
self.llm_concurrency = n
|
||
self.llm_sem = threading.BoundedSemaphore(n)
|
||
return n
|
||
|
||
def set_transcribe_concurrency(self, n):
|
||
"""动态调整转录并发数（1-4，重建信号量）"""
|
||
n = max(1, min(int(n), 4))
|
||
self.transcribe_concurrency = n
|
||
self.transcribe_lock = threading.BoundedSemaphore(n)
|
||
return n
|
||
|
||
# ---- SSE ----
|
||
def subscribe(self, q):
|
||
self.subscribers.append(q)
|
||
|
||
def unsubscribe(self, q):
|
||
if q in self.subscribers:
|
||
self.subscribers.remove(q)
|
||
|
||
def _broadcast(self, task):
|
||
ev = json.dumps({"type": "task_update", "task": task.to_dict()}, ensure_ascii=False)
|
||
for q in list(self.subscribers):
|
||
try:
|
||
q.put(ev)
|
||
except Exception:
|
||
pass
|
||
|
||
# ---- 任务管理 ----
|
||
def create_task(self, source, route, language="ja"):
|
||
# 防重复：同 source+route 已有 pending/running 任务则拒绝（避免重复处理）
|
||
with self.lock:
|
||
for t in self.tasks.values():
|
||
if t.source == source and t.route == route and t.status in ("pending", "transcribing", "fixing", "translating", "srt"):
|
||
return None, [], f"该内容正在处理中（任务 {t.id[:8]}），请等待完成后再提交"
|
||
|
||
# 去重检查：文件夹 → 找出未处理/处理过的视频清单
|
||
skipped = []
|
||
if os.path.isdir(source):
|
||
history = _load_history(source)
|
||
for f in sorted(os.listdir(source)):
|
||
if f.lower().endswith(VIDEO_EXTS):
|
||
key = os.path.join(source, f)
|
||
hkey = f"{key}|{route}"
|
||
if hkey in history:
|
||
skipped.append(key)
|
||
if skipped and len(skipped) == sum(1 for f in os.listdir(source) if f.lower().endswith(VIDEO_EXTS)):
|
||
# 全部已处理过，拒绝创建
|
||
return None, skipped, ""
|
||
elif os.path.isfile(source):
|
||
# 单文件也检查：同路线已处理过则拒绝
|
||
history = _load_history(source)
|
||
hkey = f"{source}|{route}"
|
||
if hkey in history:
|
||
return None, [source], ""
|
||
|
||
task = Task(id=uuid.uuid4().hex[:12], source=source, route=route, language=language)
|
||
with self.lock:
|
||
self.tasks[task.id] = task
|
||
self.queue.append(task.id)
|
||
self._broadcast(task)
|
||
self._maybe_start()
|
||
return task, skipped, ""
|
||
|
||
def _maybe_start(self):
|
||
with self.lock:
|
||
if self.worker and self.worker.is_alive():
|
||
return
|
||
running = [t for t in self.tasks.values() if t.status in ("transcribing", "fixing", "translating", "srt")]
|
||
if len(running) >= MAX_CONCURRENT_TASKS:
|
||
return
|
||
if not self.queue:
|
||
return
|
||
task_id = self.queue.pop(0)
|
||
task = self.tasks[task_id]
|
||
task.status = "transcribing"
|
||
task.message = "开始处理"
|
||
self.worker = threading.Thread(target=self._run_task, args=(task,), daemon=True)
|
||
self.worker.start()
|
||
self._broadcast(task)
|
||
|
||
def _run_task(self, task):
|
||
try:
|
||
self._process(task)
|
||
task.status = "done"
|
||
task.progress = 100
|
||
if task.failed_count > 0:
|
||
task.message = f"部分完成: 成功 {task.processed_count}, 失败 {task.failed_count}"
|
||
else:
|
||
task.message = f"完成 ({task.processed_count} 个视频)"
|
||
task.finished = time.time()
|
||
except Exception as ex:
|
||
task.status = "error"
|
||
task.error = str(ex)
|
||
task.message = f"失败: {ex}"
|
||
task.finished = time.time()
|
||
self._broadcast(task)
|
||
# 处理下一个
|
||
self._maybe_start()
|
||
|
||
def _update(self, task, msg, pct=None):
|
||
task.message = msg
|
||
if pct is not None:
|
||
task.progress = int(pct)
|
||
self._broadcast(task)
|
||
|
||
def _process(self, task):
|
||
# 收集视频文件
|
||
videos = []
|
||
skipped_in_run = []
|
||
if os.path.isfile(task.source):
|
||
if task.source.lower().endswith(VIDEO_EXTS):
|
||
videos = [task.source]
|
||
elif os.path.isdir(task.source):
|
||
# 过滤已处理（同路线）的视频
|
||
history = _load_history(task.source)
|
||
for f in sorted(os.listdir(task.source)):
|
||
if f.lower().endswith(VIDEO_EXTS):
|
||
key = os.path.join(task.source, f)
|
||
if f"{key}|{task.route}" in history:
|
||
skipped_in_run.append(key)
|
||
else:
|
||
videos.append(key)
|
||
if skipped_in_run:
|
||
task.message = f"跳过 {len(skipped_in_run)} 个已处理视频（同路线去重）"
|
||
self._update(task, task.message, 1)
|
||
if not videos:
|
||
raise ValueError(f"所选内容均已处理过，无需重复处理: {task.source}")
|
||
|
||
out_dir = os.path.join(OUTPUT_DIR, task.id)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
total = len(videos)
|
||
task.total_videos = total
|
||
task.started_at = time.time()
|
||
# 初始化子任务状态
|
||
for v in videos:
|
||
task.subtasks[v] = {"status": "pending", "progress": 0, "message": "等待中", "started_at": None, "elapsed": 0}
|
||
task.message = f"开始处理 {total} 个视频（转录 {MAX_CONCURRENT_TASKS} + LLM {self.llm_concurrency} 并发）"
|
||
self._update(task, task.message, 2)
|
||
|
||
# 流水线并发：每个视频一个 worker 线程
|
||
# 转录内部拿 GPU 锁（串行），修正/翻译拿 LLM 信号量（限并发）
|
||
done_count = 0
|
||
done_lock = threading.Lock()
|
||
|
||
def _update_overall():
|
||
"""更新总体进度：已完成/失败/总数"""
|
||
done_total = task.processed_count + task.failed_count
|
||
overall_pct = int(done_total / total * 100) if total else 100
|
||
# 聚合消息
|
||
running = [st for st in task.subtasks.values() if st["status"] in ("transcribing", "fixing", "translating", "srt")]
|
||
running_names = [v for v, st in task.subtasks.items() if st["status"] in ("transcribing", "fixing", "translating", "srt")]
|
||
msg = f"总体 {done_total}/{total} 完成"
|
||
if running_names:
|
||
msg += f" | 处理中: {', '.join(os.path.basename(n) for n in running_names[:3])}"
|
||
if task.failed_count:
|
||
msg += f" | 失败 {task.failed_count}"
|
||
task.message = msg
|
||
self._update(task, msg, overall_pct)
|
||
|
||
def _sub_update(video, status, progress, message):
|
||
"""更新单个视频子任务状态（含耗时）"""
|
||
cur = task.subtasks.get(video, {})
|
||
started = cur.get("started_at")
|
||
if started is None:
|
||
started = time.time()
|
||
elapsed = (time.time() - started) if started else 0
|
||
task.subtasks[video] = {
|
||
"status": status, "progress": progress, "message": message,
|
||
"started_at": started, "elapsed": round(elapsed, 1),
|
||
}
|
||
self._broadcast(task)
|
||
|
||
def _video_worker(v_idx, video):
|
||
nonlocal done_count
|
||
try:
|
||
_sub_update(video, "transcribing", 5, "转录中")
|
||
self._process_one(task, video, out_dir, v_idx, total, _sub_update)
|
||
with done_lock:
|
||
task.processed_count += 1
|
||
done_count += 1
|
||
_sub_update(video, "done", 100, "完成")
|
||
except Exception as ex:
|
||
with done_lock:
|
||
task.failed_count += 1
|
||
done_count += 1
|
||
task.video_errors.append({"video": video, "error": str(ex)[:300]})
|
||
# 记录完整 traceback 到日志（排查用）
|
||
try:
|
||
import traceback as _tb
|
||
tb = _tb.format_exc()
|
||
with open(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "_error.log"), 'a', encoding='utf-8') as _ef:
|
||
_ef.write(f"\n=== {os.path.basename(video)} ===\n{tb}\n")
|
||
except Exception:
|
||
pass
|
||
_sub_update(video, "error", 0, f"失败: {str(ex)[:80]}")
|
||
finally:
|
||
_update_overall()
|
||
|
||
# 并发执行（worker 数 = 转录锁 + LLM 并发，转录阶段大部分在等锁）
|
||
pool_size = max(MAX_CONCURRENT_TASKS, self.llm_concurrency)
|
||
with ThreadPoolExecutor(max_workers=pool_size) as pool:
|
||
futures = [pool.submit(_video_worker, i, v) for i, v in enumerate(videos)]
|
||
# 等待全部完成（含失败）
|
||
for f in futures:
|
||
f.result()
|
||
|
||
if task.failed_count > 0:
|
||
task.message = f"部分完成: 成功 {task.processed_count}, 失败 {task.failed_count}"
|
||
else:
|
||
task.message = f"完成 ({task.processed_count} 个视频)"
|
||
self._update(task, task.message, 100)
|
||
|
||
def _process_one(self, task, video, out_dir, v_idx=0, total=1, sub_update=None):
|
||
"""处理单个视频（转录→修正→翻译→SRT）——支持断点续跑
|
||
中间产物在 视频所在文件夹/.subtitle-work/（固定目录，不按视频分）
|
||
续跑：有 transcript 无 fixed → 从修正开始；有 fixed → 从翻译开始
|
||
转录拿 GPU 锁（串行），修正/翻译拿 LLM 信号量（限并发）
|
||
"""
|
||
def sub(status, progress, message):
|
||
if sub_update:
|
||
sub_update(video, status, progress, message)
|
||
|
||
base = os.path.splitext(os.path.basename(video))[0]
|
||
fname = os.path.basename(video)
|
||
# 中间产物目录：视频所在文件夹 .subtitle-work/
|
||
work_dir = os.path.join(os.path.dirname(video), WORK_DIRNAME)
|
||
os.makedirs(work_dir, exist_ok=True)
|
||
transcript_json = os.path.join(work_dir, f"{base}_transcript.json")
|
||
fixed_json = os.path.join(work_dir, f"{base}_fixed.json")
|
||
translated_json = os.path.join(work_dir, f"{base}_translated.json")
|
||
wav_path = os.path.join(work_dir, f"{base}_audio.wav")
|
||
|
||
# ===== 阶段 1: 转录（可续跑）=====
|
||
if os.path.exists(transcript_json):
|
||
sub("transcribing", 100, f"转录已存在（续跑）{fname}")
|
||
else:
|
||
# 提取音频 + 降噪（CPU，可并发——在 GPU 锁外做）
|
||
sub("transcribing", 2, f"提取音频+降噪 {fname}")
|
||
extract_and_denoise(video, wav_path, denoise=True,
|
||
progress_cb=lambda msg, pct: sub("transcribing", 2 + int(pct * 0.06), f"{msg} {fname}"))
|
||
# 转录（GPU 锁）
|
||
with self.transcribe_lock:
|
||
sub("transcribing", 10, f"转录中 {fname}")
|
||
def _tcb(msg, pct, _fname=fname):
|
||
if msg.startswith("转录中"):
|
||
sub("transcribing", 10 + int(pct * 0.3), f"{_fname}: {msg.split('转录中')[1].strip()}")
|
||
else:
|
||
sub("transcribing", 10 + int(pct * 0.3), f"{msg} {_fname}")
|
||
_result, _warnings = transcribe(video, transcript_json, language=task.language,
|
||
progress_cb=_tcb, wav_path=wav_path)
|
||
if _warnings:
|
||
task.video_errors.append({"video": video, "error": "；".join(_warnings), "warn": True})
|
||
task.transcript_path = transcript_json
|
||
|
||
if task.route == "transcribe_only":
|
||
task.files.append(transcript_json)
|
||
self._save_history_record(video, task, [transcript_json])
|
||
return
|
||
|
||
with open(transcript_json, encoding='utf-8') as f:
|
||
items = json.load(f)
|
||
|
||
# ===== 阶段 2+3: 修正 + 翻译（LLM 信号量，可续跑）=====
|
||
with self.llm_sem:
|
||
# 修正（可续跑）
|
||
if os.path.exists(fixed_json):
|
||
sub("fixing", 100, f"修正已存在（续跑）{fname}")
|
||
with open(fixed_json, encoding='utf-8') as f:
|
||
items = json.load(f)
|
||
else:
|
||
sub("fixing", 45, f"LLM 修正 {fname}")
|
||
items = fix(items, progress_cb=lambda msg, pct: sub("fixing", 45 + int(pct * 0.15), f"{msg} [{fname}]"))
|
||
with open(fixed_json, 'w', encoding='utf-8') as f:
|
||
json.dump(items, f, ensure_ascii=False, indent=1)
|
||
|
||
# 翻译（可续跑：从 translated_json 加载已有结果，只翻缺失段）
|
||
existing_zh = {}
|
||
existing_en = {}
|
||
if os.path.exists(translated_json):
|
||
try:
|
||
with open(translated_json, encoding='utf-8') as f:
|
||
prev = json.load(f)
|
||
# 提取已有翻译
|
||
for i, it in enumerate(prev):
|
||
if it.get('zh'):
|
||
existing_zh[i] = it['zh']
|
||
if it.get('en'):
|
||
existing_en[i] = it['en']
|
||
sub("translating", 30, f"翻译已有 {len(existing_zh)}/{len(items)} 段（续跑）{fname}")
|
||
except Exception:
|
||
existing_zh, existing_en = {}, {}
|
||
|
||
sub("translating", 40, f"翻译中 {fname}")
|
||
if task.route == "direct":
|
||
items = translate_direct(items,
|
||
progress_cb=lambda msg, pct: sub("translating", 40 + int(pct * 0.5), f"{msg} [{fname}]"),
|
||
existing_zh=existing_zh)
|
||
elif task.route == "via_en":
|
||
items = translate_via_en(items,
|
||
progress_cb=lambda msg, pct: sub("translating", 40 + int(pct * 0.5), f"{msg} [{fname}]"),
|
||
existing_en=existing_en, existing_zh=existing_zh)
|
||
with open(translated_json, 'w', encoding='utf-8') as f:
|
||
json.dump(items, f, ensure_ascii=False, indent=1)
|
||
|
||
# 4. SRT（写视频同目录，PotPlayer 需要；副本到 output）
|
||
sub("srt", 92, f"生成 SRT {fname}")
|
||
srt_files = write_srts(items, video, task.route, out_dir)
|
||
task.files.extend(srt_files)
|
||
sub("srt", 98, f"SRT 完成 {fname}")
|
||
|
||
# 5. 记录 history
|
||
self._save_history_record(video, task, srt_files)
|
||
|
||
def _save_history_record(self, video, task, srt_files):
|
||
"""记录 history（精确到视频，存在视频所在文件夹）"""
|
||
try:
|
||
history = _load_history(video)
|
||
hkey = f"{video}|{task.route}"
|
||
# 收集警告（transcribe 的 warnings 记录在 video_errors 里 warn=true）
|
||
warnings = []
|
||
for ve in task.video_errors:
|
||
if ve.get("video") == video and ve.get("warn"):
|
||
warnings.append(ve["error"])
|
||
history[hkey] = {
|
||
"video": video,
|
||
"route": task.route,
|
||
"files": list(srt_files),
|
||
"warnings": warnings,
|
||
"time": time.time(),
|
||
}
|
||
_save_history(video, history)
|
||
except Exception:
|
||
pass
|
||
|
||
# ---- API ----
|
||
def list_tasks(self):
|
||
return [t.to_dict() for t in sorted(self.tasks.values(), key=lambda t: -t.created)]
|
||
|
||
def get_task(self, tid):
|
||
t = self.tasks.get(tid)
|
||
return t.to_dict() if t else None
|
||
|
||
def retry_task(self, tid):
|
||
t = self.tasks.get(tid)
|
||
if not t or t.status not in ("error", "done"):
|
||
return None
|
||
# 复用 ID 重新入队
|
||
t.status = "pending"
|
||
t.error = ""
|
||
t.message = "重试"
|
||
t.progress = 0
|
||
with self.lock:
|
||
self.queue.append(t.id)
|
||
self._broadcast(t)
|
||
self._maybe_start()
|
||
return t.to_dict()
|
||
|
||
def cancel_task(self, tid):
|
||
"""从队列移除（无法中断运行中的，只移出等待队列）"""
|
||
t = self.tasks.get(tid)
|
||
if not t:
|
||
return None
|
||
if t.status == "pending":
|
||
with self.lock:
|
||
if tid in self.queue:
|
||
self.queue.remove(tid)
|
||
t.status = "cancelled"
|
||
t.message = "已取消"
|
||
self._broadcast(t)
|
||
return t.to_dict()
|
||
|
||
|
||
manager = TaskManager()
|