File size: 4,206 Bytes
711802a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce79f16
 
 
 
 
 
326cf8e
711802a
 
ce79f16
711802a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
326cf8e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
711802a
ce79f16
711802a
 
 
 
 
ce79f16
711802a
 
ce79f16
711802a
 
 
 
 
 
 
 
 
 
326cf8e
711802a
 
 
 
 
 
326cf8e
711802a
 
 
c68ecfe
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
"""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")