mirror of
https://github.com/facefusion/facefusion.git
synced 2026-07-27 12:30:55 +02:00
henryruhs
224 lines
10 KiB
Python
224 lines
10 KiB
Python
import os
|
|
import subprocess
|
|
from typing import Optional, Tuple, cast
|
|
|
|
import numpy
|
|
|
|
from facefusion import ffmpeg_builder, state_manager
|
|
from facefusion.common_helper import is_windows
|
|
from facefusion.ffprobe import extract_video_metadata
|
|
from facefusion.filesystem import get_file_format
|
|
from facefusion.temp_helper import get_temp_file_path
|
|
from facefusion.types import Fps, Resolution, VideoFormat, VideoPoolSet, VideoReader, VideoReaderBuffer, VisionFrame
|
|
|
|
VIDEO_POOL_SET : VideoPoolSet =\
|
|
{
|
|
'reader': {},
|
|
'writer': {}
|
|
}
|
|
|
|
|
|
#todo: needs review - [decoding] [critical: medium] new ffmpeg raw pipe reader with seek based start, stderr devnull hides decode errors
|
|
def create_video_reader_process(video_path : str, frame_position : int, video_fps : Fps) -> subprocess.Popen[bytes]:
|
|
color_transfer = extract_video_metadata(video_path).get('color_transfer')
|
|
commands = ffmpeg_builder.chain(
|
|
ffmpeg_builder.set_input_seek(frame_position / video_fps),
|
|
ffmpeg_builder.set_input(video_path),
|
|
ffmpeg_builder.tonemap_video(color_transfer),
|
|
ffmpeg_builder.prevent_frame_drop(),
|
|
ffmpeg_builder.enforce_pixel_format('bgr24'),
|
|
ffmpeg_builder.set_output_format('rawvideo'),
|
|
ffmpeg_builder.cast_stream()
|
|
)
|
|
commands = ffmpeg_builder.run(commands)
|
|
return subprocess.Popen(commands, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL)
|
|
|
|
|
|
#todo: needs review - [sampling] [critical: medium] select filter decodes the entire range to sample frames, runs at lowered process priority
|
|
def create_video_sampler_process(video_path : str, frame_start : int, frame_end : int, frame_stride : int) -> subprocess.Popen[bytes]:
|
|
color_transfer = extract_video_metadata(video_path).get('color_transfer')
|
|
sample_total = len([ frame_number for frame_number in range(frame_start, frame_end) if frame_number % frame_stride == 0 ])
|
|
sampler_thread_count = 4
|
|
commands = ffmpeg_builder.chain(
|
|
ffmpeg_builder.set_thread_count(sampler_thread_count),
|
|
ffmpeg_builder.set_input(video_path),
|
|
ffmpeg_builder.concat(
|
|
ffmpeg_builder.select_frame_samples(frame_start, frame_end, frame_stride),
|
|
ffmpeg_builder.tonemap_video(color_transfer)
|
|
),
|
|
ffmpeg_builder.prevent_frame_drop(),
|
|
ffmpeg_builder.set_frame_total(sample_total),
|
|
ffmpeg_builder.enforce_pixel_format('bgr24'),
|
|
ffmpeg_builder.set_output_format('rawvideo'),
|
|
ffmpeg_builder.cast_stream()
|
|
)
|
|
commands = ffmpeg_builder.run(commands)
|
|
|
|
if is_windows():
|
|
return subprocess.Popen(commands, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, creationflags = subprocess.IDLE_PRIORITY_CLASS) #type:ignore[attr-defined]
|
|
return subprocess.Popen(commands, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, preexec_fn = demote_process_priority)
|
|
|
|
|
|
#todo: needs review - [sampling] [critical: low] posix only nice fallback for the sampler priority
|
|
def demote_process_priority() -> None:
|
|
os.nice(10)
|
|
|
|
|
|
#todo: needs review - [lifecycle] [critical: medium] pooled reader per path is created without lock and only reaped via clear_video_pool
|
|
def get_video_reader(video_path : str) -> VideoReader:
|
|
if video_path not in VIDEO_POOL_SET.get('reader'):
|
|
video_metadata = extract_video_metadata(video_path)
|
|
width, height = video_metadata.get('resolution')
|
|
|
|
VIDEO_POOL_SET['reader'][video_path] =\
|
|
{
|
|
'process': create_video_reader_process(video_path, 0, video_metadata.get('fps')),
|
|
'video_path': video_path,
|
|
'width': width,
|
|
'height': height,
|
|
'fps': video_metadata.get('fps'),
|
|
'frame_total': video_metadata.get('frame_total'),
|
|
'position': 0,
|
|
'frame_buffer': {}
|
|
}
|
|
|
|
return VIDEO_POOL_SET.get('reader').get(video_path)
|
|
|
|
|
|
#todo: needs review - [seeking] [critical: high] forward skip up to 128 frames by draining the pipe, everything else restarts the ffmpeg process
|
|
def conditional_set_video_reader_position(video_reader : VideoReader, frame_position : int) -> None:
|
|
skip_margin = 128
|
|
skipping = video_reader.get('position') < frame_position and frame_position - video_reader.get('position') <= skip_margin
|
|
|
|
while skipping:
|
|
has_vision_frame, vision_frame = read_video_reader_frame(video_reader)
|
|
skipping = has_vision_frame and video_reader.get('position') < frame_position
|
|
|
|
if not video_reader.get('position') == frame_position:
|
|
restart_video_reader(video_reader, frame_position)
|
|
|
|
|
|
#todo: needs review - [seeking] [critical: high] kill and respawn per out of order seek, frequent restarts stack up on random access
|
|
def restart_video_reader(video_reader : VideoReader, frame_position : int) -> None:
|
|
video_reader.get('process').kill()
|
|
video_reader.get('process').wait()
|
|
video_reader['process'] = create_video_reader_process(video_reader.get('video_path'), frame_position, video_reader.get('fps'))
|
|
video_reader['position'] = frame_position
|
|
video_reader['frame_buffer'].clear()
|
|
|
|
|
|
#todo: needs review - [decoding] [critical: high] partial pipe read returns none and desyncs position from the actual stream
|
|
def read_video_reader_frame(video_reader : VideoReader) -> Tuple[bool, Optional[VisionFrame]]:
|
|
frame_size = video_reader.get('width') * video_reader.get('height') * 3
|
|
frame_buffer = video_reader.get('process').stdout.read(frame_size)
|
|
|
|
if len(frame_buffer) == frame_size:
|
|
video_reader['position'] = video_reader.get('position') + 1
|
|
return True, numpy.frombuffer(frame_buffer, numpy.uint8).reshape(video_reader.get('height'), video_reader.get('width'), 3)
|
|
return False, None
|
|
|
|
|
|
#todo: needs review - [memory] [critical: medium] eviction only trims below frame_start, look ahead growth is bounded by the window callers
|
|
def evict_video_reader_buffer(video_reader : VideoReader, frame_start : int, buffer_margin : int) -> None:
|
|
frame_buffer = video_reader.get('frame_buffer')
|
|
|
|
for frame_number in [ key for key in frame_buffer if key < frame_start - buffer_margin ]:
|
|
del frame_buffer[frame_number]
|
|
|
|
|
|
#todo: needs review - [memory] [critical: high] frame_buffer keeps decoded frames in ram, restart heuristic depends on buffer_margin 16
|
|
def read_video_reader_window(video_reader : VideoReader, frame_start : int, frame_end : int) -> VideoReaderBuffer:
|
|
frame_buffer = video_reader.get('frame_buffer')
|
|
buffer_margin = 16
|
|
read_start = max(frame_start, 0)
|
|
read_end = frame_end
|
|
|
|
if video_reader.get('frame_total') > 0:
|
|
read_end = min(read_end, video_reader.get('frame_total') - 1)
|
|
|
|
if read_start not in frame_buffer and (read_start < video_reader.get('position') or read_start > video_reader.get('position') + buffer_margin):
|
|
restart_video_reader(video_reader, read_start)
|
|
|
|
reading = video_reader.get('position') <= read_end
|
|
|
|
while reading:
|
|
has_vision_frame, vision_frame = read_video_reader_frame(video_reader)
|
|
|
|
if has_vision_frame:
|
|
frame_buffer[video_reader.get('position') - 1] = vision_frame
|
|
reading = has_vision_frame and video_reader.get('position') <= read_end
|
|
|
|
evict_video_reader_buffer(video_reader, read_start, buffer_margin)
|
|
return frame_buffer
|
|
|
|
|
|
#todo: needs review - [encoding] [critical: low] rounds odd dimensions to even for yuv420p compatibility
|
|
def pack_video_resolution(resolution : Resolution) -> str:
|
|
width, height = resolution
|
|
return str(round(width / 2) * 2) + 'x' + str(round(height / 2) * 2)
|
|
|
|
|
|
#todo: needs review - [color] [critical: high] output is always converted and tagged bt709, encoder_thread_count hardcoded to 16
|
|
def create_video_writer_process(target_path : str, temp_video_fps : Fps, temp_video_resolution : Resolution, output_video_resolution : Resolution, output_video_fps : Fps) -> subprocess.Popen[bytes]:
|
|
output_video_encoder = state_manager.get_item('output_video_encoder')
|
|
output_video_quality = state_manager.get_item('output_video_quality')
|
|
output_video_preset = state_manager.get_item('output_video_preset')
|
|
temp_video_path = get_temp_file_path(target_path)
|
|
temp_video_format = cast(VideoFormat, get_file_format(temp_video_path))
|
|
output_video_encoder = ffmpeg_builder.fix_video_encoder(temp_video_format, output_video_encoder)
|
|
encoder_thread_count = 16
|
|
commands = ffmpeg_builder.chain(
|
|
ffmpeg_builder.set_output_format('rawvideo'),
|
|
ffmpeg_builder.enforce_pixel_format('bgr24'),
|
|
ffmpeg_builder.set_media_resolution(pack_video_resolution(temp_video_resolution)),
|
|
ffmpeg_builder.set_input_fps(temp_video_fps),
|
|
ffmpeg_builder.set_input('pipe:0'),
|
|
ffmpeg_builder.set_media_resolution(pack_video_resolution(output_video_resolution)),
|
|
ffmpeg_builder.set_video_encoder(output_video_encoder),
|
|
ffmpeg_builder.set_thread_count(encoder_thread_count),
|
|
ffmpeg_builder.set_video_tag(output_video_encoder, temp_video_format),
|
|
ffmpeg_builder.set_video_quality(output_video_encoder, output_video_quality),
|
|
ffmpeg_builder.set_video_preset(output_video_encoder, output_video_preset),
|
|
ffmpeg_builder.concat(
|
|
ffmpeg_builder.set_video_fps(output_video_fps),
|
|
ffmpeg_builder.convert_colorspace('bt709')
|
|
),
|
|
ffmpeg_builder.set_pixel_format(output_video_encoder),
|
|
ffmpeg_builder.force_output(temp_video_path)
|
|
)
|
|
commands = ffmpeg_builder.run(commands)
|
|
return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.DEVNULL)
|
|
|
|
|
|
#todo: needs review - [lifecycle] [critical: low] pooled writer keyed by target_path
|
|
def get_video_writer(target_path : str, temp_video_fps : Fps, temp_video_resolution : Resolution, output_video_resolution : Resolution, output_video_fps : Fps) -> subprocess.Popen[bytes]:
|
|
if target_path not in VIDEO_POOL_SET.get('writer'):
|
|
VIDEO_POOL_SET['writer'][target_path] = create_video_writer_process(target_path, temp_video_fps, temp_video_resolution, output_video_resolution, output_video_fps)
|
|
|
|
return VIDEO_POOL_SET.get('writer').get(target_path)
|
|
|
|
|
|
#todo: needs review - [encoding] [critical: medium] blocking stdin write without backpressure or broken pipe handling
|
|
def write_video_writer_frame(video_writer : subprocess.Popen[bytes], vision_frame : VisionFrame) -> None:
|
|
video_writer.stdin.write(vision_frame.tobytes())
|
|
|
|
|
|
#todo: needs review - [encoding] [critical: medium] encoder failures only surface via returncode at close time
|
|
def close_video_writer(video_writer : subprocess.Popen[bytes]) -> bool:
|
|
video_writer.stdin.close()
|
|
video_writer.wait()
|
|
return video_writer.returncode == 0
|
|
|
|
|
|
#todo: needs review - [lifecycle] [critical: medium] terminates writers mid encode, readers are terminated without reaping stdout
|
|
def clear_video_pool() -> None:
|
|
for video_writer in VIDEO_POOL_SET.get('writer').values():
|
|
video_writer.terminate()
|
|
video_writer.wait()
|
|
|
|
for video_reader in VIDEO_POOL_SET.get('reader').values():
|
|
video_reader.get('process').terminate()
|
|
|
|
VIDEO_POOL_SET['reader'].clear()
|
|
VIDEO_POOL_SET['writer'].clear()
|