ffmpeg-worker / app.py
ymo2017's picture
favicon_path: AI-generated app icon
c68ecfe verified
Raw History Blame Contribute Delete
4.21 kB
"""ffmpeg-worker: private ffmpeg job queue for Muse.
Public Space, but every job must present the WORKER_TOKEN (set as a Space
secret). Jobs serialize with concurrency_limit=1 so encodes never fight.
"""
import gradio as gr
import subprocess, os, json, tempfile, shutil, hmac, time
WORKER_TOKEN = os.environ.get("WORKER_TOKEN", "")
JOB_ROOT = "/tmp/ffmpeg-worker-jobs"
os.makedirs(JOB_ROOT, exist_ok=True)
def _cleanup():
now = time.time()
for d in os.listdir(JOB_ROOT):
p = os.path.join(JOB_ROOT, d)
try:
if os.path.isdir(p) and now - os.path.getmtime(p) > 7200:
shutil.rmtree(p, ignore_errors=True)
except OSError:
pass
def _wlog(msg):
# Container stdout -> Space runtime log. Gradio return values never reach
# it, so without this successful jobs are invisible there.
print(f"[ffmpeg-worker] {msg}", flush=True)
def run_job(files, job_json, token):
_cleanup()
if not WORKER_TOKEN or not hmac.compare_digest(str(token or ""), WORKER_TOKEN):
_wlog("auth failed")
return None, "auth failed"
try:
job = json.loads(job_json)
except Exception as e:
return None, f"bad job_json: {e}"
args = job.get("args")
if not isinstance(args, list) or not args:
return None, "job.args must be a non-empty list of ffmpeg args (without leading 'ffmpeg')"
if args and args[0] == "ffmpeg":
args = args[1:]
if args and args[0] == "-y":
args = args[1:]
# block obviously dangerous patterns; this box only encodes media
for a in args:
if isinstance(a, str) and (a.startswith("http://") or a.startswith("https://")
or a.startswith("ftp://") or a.startswith("|")):
return None, "network inputs / pipes are not allowed"
timeout = min(int(job.get("timeout", 1500)), 1800)
out_name = job.get("output") or ""
if not out_name or "/" in out_name or out_name.startswith("."):
return None, "job.output must be a plain filename"
# files: one path or a list of paths from gradio
if isinstance(files, str):
files = [files]
in_paths = [f for f in (files or []) if f and os.path.exists(f)]
if not in_paths:
return None, "no input files uploaded"
def _subst(a):
if not isinstance(a, str):
return a
for i, p in enumerate(in_paths):
a = a.replace("{input%d}" % i, p)
a = a.replace("{input}", in_paths[0]) # legacy alias = first input
return a
args = [_subst(a) for a in args]
workdir = tempfile.mkdtemp(dir=JOB_ROOT)
_wlog(f"start: {len(in_paths)} input(s), timeout={timeout}s, output={out_name!r}")
t1 = time.time()
try:
p = subprocess.run(["ffmpeg", "-y"] + args, capture_output=True, text=True,
timeout=timeout, cwd=workdir)
except subprocess.TimeoutExpired:
_wlog(f"timeout after {timeout}s (output={out_name!r})")
return None, f"timeout after {timeout}s"
dt = time.time() - t1
_wlog(f"done: rc={p.returncode} elapsed={dt:.1f}s output={out_name!r}")
log = (f"rc={p.returncode} elapsed={dt:.1f}s\n"
f"--- stdout ---\n{p.stdout}\n--- stderr (tail) ---\n{p.stderr[-6000:]}")
out_path = os.path.join(workdir, out_name)
if p.returncode != 0 or not os.path.exists(out_path):
return None, log
return out_path, log
with gr.Blocks(title="ffmpeg-worker") as demo:
gr.Markdown("Private ffmpeg job queue for Muse. Token-gated; jobs run one at a time.")
files = gr.File(label="input files (one or more)", file_count="multiple")
job_json = gr.Textbox(label="job json", lines=6,
value='{"args": ["-i", "{input}", "-c", "copy", "out.mp4"], "output": "out.mp4"}')
token = gr.Textbox(label="token", type="password")
btn = gr.Button("Run", variant="primary")
out_file = gr.File(label="output")
log = gr.Textbox(label="log", lines=12)
btn.click(run_job, inputs=[files, job_json, token], outputs=[out_file, log],
api_name="run", concurrency_limit=1)
demo.queue(max_size=32)
demo.launch(favicon_path="favicon.png")