Files

275 lines
12 KiB
Python

"""Where a clip's bytes are, and whether a browser can play them.
A clip is named the way ComfyUI names a file (filename, subfolder, type) or by a
media-library id. Names come out of a saved project, so they are never trusted as
paths: every one is resolved under the output or temp folder and refused if it
leaves it.
MP4 detail that cost real debugging in v4: ComfyUI's video saver writes the `moov`
index AFTER the picture data. A browser can play such a file from the start but a
deep seek into it is undecodable. So a file is served only when its `moov` is in
front (or has been remuxed to put it there), and a file with no `moov` at all is
still being written -- never served, because the same URL would later return
different bytes under a browser that is mid-stream.
"""
from __future__ import annotations
import hashlib
import os
import shutil
import subprocess
from pathlib import Path
from ..._core import media, video_pipe
ISOBMFF = (".mp4", ".m4v", ".mov") # containers where moov/mdat order applies
REENCODE_CRF = "16" # re-encodes of a finished clip (preview window, trim, stitch): visually clean, not libx264's 23
AUDIO_RATE = 48000 # every exported part is made at this rate, so parts can be joined
ClipError, FFMPEG_MISSING, ffmpeg = video_pipe.ClipError, video_pipe.FFMPEG_MISSING, video_pipe.ffmpeg
def comfy_dir(kind: str) -> str | None:
return video_pipe.comfy_dir(kind)
def comfy_path(filename, subfolder="", kind="output") -> str | None:
return video_pipe.comfy_path(filename, subfolder, kind, comfy_dir) # through this module's comfy_dir, which tests replace
PICTURES = (".png", ".jpg", ".jpeg", ".webp")
STILL_FPS = 24
def is_picture(path: str) -> bool:
return str(path).lower().endswith(PICTURES)
def still_clip(path: str, seconds: float) -> str:
"""A picture file played as a clip: the picture held for `seconds`, made once and kept in ComfyUI's temp folder.
Everything after a render (preview segments, trims, the last frame, export) reads clips, so a picture render becomes one
here, at the one place they all ask for a file. Without sound; the sides are made even for H.264."""
seconds = max(0.1, round(float(seconds or 0), 2))
sig = signature(path)
out = temp_file("funui_still_" + hashlib.md5(f"{path}|{sig[0]}|{sig[1]}|{seconds}".encode()).hexdigest()[:16] + ".mp4")
if not os.path.isfile(out):
run([ffmpeg(), "-y", "-loop", "1", "-framerate", str(STILL_FPS), "-i", path, "-t", f"{seconds:.2f}",
"-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-c:v", "libx264", "-crf", REENCODE_CRF, "-pix_fmt", "yuv420p",
"-movflags", "+faststart", out])
if not os.path.isfile(out) or os.path.getsize(out) == 0:
raise ClipError("a picture could not be turned into a clip")
return out
def clip_path(clip: dict) -> str:
"""The file a clip spec points at: a media-library id, or a ComfyUI file (a picture comes back as a clip of the length asked)."""
ref = clip.get("bin_media_ref")
if ref:
found = media.path_for(ref) if media.is_id(ref) else None
if found is None:
raise ClipError("A media file in this clip is gone from the library: it may have been deleted.")
return str(found)
path = comfy_path(clip.get("filename"), clip.get("subfolder") or "", clip.get("type") or "output")
if not path or not os.path.isfile(path):
raise ClipError("A clip's file is not on disk: generate it again, then try this.")
if is_picture(path):
return still_clip(path, float(clip.get("in") or 0) + float(clip.get("dur") or 1))
return path
def is_isobmff(path: str) -> bool:
return os.path.splitext(path)[1].lower() in ISOBMFF
def moov_position(path: str) -> str:
""""front" (playable and seekable), "end" (complete, but undecodable on a deep seek: remux
it), or "none" (no index anywhere: still being written, or a save was aborted)."""
try:
size = os.stat(path).st_size
seen_mdat = False
offset = 0
with open(path, "rb") as f:
while offset + 8 <= size:
f.seek(offset)
header = f.read(8)
if len(header) < 8:
break
box = int.from_bytes(header[:4], "big")
kind = header[4:8]
if kind == b"moov":
return "end" if seen_mdat else "front"
if kind == b"mdat":
seen_mdat = True
if box == 1: # 64-bit size follows
ext = f.read(8)
if len(ext) < 8:
break
box = int.from_bytes(ext, "big")
elif box == 0: # runs to EOF: an mdat still being streamed out
return "none"
if box < 8:
break
offset += box
except OSError:
pass
return "none"
def signature(path: str) -> tuple:
"""(mtime_ns, size): the same pair means the same bytes, for cache keys."""
st = os.stat(path)
return (st.st_mtime_ns, st.st_size)
def run(cmd: list[str]) -> None:
proc = subprocess.run(cmd, capture_output=True, text=True)
if proc.returncode != 0:
raise ClipError((proc.stderr or "ffmpeg failed")[-1000:])
def remux_faststart(src: str, out: str) -> None:
"""Copy the streams (no re-encode) with the index moved to the front."""
run([ffmpeg(), "-y", "-i", src, "-c", "copy", "-movflags", "+faststart", out])
def trim(src: str, out: str, start=None, dur=None, *, fast=False, reverse=False, fps=None) -> None:
"""Cut [start, start+dur) out of `src` into a seekable h264/aac mp4.
`fast` is for previews, where the latency of a scrub across a clip boundary matters
more than a few percent of bitrate."""
cmd = [ffmpeg(), "-y"]
if start is not None and float(start) > 0:
cmd += ["-ss", f"{float(start):.3f}"]
if dur is not None:
cmd += ["-t", f"{float(dur):.3f}"]
cmd += ["-i", src]
if reverse:
cmd += ["-vf", "reverse", "-af", "areverse"] # after the trim: only the clip is buffered
if fps:
cmd += ["-r", str(fps)] # parts to be joined must share one rate
# A re-encode of a generation: libx264's own default (crf 23) would stack a second lossy pass on the first.
cmd += ["-c:v", "libx264", "-crf", REENCODE_CRF, "-pix_fmt", "yuv420p"]
if fast:
cmd += ["-preset", "veryfast"]
cmd += ["-c:a", "aac", "-b:a", "192k", "-ar", str(AUDIO_RATE), "-ac", "2", "-movflags", "+faststart", out]
run(cmd)
def last_frame(src: str, out_png: str, start=0.0, dur=None, fps=25.0, first=False, at_fraction=None) -> tuple[float, float]:
"""Write the last picture of [start, start+dur) of `src` as a PNG (the first one with `first`: a reversed clip ends where it began);
-> (average brightness, a near-peak brightness), both 0 (black) .. 255.
One frame before the window's end, not at it (a seek to the very end finds nothing), and never past the file's own end.
`at_fraction` (0..1) takes the picture that far into the window instead (0.5: the middle one)."""
end = float(start) + float(dur) if dur else None
at = float(start) if first or end is None else max(float(start), end - 1.5 / (fps or 25.0))
if at_fraction is not None and end is not None:
at = min(at, float(start) + max(0.0, min(1.0, float(at_fraction))) * float(dur))
have = duration(src)
if have and not first:
at = max(0.0, min(at, have - 2.0 / (fps or 25.0)))
run([ffmpeg(), "-y", "-ss", f"{at:.3f}", "-i", src, "-frames:v", "1", out_png])
if not os.path.isfile(out_png) or os.path.getsize(out_png) == 0:
raise ClipError("the clip has no picture at its end (is the window past the end of the render?)")
return judge(out_png)
def judge(out_png: str) -> tuple[float, float]:
"""-> (average brightness, a near-peak brightness) of a picture file, both 0 (black) .. 255."""
try:
from PIL import Image, ImageStat
except ImportError as exc:
raise ClipError("Pillow is needed to judge the last frame and is not installed") from exc
with Image.open(out_png) as im:
gray = im.convert("L")
mean = float(ImageStat.Stat(gray).mean[0])
hist, seen, total = gray.histogram(), 0, gray.width * gray.height
peak = 255
for level, n in enumerate(hist):
seen += n
if seen >= 0.99 * total:
peak = level
break
return mean, float(peak)
def has_audio(path: str) -> bool:
"""True when the file has a sound stream. ffprobe if there is one, else a probe run of ffmpeg."""
probe = shutil.which("ffprobe")
if probe:
proc = subprocess.run([probe, "-v", "error", "-select_streams", "a", "-show_entries",
"stream=index", "-of", "csv=p=0", path], capture_output=True, text=True)
return proc.returncode == 0 and bool(proc.stdout.strip())
proc = subprocess.run([ffmpeg(), "-i", path], capture_output=True, text=True)
return "Audio:" in (proc.stderr or "")
def duration(path: str) -> float | None:
"""Seconds of picture in `path`, or None when it has none (a window past the end of the source
trims to a file with no streams at all) or ffprobe is not here to ask."""
probe = shutil.which("ffprobe")
if not probe:
return None
proc = subprocess.run([probe, "-v", "error", "-select_streams", "v:0", "-show_entries",
"stream=duration:format=duration", "-of", "default=nw=1:nk=1", path],
capture_output=True, text=True)
for line in proc.stdout.split():
try:
return float(line)
except ValueError:
continue
return 0.0 if proc.returncode == 0 else None
def dimensions(path: str) -> tuple[int, int] | None:
"""(width, height) of the first video stream, or None."""
probe = shutil.which("ffprobe")
if not probe:
return None
proc = subprocess.run([probe, "-v", "error", "-select_streams", "v:0", "-show_entries",
"stream=width,height", "-of", "csv=p=0:s=x", path],
capture_output=True, text=True)
try:
w, h = proc.stdout.strip().split("x")[:2]
return (int(w), int(h)) if int(w) > 0 and int(h) > 0 else None
except ValueError:
return None
def frame_rate(path: str) -> str | None:
"""The first video stream's rate as ffprobe spells it ("25/1"), or None."""
probe = shutil.which("ffprobe")
if not probe:
return None
proc = subprocess.run([probe, "-v", "error", "-select_streams", "v:0", "-show_entries",
"stream=r_frame_rate", "-of", "default=nw=1:nk=1", path],
capture_output=True, text=True)
out = proc.stdout.strip()
return out if out and out != "0/0" else None
def add_silence(src: str, out: str) -> None:
"""Copy `src` with a silent stereo track added, so it can be joined to clips that have sound."""
run([ffmpeg(), "-y", "-i", src, "-f", "lavfi", "-i", f"anullsrc=r={AUDIO_RATE}:cl=stereo",
"-shortest", "-map", "0:v:0", "-map", "1:a:0", "-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
"-movflags", "+faststart", out])
def poster(src: str, out_jpg: str, height: int = 96) -> None:
"""One small still from the middle of `src`, for a thumbnail."""
mid = (duration(src) or 0.0) / 2
run([ffmpeg(), "-y", "-ss", f"{mid:.3f}", "-i", src, "-frames:v", "1", "-vf", f"scale=-2:{int(height)}", "-q:v", "5", out_jpg])
if not os.path.isfile(out_jpg) or os.path.getsize(out_jpg) == 0:
raise ClipError("the clip has no picture to take a still from")
def temp_file(name: str) -> str:
base = comfy_dir("temp")
if not base:
raise ClipError("ComfyUI's temp folder is not available here, so nothing can be written for playback.")
Path(base).mkdir(parents=True, exist_ok=True)
return os.path.join(base, name)