Spaces:
Paused
Paused
| """Frame export (last-frame extraction from a generated video). | |
| `export_last_frame` is bound to the "Export Last Frame" button in the UI. | |
| It validates that the input path lies inside the managed task temp root or the | |
| legacy `/tmp/h3-outputs` directory, then runs ffmpeg to extract the final frame. | |
| This is the only place that exposes the disk to client-supplied paths; the | |
| path validation is intentionally strict. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import subprocess | |
| import tempfile | |
| import uuid | |
| import gradio as gr | |
| from task import _TASK_TEMP_ROOT | |
| def export_last_frame(video_path): | |
| """Extract the last frame from a generated video and return it as an image file. | |
| Path is validated to lie under the managed task/temp root; arbitrary client paths are rejected. | |
| Destination is unique per call. | |
| """ | |
| if not video_path: | |
| raise gr.Error("No video to extract from.") | |
| # Resolve and constrain to managed roots (task temp or legacy h3-outputs). | |
| try: | |
| resolved = os.path.realpath(video_path) | |
| except OSError: | |
| raise gr.Error("Invalid video path.") | |
| # Only managed task/output roots — never the bare system tempdir (symlink surface). | |
| allowed_roots = [ | |
| os.path.realpath(_TASK_TEMP_ROOT), | |
| os.path.realpath(os.path.join(tempfile.gettempdir(), "h3-outputs")), | |
| ] | |
| if not any(resolved == root or resolved.startswith(root + os.sep) for root in allowed_roots): | |
| raise gr.Error("Video path is outside the managed output area.") | |
| # Unique destination under the same parent or task intermediates. | |
| parent = os.path.dirname(resolved) | |
| out = os.path.join(parent, f"lastframe-{uuid.uuid4().hex[:12]}.png") | |
| result = subprocess.run( | |
| ["ffmpeg", "-sseof", "-0.1", "-i", resolved, "-frames:v", "1", "-y", out], | |
| capture_output=True, text=True, | |
| ) | |
| if result.returncode != 0: | |
| raise gr.Error(f"ffmpeg failed: {result.stderr[-300:]}") | |
| return out | |
| # ── Durable video finalization (display + download) ───────────────────────── | |
| # Public-ish mirror of task outputs. Gradio serves these more reliably than deep | |
| # task dirs, and clients can still download after the request finishes. | |
| _H3_OUTPUT_ROOT = os.environ.get( | |
| "H3_OUTPUT_ROOT", os.path.join(tempfile.gettempdir(), "h3-outputs") | |
| ) | |
| def _ensure_output_roots(): | |
| for d in (_H3_OUTPUT_ROOT, _TASK_TEMP_ROOT): | |
| try: | |
| os.makedirs(d, exist_ok=True) | |
| except OSError: | |
| pass | |
| def _finalize_output_video(frames, audio, sampling_rate, *, task=None, fps=24): | |
| """Mux frames (+ optional audio) to a durable mp4 and return an absolute path. | |
| Hardens Gradio Video / Download against: | |
| - missing/empty files after encode | |
| - encode failures with audio (retries silent) | |
| - non-absolute paths Gradio cannot serve | |
| - deep task-only paths that disappear from the UI file server | |
| """ | |
| import shutil | |
| from diffusers.utils import encode_video as _encode_video | |
| _ensure_output_roots() | |
| # Normalize frames to a numpy-friendly object encode_video accepts. | |
| if hasattr(frames, "detach"): | |
| frames = frames.detach().cpu().numpy() | |
| if frames is None: | |
| raise gr.Error("No frames were produced — nothing to encode.") | |
| tid = (task.task_id if task is not None else None) or uuid.uuid4().hex | |
| work = (task.work_dir if task is not None and getattr(task, "work_dir", None) else None) | |
| if work: | |
| try: | |
| os.makedirs(work, exist_ok=True) | |
| except OSError: | |
| work = None | |
| primary = os.path.join(work or _H3_OUTPUT_ROOT, f"h3-{tid}.mp4") | |
| mirror = os.path.join(_H3_OUTPUT_ROOT, f"h3-{tid}.mp4") | |
| def _try_encode(path, with_audio: bool): | |
| kwargs = dict(fps=fps, output_path=path) | |
| if with_audio and audio is not None: | |
| kwargs["audio"] = audio | |
| kwargs["audio_sample_rate"] = sampling_rate | |
| _encode_video(frames, **kwargs) | |
| last_err = None | |
| encoded = False | |
| for with_audio in (True, False): | |
| try: | |
| _try_encode(primary, with_audio=with_audio) | |
| encoded = True | |
| if not with_audio and audio is not None: | |
| # Silent fallback — still deliver a playable file. | |
| pass | |
| break | |
| except Exception as e: | |
| last_err = e | |
| try: | |
| if os.path.isfile(primary): | |
| os.remove(primary) | |
| except OSError: | |
| pass | |
| continue | |
| if not encoded: | |
| raise gr.Error( | |
| f"Failed to mux the video file ({type(last_err).__name__}: {last_err}). " | |
| "Try a shorter duration or fewer steps." | |
| ) | |
| if not os.path.isfile(primary) or os.path.getsize(primary) < 1024: | |
| raise gr.Error( | |
| "Video file was not written correctly (missing or under 1 KB). " | |
| "Please retry the generation." | |
| ) | |
| # Mirror into h3-outputs so Gradio always has a stable root to serve from. | |
| try: | |
| if os.path.realpath(primary) != os.path.realpath(mirror): | |
| shutil.copy2(primary, mirror) | |
| serve_path = mirror | |
| else: | |
| serve_path = primary | |
| except OSError: | |
| serve_path = primary | |
| serve_path = os.path.realpath(serve_path) | |
| if task is not None: | |
| task.output_path = serve_path | |
| return serve_path | |