Upload app.py with huggingface_hub

app.py

@@ -1,4 +1,547 @@
-chromium
-ffmpeg
-fonts-noto-color-emoji
-libnvidia-encode1
+import gradio as gr
+import spaces
+import subprocess
+import asyncio
+import json
+import base64
+import os
+import shutil
+import socket
+import time
+import tempfile
+
+import aiohttp
+import websockets
+
+
+# ── Static FFmpeg with NVENC support ─────────────────────────────────────────
+_STATIC_FFMPEG = "/tmp/ffmpeg-static/ffmpeg"
+
+def _install_static_ffmpeg():
+    if os.path.exists(_STATIC_FFMPEG):
+        print(f"[wr] static ffmpeg already at {_STATIC_FFMPEG}", flush=True)
+        return
+    print("[wr] downloading static ffmpeg with nvenc (BtbN)...", flush=True)
+    os.makedirs("/tmp/ffmpeg-src", exist_ok=True)
+    url = "https://github.com/BtbN/FFmpeg-Builds/releases/download/latest/ffmpeg-master-latest-linux64-gpl.tar.xz"
+    subprocess.run(["wget", "-q", "-O", "/tmp/ffmpeg.tar.xz", url], check=True)
+    subprocess.run(["tar", "-xf", "/tmp/ffmpeg.tar.xz", "-C", "/tmp/ffmpeg-src"], check=True)
+    import glob as _glob
+    bins = _glob.glob("/tmp/ffmpeg-src/*/bin/ffmpeg")
+    if not bins:
+        raise RuntimeError("ffmpeg binary not found after extraction")
+    os.makedirs(os.path.dirname(_STATIC_FFMPEG), exist_ok=True)
+    shutil.copy2(bins[0], _STATIC_FFMPEG)
+    os.chmod(_STATIC_FFMPEG, 0o755)
+    print("[wr] static ffmpeg ready", flush=True)
+
+_install_static_ffmpeg()
+FFMPEG_BIN = _STATIC_FFMPEG if os.path.exists(_STATIC_FFMPEG) else "ffmpeg"
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+CHROMIUM_PATHS = [
+    "/usr/bin/chromium",
+    "/usr/bin/chromium-browser",
+    "/usr/bin/google-chrome-stable",
+    "/usr/bin/google-chrome",
+]
+
+# Normalize audio to a common format before concat to avoid channel/rate mismatches
+AUDIO_NORMALIZE = "aformat=sample_rates=44100:channel_layouts=stereo"
+
+
+def get_chromium_path():
+    for path in CHROMIUM_PATHS:
+        if os.path.exists(path):
+            return path
+    raise RuntimeError(f"Chromium not found. Tried: {CHROMIUM_PATHS}")
+
+
+def get_free_port():
+    with socket.socket() as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+def get_media_duration(path: str) -> float | None:
+    result = subprocess.run(
+        [
+            "ffprobe", "-v", "quiet",
+            "-show_entries", "format=duration",
+            "-of", "default=noprint_wrappers=1:nokey=1",
+            path,
+        ],
+        capture_output=True, text=True,
+    )
+    try:
+        return float(result.stdout.strip())
+    except ValueError:
+        return None
+
+
+def has_audio_stream(path: str) -> bool:
+    result = subprocess.run(
+        [
+            "ffprobe", "-v", "quiet",
+            "-show_streams", "-select_streams", "a:0",
+            "-of", "default=noprint_wrappers=1",
+            path,
+        ],
+        capture_output=True, text=True,
+    )
+    return bool(result.stdout.strip())
+
+
+def build_ffmpeg_cmd(
+    concat_path: str,
+    output_path: str,
+    w: int,
+    h: int,
+    rec_duration: float,
+    audio_path: str | None,
+    pre_rec_path: str | None,
+    post_rec_path: str | None,
+    postroll_path: str | None,
+) -> list[str]:
+    """
+    Build a single-pass FFmpeg command for any combination of optional segments.
+
+    Final order: pre_recording → main recording → post_recording → post_roll
+
+    FFmpeg inputs are assigned dynamically:
+      0        — main recording (concat demuxer of PNG frames)
+      1        — audio track for main (optional)
+      then any video files in the order: pre_rec, post_rec, postroll
+    """
+
+    scale_video = (
+        f"scale={w}:{h}:force_original_aspect_ratio=decrease,"
+        f"pad={w}:{h}:(ow-iw)/2:(oh-ih)/2,setsar=1"
+    )
+    scale_main = "scale=trunc(iw/2)*2:trunc(ih/2)*2,setsar=1,fps=25"
+
+    video_enc = ["-c:v", "h264_nvenc", "-preset", "p4", "-cq", "15", "-pix_fmt", "yuv420p", "-profile:v", "high", "-level:v", "4.2"]
+    audio_enc = ["-c:a", "aac", "-b:a", "192k"]
+    tail      = ["-movflags", "+faststart", "-y", output_path]
+
+    # ── Assign FFmpeg input indices ───────────────────────────────────────────
+    ffmpeg_inputs = ["-f", "concat", "-safe", "0", "-i", concat_path]
+    next_idx = 1
+
+    audio_idx = None
+    if audio_path:
+        ffmpeg_inputs += ["-i", audio_path]
+        audio_idx = next_idx
+        next_idx += 1
+
+    # Video file segments in playback order: pre_rec, post_rec, postroll
+    video_files = []   # (label, path, input_idx)
+    for label, path in [("pre", pre_rec_path), ("postrec", post_rec_path), ("postroll", postroll_path)]:
+        if path:
+            ffmpeg_inputs += ["-i", path]
+            video_files.append((label, path, next_idx))
+            next_idx += 1
+
+    pre_info     = next((v for v in video_files if v[0] == "pre"),     None)
+    postrec_info = next((v for v in video_files if v[0] == "postrec"), None)
+    postroll_info= next((v for v in video_files if v[0] == "postroll"),None)
+
+    # Probe video file durations and audio tracks
+    def probe(info):
+        if not info:
+            return None
+        _, path, _ = info
+        return {
+            "has_audio": has_audio_stream(path),
+            "duration":  get_media_duration(path) or 0.0,
+        }
+
+    pre_meta     = probe(pre_info)
+    postrec_meta = probe(postrec_info)
+    postroll_meta= probe(postroll_info)
+
+    # Segments in final playback order with all needed metadata
+    # Each entry: (video_filter, has_real_audio, audio_filter_or_null, duration)
+    segments = []
+
+    if pre_info:
+        _, _, idx = pre_info
+        segments.append({
+            "vf":        f"[{idx}:v]{scale_video}",
+            "has_audio": pre_meta["has_audio"],
+            "af_real":   f"[{idx}:a]{AUDIO_NORMALIZE},asetpts=PTS-STARTPTS",
+            "duration":  pre_meta["duration"],
+        })
+
+    # Main recording always present
+    segments.append({
+        "vf":        f"[0:v]{scale_main}",
+        "has_audio": bool(audio_path),
+        "af_real":   f"[{audio_idx}:a]{AUDIO_NORMALIZE},atrim=duration={rec_duration:.4f},asetpts=PTS-STARTPTS"
+                     if audio_path else None,
+        "duration":  rec_duration,
+    })
+
+    if postrec_info:
+        _, _, idx = postrec_info
+        segments.append({
+            "vf":        f"[{idx}:v]{scale_video}",
+            "has_audio": postrec_meta["has_audio"],
+            "af_real":   f"[{idx}:a]{AUDIO_NORMALIZE},asetpts=PTS-STARTPTS",
+            "duration":  postrec_meta["duration"],
+        })
+
+    if postroll_info:
+        _, _, idx = postroll_info
+        segments.append({
+            "vf":        f"[{idx}:v]{scale_video}",
+            "has_audio": postroll_meta["has_audio"],
+            "af_real":   f"[{idx}:a]{AUDIO_NORMALIZE},asetpts=PTS-STARTPTS",
+            "duration":  postroll_meta["duration"],
+        })
+
+    need_audio_out = any(s["has_audio"] for s in segments)
+
+    # ── Simple path: only main recording ─────────────────────────────────────
+    if len(segments) == 1:
+        cmd = [FFMPEG_BIN] + ffmpeg_inputs + ["-vf", scale_main]
+        if audio_path:
+            cmd += audio_enc + ["-shortest"]
+        return cmd + video_enc + tail
+
+    # ── Complex path: multiple segments, single filter_complex pass ───────────
+    fp = []
+    v_labels = []
+    a_labels = []
+
+    for i, seg in enumerate(segments):
+        vl = f"v{i}"
+        al = f"a{i}"
+        fp.append(f"{seg['vf']}[{vl}]")
+        v_labels.append(f"[{vl}]")
+
+        if need_audio_out:
+            if seg["has_audio"]:
+                fp.append(f"{seg['af_real']}[{al}]")
+            else:
+                d = f"{seg['duration']:.4f}"
+                fp.append(f"anullsrc=r=44100:cl=stereo,atrim=duration={d},asetpts=PTS-STARTPTS[{al}]")
+            a_labels.append(f"[{al}]")
+
+    n = len(segments)
+    fp.append(f"{''.join(v_labels)}concat=n={n}:v=1:a=0[outv]")
+    if need_audio_out:
+        fp.append(f"{''.join(a_labels)}concat=n={n}:v=0:a=1[outa]")
+
+    cmd = [FFMPEG_BIN] + ffmpeg_inputs + ["-filter_complex", ";".join(fp), "-map", "[outv]"]
+    if need_audio_out:
+        cmd += ["-map", "[outa]"] + audio_enc
+    return cmd + video_enc + tail
+
+
+@spaces.GPU(duration=120)
+def _gpu_encode(cmd: list[str]) -> tuple[int, str]:
+    """Run FFmpeg on GPU using h264_nvenc with real-time stderr logging."""
+    import glob as _glob
+
+    # Find libnvidia-encode.so.1 wherever ZeroGPU mounts it
+    found = (
+        _glob.glob("/usr/lib*/**/libnvidia-encode*", recursive=True)
+        + _glob.glob("/usr/local/**/libnvidia-encode*", recursive=True)
+        + _glob.glob("/run/**/libnvidia-encode*", recursive=True)
+        + _glob.glob("/opt/**/libnvidia-encode*", recursive=True)
+    )
+    print(f"[wr] libnvidia-encode glob: {found}", flush=True)
+    # Also check ldconfig cache
+    ldc = subprocess.run(["ldconfig", "-p"], capture_output=True, text=True)
+    for line in ldc.stdout.splitlines():
+        if "nvidia" in line.lower():
+            print(f"[wr] ldconfig: {line.strip()}", flush=True)
+
+    extra_dirs = list({os.path.dirname(p) for p in found})
+    env = dict(os.environ)
+    env["LD_LIBRARY_PATH"] = ":".join(extra_dirs) + ":" + env.get("LD_LIBRARY_PATH", "")
+    print(f"[wr] LD_LIBRARY_PATH={env['LD_LIBRARY_PATH']}", flush=True)
+
+    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, bufsize=0, env=env)
+    chunks = []
+    while True:
+        chunk = proc.stderr.read(512)
+        if not chunk:
+            break
+        chunks.append(chunk)
+        print(chunk.decode("utf-8", errors="replace"), end="", flush=True)
+    proc.wait()
+    return proc.returncode, b"".join(chunks).decode("utf-8", errors="replace")
+
+
+async def _record_async(
+    url: str, w: int, h: int, duration: int,
+    output_path: str,
+    audio_path: str | None,
+    pre_rec_path: str | None,
+    post_rec_path: str | None,
+    postroll_path: str | None,
+):
+    t0 = time.monotonic()
+    def log(msg):
+        print(f"[wr {time.monotonic()-t0:.1f}s] {msg}", flush=True)
+
+    port = get_free_port()
+    chromium = get_chromium_path()
+    log(f"launching chromium on port {port}, url={url}, size={w}x{h}, duration={duration}")
+
+    proc = subprocess.Popen(
+        [
+            chromium,
+            f"--remote-debugging-port={port}",
+            "--headless",
+            "--no-sandbox",
+            "--disable-dev-shm-usage",
+            "--disable-setuid-sandbox",
+            "--use-gl=swiftshader",
+            "--enable-webgl",
+            "--hide-scrollbars",
+            "--autoplay-policy=no-user-gesture-required",
+            f"--window-size={w},{h}",
+        ],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+    try:
+        ws_url = None
+        async with aiohttp.ClientSession() as session:
+            for attempt in range(20):
+                await asyncio.sleep(0.5)
+                try:
+                    async with session.get(f"http://localhost:{port}/json/list") as resp:
+                        pages = await resp.json()
+                        if pages:
+                            ws_url = pages[0]["webSocketDebuggerUrl"]
+                            log(f"chrome devtools ready after {attempt+1} attempts")
+                            break
+                except Exception as e:
+                    log(f"devtools attempt {attempt+1} failed: {e}")
+                    continue
+
+        if not ws_url:
+            raise gr.Error("Could not connect to Chrome DevTools.")
+
+        frames = []
+        msg_id = 1
+
+        async with websockets.connect(ws_url, max_size=50_000_000) as ws:
+
+            async def send(method, params=None):
+                nonlocal msg_id
+                current_id = msg_id
+                await ws.send(json.dumps({"id": current_id, "method": method, "params": params or {}}))
+                msg_id += 1
+                return current_id
+
+            await send("Page.enable")
+            await send("Emulation.setDeviceMetricsOverride", {
+                "width": w, "height": h, "deviceScaleFactor": 1, "mobile": False,
+            })
+
+            log(f"navigating to {url}")
+            nav_id = await send("Page.navigate", {"url": url})
+            deadline = time.monotonic() + 10
+            while time.monotonic() < deadline:
+                try:
+                    raw = await asyncio.wait_for(ws.recv(), timeout=1.0)
+                    if json.loads(raw).get("id") == nav_id:
+                        log("Page.navigate acknowledged")
+                        break
+                except asyncio.TimeoutError:
+                    continue
+
+            log("waiting for Page.loadEventFired (max 15s)")
+            deadline = time.monotonic() + 15
+            while time.monotonic() < deadline:
+                try:
+                    raw = await asyncio.wait_for(ws.recv(), timeout=1.0)
+                    if json.loads(raw).get("method") == "Page.loadEventFired":
+                        log("Page.loadEventFired received")
+                        break
+                except asyncio.TimeoutError:
+                    continue
+            else:
+                log("load timeout — starting screencast anyway")
+
+            await send("Runtime.enable")
+            await send("Runtime.evaluate", {
+                "expression": (
+                    "document.querySelectorAll('audio,video')"
+                    ".forEach(m => { m.muted = false; m.play().catch(() => {}); })"
+                ),
+                "awaitPromise": False,
+            })
+
+            log(f"starting screencast for {duration}s")
+            await send("Page.startScreencast", {
+                "format": "png",
+                "maxWidth": w,
+                "maxHeight": h,
+                "everyNthFrame": 1,
+            })
+
+            start = time.monotonic()
+            last_log = start
+            while time.monotonic() - start < duration:
+                try:
+                    raw = await asyncio.wait_for(ws.recv(), timeout=1.0)
+                    obj = json.loads(raw)
+                    if obj.get("method") == "Page.screencastFrame":
+                        ts = time.monotonic() - start
+                        frames.append((ts, obj["params"]["data"]))
+                        await send("Page.screencastFrameAck", {
+                            "sessionId": obj["params"]["sessionId"]
+                        })
+                        now = time.monotonic()
+                        if now - last_log >= 5:
+                            log(f"recording... {ts:.1f}s elapsed, {len(frames)} frames")
+                            last_log = now
+                except asyncio.TimeoutError:
+                    continue
+
+            log(f"screencast done: {len(frames)} frames captured")
+            await send("Page.stopScreencast")
+
+    finally:
+        proc.terminate()
+        proc.wait()
+        log("chromium terminated")
+
+    if not frames:
+        raise gr.Error("No frames captured. Check the URL.")
+
+    log(f"writing {len(frames)} frames to disk")
+    with tempfile.TemporaryDirectory() as tmpdir:
+        for i, (ts, data) in enumerate(frames):
+            with open(os.path.join(tmpdir, f"frame_{i:06d}.png"), "wb") as f:
+                f.write(base64.b64decode(data))
+
+        concat_path = os.path.join(tmpdir, "concat.txt")
+        with open(concat_path, "w") as f:
+            f.write("ffconcat version 1.0\n")
+            for i, (ts, _) in enumerate(frames):
+                dur = (frames[i + 1][0] - ts) if i < len(frames) - 1 else (duration - ts)
+                f.write(f"file 'frame_{i:06d}.png'\n")
+                f.write(f"duration {dur:.4f}\n")
+            f.write(f"file 'frame_{len(frames) - 1:06d}.png'\n")
+
+        cmd = build_ffmpeg_cmd(
+            concat_path, output_path, w, h, duration,
+            audio_path, pre_rec_path, post_rec_path, postroll_path,
+        )
+        log(f"running ffmpeg on GPU: {' '.join(cmd[:8])}...")
+        returncode, stderr = _gpu_encode(cmd)
+        if returncode != 0:
+            log(f"ffmpeg stderr: {stderr[-1000:]}")
+            raise gr.Error(f"FFmpeg failed: {stderr[-2000:]}")
+        log("ffmpeg done")
+
+    log("complete")
+    return output_path
+
+
+def _cp(path: str | None) -> str | None:
+    """Copy an uploaded file to /tmp before asyncio yields give Gradio a cleanup opportunity."""
+    if not path:
+        return None
+    dst = f"/tmp/wr_{os.getpid()}_{int(time.time() * 1000)}_{os.path.basename(path)}"
+    shutil.copy2(path, dst)
+    return dst
+
+
+async def record(
+    url: str,
+    width: float,
+    height: float,
+    duration: int,
+    audio_file: str | None,
+    pre_rec_file: str | None,
+    post_rec_file: str | None,
+    postroll_file: str | None,
+) -> str:
+    if not url.strip():
+        raise gr.Error("Please enter a URL.")
+    if not url.startswith(("http://", "https://")):
+        url = "https://" + url
+
+    # Copy uploaded files synchronously before the first await, so Gradio's
+    # event-loop cleanup (triggered at await yield points) can't delete them.
+    audio_file    = _cp(audio_file)
+    pre_rec_file  = _cp(pre_rec_file)
+    post_rec_file = _cp(post_rec_file)
+    postroll_file = _cp(postroll_file)
+    print(f"[wr] files copied: audio={audio_file} pre={pre_rec_file} post={post_rec_file} postroll={postroll_file}", flush=True)
+
+    w = max(2, (int(width) // 2) * 2)
+    h = max(2, (int(height) // 2) * 2)
+
+    # If audio uploaded, its duration overrides the slider (capped at 2 min)
+    if audio_file:
+        audio_dur = get_media_duration(audio_file)
+        if audio_dur:
+            duration = int(min(audio_dur, 120))
+
+    output_path = f"/tmp/recording_{os.getpid()}_{int(time.time())}.mp4"
+    return await _record_async(
+        url, w, h, duration, output_path,
+        audio_file, pre_rec_file, post_rec_file, postroll_file,
+    )
+
+
+with gr.Blocks(title="Website Renderer") as demo:
+    gr.Markdown(
+        "# Website Renderer\n"
+        "Recording starts after full page load. "
+        "Uploading audio overrides the duration slider (max 2 min)."
+    )
+
+    url_input = gr.Textbox(label="URL", placeholder="https://example.com")
+
+    with gr.Row():
+        width_input  = gr.Number(label="Width",  value=540, precision=0, minimum=100, maximum=3840)
+        height_input = gr.Number(label="Height", value=960, precision=0, minimum=100, maximum=3840)
+        duration_input = gr.Slider(minimum=5, maximum=120, value=30, step=5, label="Duration (seconds)")
+
+    with gr.Row():
+        audio_input = gr.Audio(
+            label="Audio track (optional — overrides duration)",
+            type="filepath",
+            sources=["upload"],
+        )
+
+    with gr.Row():
+        pre_rec_input  = gr.File(
+            label="Pre-recording video (optional)",
+            file_types=[".mp4", ".mov", ".webm", ".avi", ".mkv"],
+        )
+        post_rec_input = gr.File(
+            label="Post-recording video (optional)",
+            file_types=[".mp4", ".mov", ".webm", ".avi", ".mkv"],
+        )
+        postroll_input = gr.File(
+            label="Post-roll video (optional)",
+            file_types=[".mp4", ".mov", ".webm", ".avi", ".mkv"],
+        )
+
+    record_btn   = gr.Button("Record", variant="primary")
+    video_output = gr.Video(label="Recorded Video")
+
+    record_btn.click(
+        fn=record,
+        inputs=[
+            url_input, width_input, height_input, duration_input,
+            audio_input, pre_rec_input, post_rec_input, postroll_input,
+        ],
+        outputs=video_output,
+    )
+
+if __name__ == "__main__":
+    demo.launch()