200 lines
8.8 KiB
Python
200 lines
8.8 KiB
Python
"""Stream a video file through a per-chunk function into a new file, keeping its sound.
|
|
|
|
Frames go through ffmpeg in small chunks (never the whole clip in RAM). Shared by the post-render
|
|
steps that rewrite a finished video (upscale, interpolate), which cannot import each other.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
|
|
CHUNK = 8 # frames per call, at most
|
|
CHUNK_PIXELS = 2_000_000 # ...and at most this many input pixels per call: a big upscale of 1080p is gigabytes per frame
|
|
|
|
FFMPEG_MISSING = "ffmpeg was not found on PATH: install it to preview, export or render clips."
|
|
|
|
|
|
class ClipError(Exception):
|
|
"""A clip cannot be used, with the sentence the person should read."""
|
|
|
|
|
|
def ffmpeg() -> str:
|
|
path = shutil.which("ffmpeg")
|
|
if not path:
|
|
raise ClipError(FFMPEG_MISSING)
|
|
return path
|
|
|
|
|
|
def comfy_dir(kind: str) -> str | None:
|
|
"""ComfyUI's output or temp folder, or None when ComfyUI is not here to ask."""
|
|
try:
|
|
import folder_paths
|
|
return folder_paths.get_output_directory() if kind == "output" else folder_paths.get_temp_directory()
|
|
except Exception: # noqa: BLE001 -- tests, headless
|
|
return None
|
|
|
|
|
|
def comfy_path(filename, subfolder="", kind="output", dir_of=None) -> str | None:
|
|
"""The file ComfyUI calls (filename, subfolder, kind), or None if the name does not
|
|
resolve to a place under that folder."""
|
|
if not isinstance(filename, str) or not filename:
|
|
return None
|
|
base = (dir_of or comfy_dir)("output" if kind != "temp" else "temp")
|
|
if not base:
|
|
return None
|
|
base = os.path.realpath(base)
|
|
full = os.path.realpath(os.path.join(base, subfolder if isinstance(subfolder, str) else "",
|
|
os.path.basename(filename)))
|
|
return full if os.path.commonpath((full, base)) == base else None
|
|
|
|
|
|
def probe(path):
|
|
"""-> (width, height, fps string, has sound)."""
|
|
ffprobe = shutil.which("ffprobe")
|
|
if not ffprobe:
|
|
raise ClipError("ffprobe was not found on PATH: install ffmpeg to upscale videos.")
|
|
out = subprocess.run([ffprobe, "-v", "error", "-show_entries",
|
|
"stream=codec_type,width,height,r_frame_rate", "-of", "json", path],
|
|
capture_output=True, text=True)
|
|
if out.returncode:
|
|
raise ClipError(f"{os.path.basename(path)} could not be read: {out.stderr.strip()[-300:] or 'ffprobe failed'}")
|
|
streams = json.loads(out.stdout).get("streams") or []
|
|
video = next((s for s in streams if s.get("codec_type") == "video"), None)
|
|
if not video:
|
|
raise ClipError(f"no video stream in {os.path.basename(path)}")
|
|
return (int(video["width"]), int(video["height"]), video.get("r_frame_rate") or "25",
|
|
any(s.get("codec_type") == "audio" for s in streams))
|
|
|
|
|
|
def _read(stream, n):
|
|
buf = bytearray()
|
|
while len(buf) < n:
|
|
part = stream.read(n - len(buf))
|
|
if not part:
|
|
break
|
|
buf += part
|
|
return bytes(buf)
|
|
|
|
|
|
def scaled_fps(fps, k):
|
|
""""24/1" -> "48/1", "30000/1001" -> "60000/1001": the frame rate after every frame became k."""
|
|
num, _, den = str(fps).partition("/")
|
|
return f"{int(num) * k}/{den or 1}"
|
|
|
|
|
|
def stream_file(src, dst, upscale, interrupted=lambda: None, fps_factor=1):
|
|
"""Stream `src` through `upscale` ([B,H,W,3] float 0-1 -> same, larger, or more frames when `fps_factor` > 1)
|
|
into `dst`, played at `fps_factor` times the source's frame rate so the length and the sound stay in step.
|
|
Returns the frames written. `interrupted()` raises to stop between chunks."""
|
|
import numpy as np
|
|
import torch
|
|
w, h, fps, has_sound = probe(src)
|
|
frame = w * h * 3
|
|
chunk = max(1, min(CHUNK, CHUNK_PIXELS // max(1, w * h)))
|
|
exe = ffmpeg()
|
|
# The decoder's complaints go to a file, not away: ffmpeg exits 0 on a video that was cut short, and the only
|
|
# sign is what it printed ("partial file"). Read after the last frame; a file avoids a full pipe stalling it.
|
|
problems = tempfile.TemporaryFile()
|
|
dec = subprocess.Popen([exe, "-v", "error", "-i", src, "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
|
|
stdout=subprocess.PIPE, stderr=problems)
|
|
enc, frames = None, 0
|
|
try:
|
|
while True:
|
|
interrupted()
|
|
raw = _read(dec.stdout, frame * chunk)
|
|
n = len(raw) // frame
|
|
if n == 0:
|
|
break
|
|
x = torch.from_numpy(np.frombuffer(raw[:n * frame], dtype=np.uint8)
|
|
.reshape(n, h, w, 3).copy()).float() / 255.0
|
|
y = (upscale(x).clamp(0, 1) * 255.0).round().to(torch.uint8).cpu().numpy()
|
|
if enc is None:
|
|
oh, ow = y.shape[1], y.shape[2]
|
|
cmd = [exe, "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
|
|
"-s", f"{ow}x{oh}", "-r", scaled_fps(fps, fps_factor), "-i", "-"]
|
|
if has_sound:
|
|
cmd += ["-i", src, "-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac", "-b:a", "192k"]
|
|
cmd += ["-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-c:v", "libx264", "-crf", "16",
|
|
"-pix_fmt", "yuv420p", "-movflags", "+faststart", dst]
|
|
enc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
try:
|
|
enc.stdin.write(y.tobytes())
|
|
except BrokenPipeError: # the encoder died: its own message says why
|
|
enc.wait()
|
|
raise ClipError(f"ffmpeg could not encode the video: {enc.stderr.read().decode(errors='replace')[-500:]}") from None
|
|
frames += y.shape[0]
|
|
tail = y[-1:]
|
|
if enc is not None and fps_factor > 1: # the last picture has no partner to blend with: hold it, so frames = source x factor and the length stays
|
|
hold = np.repeat(tail, fps_factor - 1, axis=0)
|
|
try:
|
|
enc.stdin.write(hold.tobytes())
|
|
except BrokenPipeError:
|
|
enc.wait()
|
|
raise ClipError(f"ffmpeg could not encode the video: {enc.stderr.read().decode(errors='replace')[-500:]}") from None
|
|
frames += hold.shape[0]
|
|
dec.wait()
|
|
problems.seek(0)
|
|
said = problems.read().decode(errors="replace").strip()
|
|
if dec.returncode or said:
|
|
raise ClipError(f"{os.path.basename(src)} is damaged or cut short (read {frames} frame(s) before the problem): {said[-300:] or 'ffmpeg failed'}")
|
|
if enc is None:
|
|
raise ClipError(f"no frames could be read from {os.path.basename(src)}")
|
|
interrupted() # a stop that came during the last read is still a stop
|
|
enc.stdin.close()
|
|
err = enc.stderr.read().decode(errors="replace")
|
|
if enc.wait() != 0:
|
|
raise ClipError(f"ffmpeg could not encode the video: {err[-500:]}")
|
|
interrupted() # ... and one that came while the encoder flushed: the file is removed below
|
|
problems.close()
|
|
return frames
|
|
except BaseException:
|
|
problems.close()
|
|
for p in (dec, enc):
|
|
if p is not None and p.poll() is None:
|
|
p.kill()
|
|
if p is not None:
|
|
p.wait()
|
|
for pipe in (p.stdin, p.stdout, p.stderr):
|
|
try:
|
|
if pipe:
|
|
pipe.close()
|
|
except OSError: # a killed encoder's unflushed pipe
|
|
pass
|
|
if os.path.exists(dst):
|
|
os.remove(dst)
|
|
raise
|
|
finally:
|
|
if dec.poll() is None:
|
|
dec.kill()
|
|
|
|
|
|
def last_frame_name(video_file: str) -> str:
|
|
"""The lossless copy of a render's last picture, kept in ComfyUI's temp folder (Save Video writes it; the carry reads it)."""
|
|
return "funui_last_" + os.path.splitext(os.path.basename(video_file))[0] + ".png"
|
|
|
|
|
|
def inherit_last_frame(old_file: str, new_file: str) -> None:
|
|
"""A file made from `old_file` (upscaled, interpolated) continues from the same picture: the one the sampler made, taken before
|
|
this step touched it. The copy is under the new file's name so the carry finds it. Nothing to inherit, or no temp folder: nothing happens."""
|
|
try:
|
|
folder = comfy_dir("temp")
|
|
src = os.path.join(folder or "", last_frame_name(old_file))
|
|
if folder and os.path.isfile(src):
|
|
shutil.copyfile(src, os.path.join(folder, last_frame_name(new_file)))
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def free_name(out_dir, filename, model):
|
|
stem = os.path.splitext(os.path.basename(filename))[0]
|
|
tag = os.path.splitext(os.path.basename(model))[0]
|
|
n, name = 0, f"{stem}_{tag}.mp4"
|
|
while os.path.exists(os.path.join(out_dir, name)):
|
|
n += 1
|
|
name = f"{stem}_{tag}_{n}.mp4"
|
|
return name
|