mirror of
https://github.com/facefusion/facefusion.git
synced 2026-07-27 12:30:55 +02:00
Add --workflow-mode (disk|stream), drop --keep-temp, ffmpeg video writer
Two selectable workflows share one swap loop (process_target_frame): - disk (default): extract -> process -> merge; temp frames on disk, supports fps resampling, robust/inspectable. - stream: read -> swap -> encode over pipes; no temp frames, ~2x faster wall. Remove --keep-temp entirely (arg, config, state, ui, locale). Replace the dead cv2.VideoWriter in video_manager with an ffmpeg rawvideo encoder held in the pool (symmetric to the reader): get_video_writer / write_video_writer_frame / close_video_writer, with fix_video_encoder relocated to ffmpeg_builder. The encoder must be closed before processor post_process(), which clears the video pool and would otherwise SIGTERM the still-open writer. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
henryruhs
+1
-1
@@ -51,7 +51,7 @@ voice_extractor_model =
|
||||
trim_frame_start =
|
||||
trim_frame_end =
|
||||
temp_frame_format =
|
||||
keep_temp =
|
||||
workflow_mode =
|
||||
|
||||
[frame_distribution]
|
||||
target_frame_amount =
|
||||
|
||||
+1
-1
@@ -45,7 +45,7 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None:
|
||||
apply_state_item('trim_frame_start', args.get('trim_frame_start'))
|
||||
apply_state_item('trim_frame_end', args.get('trim_frame_end'))
|
||||
apply_state_item('temp_frame_format', args.get('temp_frame_format'))
|
||||
apply_state_item('keep_temp', args.get('keep_temp'))
|
||||
apply_state_item('workflow_mode', args.get('workflow_mode'))
|
||||
apply_state_item('target_frame_amount', args.get('target_frame_amount'))
|
||||
apply_state_item('output_image_quality', args.get('output_image_quality'))
|
||||
apply_state_item('output_image_scale', args.get('output_image_scale'))
|
||||
|
||||
@@ -2,7 +2,7 @@ import logging
|
||||
from typing import List, Sequence, get_args
|
||||
|
||||
from facefusion.common_helper import create_float_range, create_int_range
|
||||
from facefusion.types import Angle, AudioEncoder, AudioFormat, AudioTypeSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, EncoderSet, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorGender, FaceSelectorMode, FaceSelectorOrder, FaceSelectorRace, Gender, ImageFormat, ImageTypeSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, UiWorkflow, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoTypeSet, VoiceExtractorModel
|
||||
from facefusion.types import Angle, AudioEncoder, AudioFormat, AudioTypeSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, EncoderSet, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorGender, FaceSelectorMode, FaceSelectorOrder, FaceSelectorRace, Gender, ImageFormat, ImageTypeSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, UiWorkflow, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoTypeSet, VoiceExtractorModel, WorkflowMode
|
||||
|
||||
face_detector_set : FaceDetectorSet =\
|
||||
{
|
||||
@@ -82,6 +82,7 @@ video_formats : List[VideoFormat] = list(get_args(VideoFormat))
|
||||
temp_frame_formats : List[TempFrameFormat] = list(get_args(TempFrameFormat))
|
||||
|
||||
output_audio_encoders : List[AudioEncoder] = list(get_args(AudioEncoder))
|
||||
workflow_modes : List[WorkflowMode] = list(get_args(WorkflowMode))
|
||||
output_video_encoders : List[VideoEncoder] = list(get_args(VideoEncoder))
|
||||
output_encoder_set : EncoderSet =\
|
||||
{
|
||||
|
||||
+2
-39
@@ -10,7 +10,7 @@ import facefusion.choices
|
||||
from facefusion import ffmpeg_builder, logger, process_manager, state_manager, translator
|
||||
from facefusion.filesystem import get_file_format, remove_file
|
||||
from facefusion.temp_helper import get_temp_file_path, get_temp_frame_pattern
|
||||
from facefusion.types import AudioBuffer, AudioEncoder, Command, EncoderSet, Fps, Resolution, UpdateProgress, VideoEncoder, VideoFormat
|
||||
from facefusion.types import AudioBuffer, AudioEncoder, Command, EncoderSet, Fps, Resolution, UpdateProgress, VideoFormat
|
||||
from facefusion.vision import detect_video_duration, detect_video_fps, pack_resolution, predict_video_frame_total
|
||||
|
||||
|
||||
@@ -70,33 +70,6 @@ def open_ffmpeg(commands : List[Command]) -> subprocess.Popen[bytes]:
|
||||
return subprocess.Popen(commands, stdin = subprocess.PIPE, stdout = subprocess.PIPE)
|
||||
|
||||
|
||||
def open_video_encoder(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 = fix_video_encoder(temp_video_format, output_video_encoder)
|
||||
commands = ffmpeg_builder.chain(
|
||||
ffmpeg_builder.set_output_format('rawvideo'),
|
||||
ffmpeg_builder.enforce_pixel_format('bgr24'),
|
||||
ffmpeg_builder.set_media_resolution(pack_resolution(temp_video_resolution)),
|
||||
ffmpeg_builder.set_input_fps(temp_video_fps),
|
||||
ffmpeg_builder.set_input('pipe:0'),
|
||||
ffmpeg_builder.set_media_resolution(pack_resolution(output_video_resolution)),
|
||||
ffmpeg_builder.set_video_encoder(output_video_encoder),
|
||||
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.set_video_fps(output_video_fps),
|
||||
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)
|
||||
|
||||
|
||||
def log_debug(process : subprocess.Popen[bytes]) -> None:
|
||||
_, stderr = process.communicate()
|
||||
errors = stderr.decode().split(os.linesep)
|
||||
@@ -254,7 +227,7 @@ def merge_video(target_path : str, temp_video_fps : Fps, output_video_resolution
|
||||
temp_video_format = cast(VideoFormat, get_file_format(temp_video_path))
|
||||
temp_frame_pattern = get_temp_frame_pattern(target_path, '%08d')
|
||||
|
||||
output_video_encoder = fix_video_encoder(temp_video_format, output_video_encoder)
|
||||
output_video_encoder = ffmpeg_builder.fix_video_encoder(temp_video_format, output_video_encoder)
|
||||
commands = ffmpeg_builder.chain(
|
||||
ffmpeg_builder.set_input_fps(temp_video_fps),
|
||||
ffmpeg_builder.set_start_number(trim_frame_start),
|
||||
@@ -317,13 +290,3 @@ def fix_audio_encoder(video_format : VideoFormat, audio_encoder : AudioEncoder)
|
||||
return audio_encoder
|
||||
|
||||
|
||||
def fix_video_encoder(video_format : VideoFormat, video_encoder : VideoEncoder) -> VideoEncoder:
|
||||
if video_format in [ 'm4v', 'mpeg', 'mxf', 'wmv' ]:
|
||||
return 'libx264'
|
||||
if video_format in [ 'mkv', 'mp4' ] and video_encoder == 'rawvideo':
|
||||
return 'libx264'
|
||||
if video_format == 'mov' and video_encoder == 'libvpx-vp9':
|
||||
return 'libx264'
|
||||
if video_format == 'webm':
|
||||
return 'libvpx-vp9'
|
||||
return video_encoder
|
||||
|
||||
@@ -199,6 +199,18 @@ def copy_video_encoder() -> List[Command]:
|
||||
return set_video_encoder('copy')
|
||||
|
||||
|
||||
def fix_video_encoder(video_format : VideoFormat, video_encoder : VideoEncoder) -> VideoEncoder:
|
||||
if video_format in [ 'm4v', 'mpeg', 'mxf', 'wmv' ]:
|
||||
return 'libx264'
|
||||
if video_format in [ 'mkv', 'mp4' ] and video_encoder == 'rawvideo':
|
||||
return 'libx264'
|
||||
if video_format == 'mov' and video_encoder == 'libvpx-vp9':
|
||||
return 'libx264'
|
||||
if video_format == 'webm':
|
||||
return 'libvpx-vp9'
|
||||
return video_encoder
|
||||
|
||||
|
||||
def set_faststart(video_format : VideoFormat) -> List[Command]:
|
||||
if video_format in [ 'm4v', 'mov', 'mp4' ]:
|
||||
return [ '-movflags', '+faststart' ]
|
||||
|
||||
@@ -136,7 +136,7 @@ LOCALES : Locales =\
|
||||
'trim_frame_start': 'specify the starting frame of the target video',
|
||||
'trim_frame_end': 'specify the ending frame of the target video',
|
||||
'temp_frame_format': 'specify the temporary resources format',
|
||||
'keep_temp': 'keep the temporary resources after processing',
|
||||
'workflow_mode': 'specify how frames flow through processing (disk extracts and merges temp frames, stream pipes frames through ffmpeg)',
|
||||
'target_frame_amount': 'specify the amount of target frames forwarded to the processor',
|
||||
'output_image_quality': 'specify the image quality which translates to the image compression',
|
||||
'output_image_scale': 'specify the image scale based on the target image',
|
||||
|
||||
@@ -169,8 +169,8 @@ def create_frame_extraction_program() -> ArgumentParser:
|
||||
group_frame_extraction.add_argument('--trim-frame-start', help = translator.get('help.trim_frame_start'), type = int, default = facefusion.config.get_int_value('frame_extraction', 'trim_frame_start'))
|
||||
group_frame_extraction.add_argument('--trim-frame-end', help = translator.get('help.trim_frame_end'), type = int, default = facefusion.config.get_int_value('frame_extraction', 'trim_frame_end'))
|
||||
group_frame_extraction.add_argument('--temp-frame-format', help = translator.get('help.temp_frame_format'), default = config.get_str_value('frame_extraction', 'temp_frame_format', 'png'), choices = facefusion.choices.temp_frame_formats)
|
||||
group_frame_extraction.add_argument('--keep-temp', help = translator.get('help.keep_temp'), action = 'store_true', default = config.get_bool_value('frame_extraction', 'keep_temp'))
|
||||
job_store.register_step_keys([ 'trim_frame_start', 'trim_frame_end', 'temp_frame_format', 'keep_temp' ])
|
||||
group_frame_extraction.add_argument('--workflow-mode', help = translator.get('help.workflow_mode'), default = config.get_str_value('frame_extraction', 'workflow_mode', 'disk'), choices = facefusion.choices.workflow_modes)
|
||||
job_store.register_step_keys([ 'trim_frame_start', 'trim_frame_end', 'temp_frame_format', 'workflow_mode' ])
|
||||
return program
|
||||
|
||||
|
||||
|
||||
@@ -43,7 +43,5 @@ def create_temp_directory(file_path : str) -> bool:
|
||||
|
||||
|
||||
def clear_temp_directory(file_path : str) -> bool:
|
||||
if not state_manager.get_item('keep_temp'):
|
||||
temp_directory_path = get_temp_directory_path(file_path)
|
||||
return remove_directory(temp_directory_path)
|
||||
return True
|
||||
temp_directory_path = get_temp_directory_path(file_path)
|
||||
return remove_directory(temp_directory_path)
|
||||
|
||||
+4
-3
@@ -65,7 +65,7 @@ Locales : TypeAlias = Dict[Language, Dict[str, Any]]
|
||||
LocalePoolSet : TypeAlias = Dict[str, Locales]
|
||||
|
||||
VideoCaptureSet : TypeAlias = Dict[str, cv2.VideoCapture]
|
||||
VideoWriterSet : TypeAlias = Dict[str, cv2.VideoWriter]
|
||||
VideoWriterSet : TypeAlias = Dict[str, subprocess.Popen]
|
||||
CameraCaptureSet : TypeAlias = Dict[str, cv2.VideoCapture]
|
||||
CameraPoolSet = TypedDict('CameraPoolSet',
|
||||
{
|
||||
@@ -169,6 +169,7 @@ AudioFormat = Literal['flac', 'm4a', 'mp3', 'ogg', 'opus', 'wav']
|
||||
ImageFormat = Literal['bmp', 'jpeg', 'png', 'tiff', 'webp']
|
||||
VideoFormat = Literal['avi', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mxf', 'webm', 'wmv']
|
||||
TempFrameFormat = Literal['bmp', 'jpeg', 'png', 'tiff']
|
||||
WorkflowMode = Literal['disk', 'stream']
|
||||
AudioTypeSet : TypeAlias = Dict[AudioFormat, str]
|
||||
ImageTypeSet : TypeAlias = Dict[ImageFormat, str]
|
||||
VideoTypeSet : TypeAlias = Dict[VideoFormat, str]
|
||||
@@ -339,7 +340,7 @@ StateKey = Literal\
|
||||
'trim_frame_start',
|
||||
'trim_frame_end',
|
||||
'temp_frame_format',
|
||||
'keep_temp',
|
||||
'workflow_mode',
|
||||
'target_frame_amount',
|
||||
'output_image_quality',
|
||||
'output_image_scale',
|
||||
@@ -410,7 +411,7 @@ State = TypedDict('State',
|
||||
'trim_frame_start' : int,
|
||||
'trim_frame_end' : int,
|
||||
'temp_frame_format' : TempFrameFormat,
|
||||
'keep_temp' : bool,
|
||||
'workflow_mode' : WorkflowMode,
|
||||
'target_frame_amount' : int,
|
||||
'output_image_quality' : int,
|
||||
'output_image_scale' : Scale,
|
||||
|
||||
@@ -6,7 +6,7 @@ from facefusion.uis.types import JobManagerAction, JobRunnerAction, PreviewMode
|
||||
job_manager_actions : List[JobManagerAction] = [ 'job-create', 'job-submit', 'job-delete', 'job-add-step', 'job-remix-step', 'job-insert-step', 'job-remove-step' ]
|
||||
job_runner_actions : List[JobRunnerAction] = [ 'job-run', 'job-run-all', 'job-retry', 'job-retry-all' ]
|
||||
|
||||
common_options : List[str] = [ 'keep-temp' ]
|
||||
common_options : List[str] = []
|
||||
|
||||
preview_modes : List[PreviewMode] = [ 'default', 'frame-by-frame', 'face-by-face' ]
|
||||
preview_resolutions : List[str] = [ '512x512', '768x768', '1024x1024' ]
|
||||
|
||||
@@ -2,7 +2,7 @@ from typing import List, Optional
|
||||
|
||||
import gradio
|
||||
|
||||
from facefusion import state_manager, translator
|
||||
from facefusion import translator
|
||||
from facefusion.uis import choices as uis_choices
|
||||
|
||||
COMMON_OPTIONS_CHECKBOX_GROUP : Optional[gradio.Checkboxgroup] = None
|
||||
@@ -13,9 +13,6 @@ def render() -> None:
|
||||
|
||||
common_options = []
|
||||
|
||||
if state_manager.get_item('keep_temp'):
|
||||
common_options.append('keep-temp')
|
||||
|
||||
COMMON_OPTIONS_CHECKBOX_GROUP = gradio.Checkboxgroup(
|
||||
label = translator.get('uis.common_options_checkbox_group'),
|
||||
choices = uis_choices.common_options,
|
||||
@@ -28,5 +25,4 @@ def listen() -> None:
|
||||
|
||||
|
||||
def update(common_options : List[str]) -> None:
|
||||
keep_temp = 'keep-temp' in common_options
|
||||
state_manager.set_item('keep_temp', keep_temp)
|
||||
pass
|
||||
|
||||
+50
-10
@@ -1,12 +1,14 @@
|
||||
import subprocess
|
||||
from typing import Optional, Tuple
|
||||
from typing import Optional, Tuple, cast
|
||||
|
||||
import cv2
|
||||
import numpy
|
||||
|
||||
from facefusion import ffmpeg_builder
|
||||
from facefusion import ffmpeg_builder, state_manager
|
||||
from facefusion.ffprobe import extract_video_metadata
|
||||
from facefusion.types import Fps, VideoPoolSet, VideoReader, VideoReaderBuffer, VisionFrame
|
||||
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 =\
|
||||
{
|
||||
@@ -119,14 +121,52 @@ def read_video_reader_window(video_reader : VideoReader, frame_start : int, fram
|
||||
return frame_buffer
|
||||
|
||||
|
||||
def get_video_writer(video_path : str) -> cv2.VideoWriter:
|
||||
if video_path not in VIDEO_POOL_SET.get('writer'):
|
||||
video_writer = cv2.VideoWriter()
|
||||
def pack_video_resolution(resolution : Resolution) -> str:
|
||||
width, height = resolution
|
||||
return str(round(width / 2) * 2) + 'x' + str(round(height / 2) * 2)
|
||||
|
||||
if video_writer.isOpened():
|
||||
VIDEO_POOL_SET['writer'][video_path] = video_writer
|
||||
|
||||
return VIDEO_POOL_SET.get('writer').get(video_path)
|
||||
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)
|
||||
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_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.set_video_fps(output_video_fps),
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def write_video_writer_frame(video_writer : subprocess.Popen[bytes], vision_frame : VisionFrame) -> None:
|
||||
video_writer.stdin.write(vision_frame.tobytes())
|
||||
|
||||
|
||||
def close_video_writer(video_writer : subprocess.Popen[bytes]) -> bool:
|
||||
video_writer.stdin.close()
|
||||
video_writer.wait()
|
||||
return video_writer.returncode == 0
|
||||
|
||||
|
||||
def clear_video_pool() -> None:
|
||||
@@ -134,7 +174,7 @@ def clear_video_pool() -> None:
|
||||
video_capture.release()
|
||||
|
||||
for video_writer in VIDEO_POOL_SET.get('writer').values():
|
||||
video_writer.release()
|
||||
video_writer.terminate()
|
||||
|
||||
for video_reader in VIDEO_POOL_SET.get('reader').values():
|
||||
video_reader.get('process').terminate()
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from functools import partial
|
||||
from typing import Tuple
|
||||
from typing import List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy
|
||||
@@ -13,10 +13,10 @@ from facefusion.common_helper import get_first, get_middle
|
||||
from facefusion.content_analyser import analyse_video
|
||||
from facefusion.filesystem import filter_audio_paths, is_video
|
||||
from facefusion.processors.core import get_processors_modules
|
||||
from facefusion.temp_helper import clear_temp_directory, create_temp_directory, move_temp_file
|
||||
from facefusion.temp_helper import clear_temp_directory, create_temp_directory, move_temp_file, resolve_temp_frame_set
|
||||
from facefusion.time_helper import calculate_end_time
|
||||
from facefusion.types import ErrorCode, Resolution, VisionFrame
|
||||
from facefusion.vision import conditional_merge_vision_mask, detect_video_resolution, extract_vision_mask, pack_resolution, read_static_images, read_static_video_frame, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames
|
||||
from facefusion.types import ErrorCode, FrameSet, Resolution, VisionFrame
|
||||
from facefusion.vision import conditional_merge_vision_mask, detect_video_resolution, extract_vision_mask, pack_resolution, read_static_image, read_static_images, read_static_video_frame, restrict_trim_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames, write_image
|
||||
from facefusion.workflows.core import is_process_stopping
|
||||
|
||||
|
||||
@@ -24,10 +24,21 @@ def process(start_time : float) -> ErrorCode:
|
||||
tasks =\
|
||||
[
|
||||
setup,
|
||||
process_video,
|
||||
extract_frames,
|
||||
process_disk_frames,
|
||||
merge_frames,
|
||||
restore_audio,
|
||||
partial(finalize_video, start_time)
|
||||
]
|
||||
|
||||
if state_manager.get_item('workflow_mode') == 'stream':
|
||||
tasks =\
|
||||
[
|
||||
setup,
|
||||
process_stream_frames,
|
||||
restore_audio,
|
||||
partial(finalize_video, start_time)
|
||||
]
|
||||
process_manager.start()
|
||||
|
||||
for task in tasks:
|
||||
@@ -73,7 +84,7 @@ def extract_frames() -> ErrorCode:
|
||||
return 0
|
||||
|
||||
|
||||
def process_video() -> ErrorCode:
|
||||
def process_stream_frames() -> ErrorCode:
|
||||
trim_frame_start, trim_frame_end = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end'))
|
||||
output_video_resolution = scale_resolution(detect_video_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale'))
|
||||
temp_video_resolution = restrict_video_resolution(state_manager.get_item('target_path'), output_video_resolution)
|
||||
@@ -81,7 +92,7 @@ def process_video() -> ErrorCode:
|
||||
frame_range = range(trim_frame_start, trim_frame_end)
|
||||
|
||||
if frame_range:
|
||||
video_encoder = ffmpeg.open_video_encoder(state_manager.get_item('target_path'), temp_video_fps, temp_video_resolution, output_video_resolution, state_manager.get_item('output_video_fps'))
|
||||
video_writer = video_manager.get_video_writer(state_manager.get_item('target_path'), temp_video_fps, temp_video_resolution, output_video_resolution, state_manager.get_item('output_video_fps'))
|
||||
|
||||
with tqdm(total = len(frame_range), desc = translator.get('processing'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress:
|
||||
progress.set_postfix(execution_providers = state_manager.get_item('execution_providers'))
|
||||
@@ -95,24 +106,23 @@ def process_video() -> ErrorCode:
|
||||
|
||||
if len(futures) >= frame_look_ahead and is_process_stopping() is False:
|
||||
_, vision_frame = futures.pop(0).result()
|
||||
video_encoder.stdin.write(vision_frame.tobytes())
|
||||
video_manager.write_video_writer_frame(video_writer, vision_frame)
|
||||
progress.update()
|
||||
|
||||
for future in futures:
|
||||
if is_process_stopping() is False:
|
||||
_, vision_frame = future.result()
|
||||
video_encoder.stdin.write(vision_frame.tobytes())
|
||||
video_manager.write_video_writer_frame(video_writer, vision_frame)
|
||||
progress.update()
|
||||
|
||||
close_succeeded = video_manager.close_video_writer(video_writer)
|
||||
|
||||
for processor_module in get_processors_modules(state_manager.get_item('processors')):
|
||||
processor_module.post_process()
|
||||
|
||||
video_encoder.stdin.close()
|
||||
video_encoder.wait()
|
||||
|
||||
if is_process_stopping():
|
||||
return 4
|
||||
if not video_encoder.returncode == 0:
|
||||
if not close_succeeded:
|
||||
return 1
|
||||
else:
|
||||
logger.error(translator.get('temp_frames_not_found'), __name__)
|
||||
@@ -167,18 +177,12 @@ def restore_audio() -> ErrorCode:
|
||||
return 0
|
||||
|
||||
|
||||
def process_stream_frame(frame_number : int, temp_video_resolution : Resolution) -> Tuple[int, VisionFrame]:
|
||||
def process_target_frame(frame_number : int, target_vision_frames : List[VisionFrame], temp_vision_frame : VisionFrame) -> VisionFrame:
|
||||
trim_frame_start, _ = restrict_trim_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end'))
|
||||
reference_vision_frame = read_static_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number'))
|
||||
source_vision_frames = read_static_images(state_manager.get_item('source_paths'))
|
||||
source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths')))
|
||||
target_vision_frames = select_video_frames(state_manager.get_item('target_path'), frame_number, state_manager.get_item('target_frame_amount'))
|
||||
target_vision_frame = get_middle(target_vision_frames)
|
||||
temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps'))
|
||||
temp_vision_frame = target_vision_frame.copy()
|
||||
|
||||
if not (target_vision_frame.shape[1], target_vision_frame.shape[0]) == temp_video_resolution:
|
||||
temp_vision_frame = cv2.resize(target_vision_frame, temp_video_resolution)
|
||||
temp_vision_mask = extract_vision_mask(temp_vision_frame)
|
||||
|
||||
source_audio_frame = get_audio_frame(source_audio_path, temp_video_fps, frame_number - trim_frame_start)
|
||||
@@ -201,7 +205,66 @@ def process_stream_frame(frame_number : int, temp_video_resolution : Resolution)
|
||||
'temp_vision_mask': temp_vision_mask
|
||||
})
|
||||
|
||||
temp_vision_frame = conditional_merge_vision_mask(temp_vision_frame, temp_vision_mask)
|
||||
return conditional_merge_vision_mask(temp_vision_frame, temp_vision_mask)
|
||||
|
||||
|
||||
def resolve_disk_target_frames(frame_number : int, frame_amount : int, temp_frame_set : FrameSet) -> List[VisionFrame]:
|
||||
target_vision_frames = []
|
||||
|
||||
for current_number in range(frame_number - frame_amount, frame_number + frame_amount + 1):
|
||||
target_vision_frames.append(read_static_image(temp_frame_set.get(current_number)))
|
||||
|
||||
return target_vision_frames
|
||||
|
||||
|
||||
def process_disk_frame(temp_frame_path : str, frame_number : int, temp_frame_set : FrameSet) -> bool:
|
||||
target_vision_frames = resolve_disk_target_frames(frame_number, state_manager.get_item('target_frame_amount'), temp_frame_set)
|
||||
temp_vision_frame = read_static_image(temp_frame_path, 'rgba')
|
||||
temp_vision_frame = process_target_frame(frame_number, target_vision_frames, temp_vision_frame)
|
||||
return write_image(temp_frame_path, temp_vision_frame)
|
||||
|
||||
|
||||
def process_disk_frames() -> ErrorCode:
|
||||
temp_frame_set = resolve_temp_frame_set(state_manager.get_item('target_path'))
|
||||
|
||||
if temp_frame_set:
|
||||
with tqdm(total = len(temp_frame_set), desc = translator.get('processing'), unit = 'frame', ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress:
|
||||
progress.set_postfix(execution_providers = state_manager.get_item('execution_providers'))
|
||||
|
||||
with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor:
|
||||
futures = []
|
||||
|
||||
for frame_number, temp_frame_path in temp_frame_set.items():
|
||||
futures.append(executor.submit(process_disk_frame, temp_frame_path, frame_number, temp_frame_set))
|
||||
|
||||
for future in as_completed(futures):
|
||||
if is_process_stopping():
|
||||
for __future__ in futures:
|
||||
__future__.cancel()
|
||||
|
||||
if not future.cancelled():
|
||||
future.result()
|
||||
progress.update()
|
||||
|
||||
for processor_module in get_processors_modules(state_manager.get_item('processors')):
|
||||
processor_module.post_process()
|
||||
|
||||
if is_process_stopping():
|
||||
return 4
|
||||
else:
|
||||
logger.error(translator.get('temp_frames_not_found'), __name__)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
def process_stream_frame(frame_number : int, temp_video_resolution : Resolution) -> Tuple[int, VisionFrame]:
|
||||
target_vision_frames = select_video_frames(state_manager.get_item('target_path'), frame_number, state_manager.get_item('target_frame_amount'))
|
||||
target_vision_frame = get_middle(target_vision_frames)
|
||||
temp_vision_frame = target_vision_frame.copy()
|
||||
|
||||
if not (target_vision_frame.shape[1], target_vision_frame.shape[0]) == temp_video_resolution:
|
||||
temp_vision_frame = cv2.resize(target_vision_frame, temp_video_resolution)
|
||||
temp_vision_frame = process_target_frame(frame_number, target_vision_frames, temp_vision_frame)
|
||||
return frame_number, numpy.ascontiguousarray(temp_vision_frame[:, :, :3])
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user