| import os
|
| import json
|
| import shutil
|
| import tempfile
|
| import subprocess
|
| import time
|
| import re
|
| import uuid
|
| import threading
|
| from concurrent.futures import ThreadPoolExecutor
|
|
|
| import gradio as gr
|
|
|
|
|
|
|
|
|
| HF_TOKEN = os.environ.get("HF_TOKEN", "")
|
| DATASET_ID = os.environ.get("DATASET_ID", "")
|
| MS_TOKEN = os.environ.get("MS_TOKEN", "")
|
| MS_DATASET_ID = os.environ.get("MS_DATASET_ID", "")
|
| ACCESS_PASSWORD = os.environ.get("ACCESS_PASSWORD", "")
|
|
|
| RESULT_KEEP_HOURS = float(os.environ.get("RESULT_KEEP_HOURS", "24"))
|
| STATE_KEEP_DAYS = float(os.environ.get("STATE_KEEP_DAYS", "7"))
|
|
|
|
|
| MAX_WORKERS = int(os.environ.get("MAX_WORKERS", "2"))
|
| UPLOAD_RETRY = int(os.environ.get("UPLOAD_RETRY", "3"))
|
| FFMPEG_TIMEOUT = int(os.environ.get("FFMPEG_TIMEOUT", "1800"))
|
| MIN_FREE_GB = float(os.environ.get("MIN_FREE_GB", "2"))
|
|
|
|
|
|
|
| SRT_FORCE_STYLE = os.environ.get(
|
| "SRT_FORCE_STYLE",
|
| "FontName=微软雅黑,FontSize=64,Bold=1,"
|
| "PrimaryColour=&H0008d76c,OutlineColour=&H00000000,"
|
| "BorderStyle=1,Outline=2.2,Shadow=0,Alignment=2,"
|
| "MarginV=8,Spacing=3.4"
|
| )
|
|
|
|
|
| FONTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fonts")
|
|
|
| BASE_DIR = "/data" if os.path.isdir("/data") and os.access("/data", os.W_OK) \
|
| else tempfile.gettempdir()
|
| JOBS_DIR = os.path.join(BASE_DIR, "jobs")
|
| STATE_DIR = os.path.join(BASE_DIR, "job_state")
|
| os.makedirs(JOBS_DIR, exist_ok=True)
|
| os.makedirs(STATE_DIR, exist_ok=True)
|
|
|
| VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv",
|
| ".webm", ".m4v", ".3gp", ".ts", ".rmvb"}
|
| SUB_EXTS = {".ass", ".ssa", ".srt"}
|
|
|
|
|
|
|
|
|
| EXECUTOR = ThreadPoolExecutor(max_workers=MAX_WORKERS)
|
| TASKS = {}
|
| _tasks_lock = threading.Lock()
|
| _ms_login_lock = threading.Lock()
|
| _ms_logged_in = {"done": False}
|
|
|
|
|
| def _state_path(tid):
|
| return os.path.join(STATE_DIR, f"{tid}.json")
|
|
|
|
|
| def _save_state(tid):
|
| try:
|
| with _tasks_lock:
|
| data = dict(TASKS.get(tid, {}))
|
| with open(_state_path(tid), "w", encoding="utf-8") as f:
|
| json.dump(data, f, ensure_ascii=False)
|
| except Exception:
|
| pass
|
|
|
|
|
| def _load_states():
|
| for fn in os.listdir(STATE_DIR):
|
| if not fn.endswith(".json"):
|
| continue
|
| try:
|
| with open(os.path.join(STATE_DIR, fn), encoding="utf-8") as f:
|
| d = json.load(f)
|
| if d.get("status") in ("queued", "running"):
|
| d["status"] = "interrupted"
|
| d["logs"] = d.get("logs", []) + ["⚠ Space 重启，任务被中断"]
|
| TASKS[d["id"]] = d
|
| except Exception:
|
| pass
|
|
|
|
|
| def new_task():
|
| tid = uuid.uuid4().hex[:12]
|
| with _tasks_lock:
|
| TASKS[tid] = {
|
| "id": tid, "status": "queued", "logs": [],
|
| "total": 0, "done": 0, "success": 0,
|
| "outputs": [], "created": time.time(),
|
| }
|
| _save_state(tid)
|
| return tid
|
|
|
|
|
| def t_log(tid, msg, save=True):
|
| ts = time.strftime("%H:%M:%S")
|
| with _tasks_lock:
|
| if tid in TASKS:
|
| TASKS[tid]["logs"].append(f"[{ts}] {msg}")
|
| if len(TASKS[tid]["logs"]) > 1000:
|
| TASKS[tid]["logs"] = TASKS[tid]["logs"][-800:]
|
| if save:
|
| _save_state(tid)
|
|
|
|
|
| def t_set(tid, **kw):
|
| with _tasks_lock:
|
| if tid in TASKS:
|
| TASKS[tid].update(kw)
|
| _save_state(tid)
|
|
|
|
|
| def t_get(tid):
|
| with _tasks_lock:
|
| return dict(TASKS[tid]) if tid in TASKS else None
|
|
|
|
|
|
|
|
|
|
|
| def sanitize_name(name, max_len=80):
|
| if not name:
|
| return ""
|
| name = os.path.basename(name)
|
| name = re.sub(r'[\\/:*?"<>|\r\n\t]', "_", name)
|
| name = name.strip().strip(".").strip()
|
| name = re.sub(r"\s+", " ", name)
|
| return name[:max_len].rstrip() if len(name) > max_len else name
|
|
|
| def strip_english_from_srt(src_path, dst_path):
|
| """读取 SRT，每条字幕只保留含中文的行，去掉纯英文行"""
|
| def has_chinese(s):
|
| return any('\u4e00' <= ch <= '\u9fff' for ch in s)
|
|
|
| with open(src_path, "r", encoding="utf-8-sig", errors="ignore") as f:
|
| content = f.read()
|
|
|
|
|
| blocks = re.split(r"\n\s*\n", content.strip())
|
| out_blocks = []
|
| for block in blocks:
|
| lines = block.splitlines()
|
| if len(lines) < 2:
|
| continue
|
| idx = lines[0].strip()
|
| timeline = lines[1].strip()
|
| text_lines = lines[2:]
|
|
|
| kept = [ln for ln in text_lines if has_chinese(ln)]
|
| if not kept:
|
| continue
|
| out_blocks.append(f"{idx}\n{timeline}\n" + "\n".join(kept))
|
|
|
| with open(dst_path, "w", encoding="utf-8") as f:
|
| f.write("\n\n".join(out_blocks) + "\n")
|
|
|
| def cleanup_once():
|
| """清理过期的成品目录与状态文件"""
|
| now = time.time()
|
|
|
| try:
|
| for tid in os.listdir(JOBS_DIR):
|
| d = os.path.join(JOBS_DIR, tid)
|
| if not os.path.isdir(d):
|
| continue
|
| age_h = (now - os.path.getmtime(d)) / 3600
|
| if age_h > RESULT_KEEP_HOURS:
|
| shutil.rmtree(d, ignore_errors=True)
|
|
|
| with _tasks_lock:
|
| if tid in TASKS:
|
| TASKS[tid]["outputs"] = []
|
| except Exception:
|
| pass
|
|
|
| try:
|
| for fn in os.listdir(STATE_DIR):
|
| if not fn.endswith(".json"):
|
| continue
|
| p = os.path.join(STATE_DIR, fn)
|
| age_d = (now - os.path.getmtime(p)) / 86400
|
| if age_d > STATE_KEEP_DAYS:
|
| os.remove(p)
|
| with _tasks_lock:
|
| TASKS.pop(fn[:-5], None)
|
| except Exception:
|
| pass
|
|
|
|
|
| def cleanup_loop():
|
| while True:
|
| time.sleep(3600)
|
| cleanup_once()
|
|
|
|
|
| def start_cleanup():
|
| cleanup_once()
|
| t = threading.Thread(target=cleanup_loop, daemon=True)
|
| t.start()
|
|
|
|
|
|
|
| def free_space_gb(path=None):
|
| try:
|
| return shutil.disk_usage(path or BASE_DIR)[2] / (1024 ** 3)
|
| except Exception:
|
| return 999
|
|
|
|
|
| def _run_ffmpeg_live(cmd, tid, tag, output_path, action):
|
| t_log(tid, f"[{tag}] {action}…")
|
| try:
|
| proc = subprocess.Popen(
|
| cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
| universal_newlines=True, bufsize=1,
|
| )
|
| except Exception as e:
|
| t_log(tid, f"[{tag}] ❌ 启动 ffmpeg 失败: {e}")
|
| return False
|
|
|
| start = time.time()
|
| last_log = 0
|
| try:
|
| for line in proc.stdout:
|
| line = line.strip()
|
| if line and any(k in line for k in ("frame=", "time=", "size=")):
|
| now = time.time()
|
| if now - last_log >= 2:
|
| t_log(tid, f"[{tag}] {line}", save=False)
|
| last_log = now
|
| elif line and any(k in line.lower() for k in ("error", "invalid", "failed")):
|
| t_log(tid, f"[{tag}] {line}", save=False)
|
| if time.time() - start > FFMPEG_TIMEOUT:
|
| proc.kill()
|
| t_log(tid, f"[{tag}] ❌ {action}超时，已终止")
|
| return False
|
| proc.wait()
|
| except Exception as e:
|
| t_log(tid, f"[{tag}] ❌ {action}异常: {e}")
|
| return False
|
| finally:
|
| _save_state(tid)
|
|
|
| if proc.returncode != 0:
|
| t_log(tid, f"[{tag}] ❌ {action}失败（码 {proc.returncode}）")
|
| return False
|
| if not (os.path.exists(output_path) and os.path.getsize(output_path) > 0):
|
| t_log(tid, f"[{tag}] ❌ {action}产物为空")
|
| return False
|
| t_log(tid, f"[{tag}] ✅ {action}完成 ({os.path.getsize(output_path)/1048576:.1f} MB)")
|
| return True
|
|
|
|
|
| def burn_subtitle(video_path, sub_path, output_path, fast_mode, tid, tag):
|
| ext = os.path.splitext(sub_path)[1].lower()
|
| if ext == ".srt":
|
| vf = (f"subtitles='{sub_path}':fontsdir='{FONTS_DIR}'"
|
| f":force_style='{SRT_FORCE_STYLE}'")
|
| else:
|
| vf = f"subtitles='{sub_path}':fontsdir='{FONTS_DIR}'"
|
|
|
| preset, crf = ("veryfast", "23") if fast_mode else ("medium", "20")
|
| cmd = ["ffmpeg", "-y", "-i", video_path, "-vf", vf,
|
| "-c:v", "libx264", "-preset", preset, "-crf", crf,
|
| "-c:a", "copy", output_path]
|
| return _run_ffmpeg_live(
|
| cmd, tid, tag, output_path,
|
| f"烧录字幕（{'快速' if fast_mode else '标准'}模式）")
|
|
|
|
|
| def convert_to_amv(input_path, output_path, tid, tag):
|
| cmd = ["ffmpeg", "-y", "-i", input_path,
|
| "-s", "320x240", "-r", "15",
|
| "-c:v", "amv", "-c:a", "adpcm_ima_amv",
|
| "-ar", "22050", "-ac", "1", "-block_size", "1470",
|
| output_path]
|
| return _run_ffmpeg_live(cmd, tid, tag, output_path, "转换 AMV")
|
|
|
|
|
| def _retry(fn, tid, tag, what):
|
| for attempt in range(1, UPLOAD_RETRY + 1):
|
| try:
|
| fn()
|
| return True
|
| except Exception as e:
|
| if attempt < UPLOAD_RETRY:
|
| wait = 2 ** attempt
|
| t_log(tid, f"[{tag}] ⚠ {what}第{attempt}次失败: {e}，{wait}s后重试")
|
| time.sleep(wait)
|
| else:
|
| t_log(tid, f"[{tag}] ❌ {what}最终失败: {e}")
|
| return False
|
|
|
|
|
| def upload_to_hf(fp, repo_path, tid, tag):
|
| if not (HF_TOKEN and DATASET_ID):
|
| t_log(tid, f"[{tag}] ⏭ 未配置 HF，跳过")
|
| return
|
| def _do():
|
| from huggingface_hub import HfApi
|
| HfApi(token=HF_TOKEN).upload_file(
|
| path_or_fileobj=fp, path_in_repo=repo_path,
|
| repo_id=DATASET_ID, repo_type="dataset")
|
| if _retry(_do, tid, tag, f"HF上传({os.path.basename(repo_path)})"):
|
| t_log(tid, f"[{tag}] ✅ 已上传 HF: {repo_path}")
|
|
|
|
|
| def upload_to_modelscope(fp, repo_path, tid, tag):
|
| if not (MS_TOKEN and MS_DATASET_ID):
|
| t_log(tid, f"[{tag}] ⏭ 未配置魔搭，跳过")
|
| return
|
| def _do():
|
| from modelscope.hub.api import HubApi
|
| api = HubApi()
|
| with _ms_login_lock:
|
| if not _ms_logged_in["done"]:
|
| api.login(MS_TOKEN)
|
| _ms_logged_in["done"] = True
|
| api.upload_file(path_or_fileobj=fp, path_in_repo=repo_path,
|
| repo_id=MS_DATASET_ID, repo_type="dataset",
|
| commit_message=f"上传: {repo_path}")
|
| if _retry(_do, tid, tag, f"魔搭上传({os.path.basename(repo_path)})"):
|
| t_log(tid, f"[{tag}] ✅ 已上传魔搭: {repo_path}")
|
|
|
|
|
| def pair_files(files):
|
| videos, subs = {}, {}
|
| for f in files:
|
| ext = os.path.splitext(f)[1].lower()
|
| stem = os.path.splitext(os.path.basename(f))[0]
|
| if ext in VIDEO_EXTS:
|
| videos[stem] = f
|
| elif ext in SUB_EXTS:
|
| subs[stem] = f
|
| pairs, unmatched = [], []
|
| for stem, v in videos.items():
|
| if stem in subs:
|
| pairs.append((stem, v, subs[stem]))
|
| else:
|
| unmatched.append(f"视频「{stem}」缺少同名字幕")
|
| for stem in subs:
|
| if stem not in videos:
|
| unmatched.append(f"字幕「{stem}」缺少同名视频")
|
| return pairs, unmatched
|
|
|
|
|
| def handle_one(stem, video, subtitle, fast_mode, folder, date_dir, tid, out_dir):
|
| tag = stem if len(stem) <= 20 else stem[:18] + "…"
|
| if free_space_gb() < MIN_FREE_GB:
|
| t_log(tid, f"[{tag}] ❌ 磁盘不足，跳过")
|
| return None
|
| work = tempfile.mkdtemp()
|
| results = []
|
| try:
|
| lv = os.path.join(work, "input" + os.path.splitext(video)[1])
|
| shutil.copy(video, lv)
|
| sub_ext = os.path.splitext(subtitle)[1].lower()
|
| ls = os.path.join(work, "sub" + sub_ext)
|
| if sub_ext == ".srt":
|
|
|
| strip_english_from_srt(subtitle, ls)
|
| t_log(tid, f"[{tag}] 已清洗 SRT，仅保留中文行")
|
| else:
|
| shutil.copy(subtitle, ls)
|
|
|
| clean = sanitize_name(stem)
|
| mp4_name = f"{clean}_subbed.mp4"
|
| amv_name = f"{clean}_subbed.amv"
|
| mp4_tmp = os.path.join(work, mp4_name)
|
| amv_tmp = os.path.join(work, amv_name)
|
|
|
|
|
| if not burn_subtitle(lv, ls, mp4_tmp, fast_mode, tid, tag):
|
| return results
|
|
|
| mp4_repo = f"{folder}/mp4/{date_dir}/{mp4_name}"
|
| upload_to_hf(mp4_tmp, mp4_repo, tid, tag)
|
| upload_to_modelscope(mp4_tmp, mp4_repo, tid, tag)
|
| mp4_keep = os.path.join(out_dir, mp4_name)
|
| shutil.copy(mp4_tmp, mp4_keep)
|
| results.append(mp4_keep)
|
|
|
|
|
| if convert_to_amv(mp4_tmp, amv_tmp, tid, tag):
|
| amv_repo = f"{folder}/amv/{date_dir}/{amv_name}"
|
| upload_to_hf(amv_tmp, amv_repo, tid, tag)
|
| upload_to_modelscope(amv_tmp, amv_repo, tid, tag)
|
| amv_keep = os.path.join(out_dir, amv_name)
|
| shutil.copy(amv_tmp, amv_keep)
|
| results.append(amv_keep)
|
|
|
| return results
|
| except Exception as e:
|
| t_log(tid, f"[{tag}] ❌ 处理异常: {e}")
|
| return results
|
| finally:
|
| shutil.rmtree(work, ignore_errors=True)
|
|
|
|
|
| def run_job(tid, saved_files, fast_mode, folder):
|
| try:
|
| t_set(tid, status="running")
|
| folder = sanitize_name(folder) or "video-srt"
|
| pairs, unmatched = pair_files(saved_files)
|
| for u in unmatched:
|
| t_log(tid, f"⚠ {u}")
|
| if not pairs:
|
| t_log(tid, "❌ 无同名配对，任务结束")
|
| t_set(tid, status="error")
|
| return
|
|
|
| out_dir = os.path.join(JOBS_DIR, tid)
|
| os.makedirs(out_dir, exist_ok=True)
|
| date_dir = time.strftime("%Y%m%d")
|
| t_set(tid, total=len(pairs))
|
| t_log(tid, f"共 {len(pairs)} 组 | 目标 {folder}/(mp4,amv) | 磁盘 {free_space_gb():.1f}GB")
|
|
|
| inner = ThreadPoolExecutor(max_workers=MAX_WORKERS)
|
| futs = [inner.submit(handle_one, stem, v, s, fast_mode,
|
| folder, date_dir, tid, out_dir)
|
| for stem, v, s in pairs]
|
| for fut in futs:
|
| files = fut.result() or []
|
| with _tasks_lock:
|
| TASKS[tid]["done"] += 1
|
| if files:
|
| TASKS[tid]["success"] += 1
|
| TASKS[tid]["outputs"].extend(files)
|
| _save_state(tid)
|
| inner.shutdown(wait=True)
|
|
|
| with _tasks_lock:
|
| s = TASKS[tid]["success"]
|
| t_log(tid, f"✅ 全部完成，成功 {s}/{len(pairs)}")
|
| t_set(tid, status="done")
|
| except Exception as e:
|
| t_log(tid, f"❌ 任务异常: {e}")
|
| t_set(tid, status="error")
|
| finally:
|
|
|
| dirs = {os.path.dirname(f) for f in saved_files}
|
| for d in dirs:
|
| shutil.rmtree(d, ignore_errors=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
| def submit(files, fast_mode, folder):
|
| if not files:
|
| raise gr.Error("请上传视频和字幕文件")
|
| persist = tempfile.mkdtemp(prefix="upload_", dir=BASE_DIR)
|
| saved = []
|
| for f in files:
|
| dst = os.path.join(persist, os.path.basename(f))
|
| shutil.copy(f, dst)
|
| saved.append(dst)
|
| tid = new_task()
|
| EXECUTOR.submit(run_job, tid, saved, fast_mode, folder)
|
| return (tid,
|
| f"✅ 任务已提交！任务ID：{tid}\n"
|
| f"可关闭页面，后台继续执行。到「查询进度」页粘贴此 ID 并勾选自动刷新即可看实时日志。")
|
|
|
|
|
| def query(tid):
|
| tid = (tid or "").strip()
|
| if not tid:
|
| return "请输入任务ID", None
|
| t = t_get(tid)
|
| if not t:
|
| return f"未找到任务 {tid}", None
|
| status_map = {"queued": "排队中", "running": "处理中", "done": "已完成",
|
| "error": "出错", "interrupted": "被中断(Space重启)"}
|
| head = (f"任务 {tid} | 状态：{status_map.get(t['status'], t['status'])} | "
|
| f"进度 {t['done']}/{t['total']}（成功 {t['success']}）\n"
|
| + "─" * 40 + "\n")
|
| body = "\n".join(t["logs"][-300:])
|
| outs = None
|
| if t["status"] in ("done", "running"):
|
| outs = [p for p in t["outputs"] if os.path.exists(p)] or None
|
| return head + body, outs
|
|
|
|
|
| _load_states()
|
| start_cleanup()
|
|
|
|
|
|
|
|
|
| CSS = """
|
| .gradio-container {max-width: 920px !important; margin: auto;}
|
| #title-box {text-align:center; padding:18px;
|
| background:linear-gradient(135deg,#6366f1,#8b5cf6);
|
| border-radius:14px; color:#fff; margin-bottom:14px;}
|
| #title-box h1 {margin:0; font-size:26px;}
|
| #title-box p {margin:6px 0 0; opacity:.92; font-size:14px;}
|
| """
|
|
|
| with gr.Blocks(title="ASS/SRT 硬字幕批量烧录", theme=gr.themes.Soft(
|
| primary_hue="indigo", secondary_hue="violet"), css=CSS) as demo:
|
|
|
| gr.HTML('<div id="title-box"><h1>🎬 ASS / SRT 硬字幕批量烧录工具</h1>'
|
| '<p>实时日志 · 后台任务 · MP4+AMV 双格式 · 多线程 · 失败重试 · 上传 HF / 魔搭</p></div>')
|
|
|
| with gr.Tab("📤 提交任务"):
|
| files_in = gr.File(label="📁 视频 + 字幕(.srt/.ass)（可多选，同名自动配对）",
|
| file_count="multiple")
|
| with gr.Row():
|
| folder_in = gr.Textbox(label="📂 数据集文件夹名", value="video-srt", scale=3)
|
| fast_in = gr.Checkbox(label="⚡ 快速模式", value=False, scale=1)
|
| gr.Markdown(
|
| "> SRT 会统一套用预设样式（微软雅黑/亮黄绿/大字）；ASS 用文件自带样式。\n"
|
| "> 产物上传到：`文件夹名/mp4/日期/` 和 `文件夹名/amv/日期/`")
|
| submit_btn = gr.Button("🚀 提交后台任务", variant="primary")
|
| tid_box = gr.Textbox(label="任务ID（请复制保存）", interactive=False)
|
| submit_msg = gr.Textbox(label="提交结果", lines=3, interactive=False)
|
| submit_btn.click(submit, [files_in, fast_in, folder_in],
|
| [tid_box, submit_msg])
|
|
|
| with gr.Tab("🔍 查询进度 / 实时日志"):
|
| gr.Markdown("输入任务ID，**勾选自动刷新**即可每 3 秒实时更新日志与进度。")
|
| with gr.Row():
|
| q_tid = gr.Textbox(label="任务ID", scale=3)
|
| auto = gr.Checkbox(label="🔄 自动刷新(3s)", value=True, scale=1)
|
| q_btn = gr.Button("查询", variant="primary")
|
| q_log = gr.Textbox(label="实时状态与日志", lines=20, show_copy_button=True)
|
| q_files = gr.File(label="📥 结果文件（MP4 / AMV）", file_count="multiple")
|
| q_btn.click(query, [q_tid], [q_log, q_files])
|
| timer = gr.Timer(3, active=True)
|
| auto.change(lambda a: gr.Timer(active=a), auto, timer)
|
| timer.tick(query, [q_tid], [q_log, q_files])
|
|
|
|
|
|
|
| _auth = ("user", ACCESS_PASSWORD) if ACCESS_PASSWORD else None
|
|
|
| demo.launch(
|
| server_name="0.0.0.0",
|
| server_port=7860,
|
| auth=_auth,
|
| auth_message="请输入访问密码（用户名：user）",
|
| )
|
|
|
|
|
