diff --git a/facefusion/_preview.py b/facefusion/_preview.py index 873e1f5d..554f3d74 100644 --- a/facefusion/_preview.py +++ b/facefusion/_preview.py @@ -1,4 +1,5 @@ -from time import time +from typing import Optional + import cv2 import numpy @@ -15,7 +16,7 @@ from facefusion.typing import AudioFrame, Face, FaceSet, VisionFrame from facefusion.vision import get_video_frame, read_static_image, read_static_images, resize_frame_resolution -def process_frame(frame_number : int = 0) -> None: +def process_frame(frame_number : int = 0) -> Optional[VisionFrame]: core.conditional_append_reference_faces() reference_faces = get_reference_faces() if 'reference' in state_manager.get_item('face_selector_mode') else None source_frames = read_static_images(state_manager.get_item('source_paths')) @@ -48,6 +49,8 @@ def process_frame(frame_number : int = 0) -> None: preview_vision_frame = process_preview_frame(reference_faces, source_face, source_audio_frame, temp_vision_frame) return preview_vision_frame + return None + def process_preview_frame(reference_faces : FaceSet, source_face : Face, source_audio_frame : AudioFrame, target_vision_frame : VisionFrame) -> VisionFrame: target_vision_frame = resize_frame_resolution(target_vision_frame, (1024, 1024)) diff --git a/facefusion/api.py b/facefusion/api.py index 6e735788..1169fa5e 100644 --- a/facefusion/api.py +++ b/facefusion/api.py @@ -1,13 +1,13 @@ import asyncio -import time -from io import BytesIO +import json from typing import Any, List import cv2 import uvicorn -from litestar import Litestar, WebSocket, get as read, websocket as stream +from litestar import Litestar, WebSocket, get as read, websocket as stream, websocket_listener +from litestar.static_files import create_static_files_router -from facefusion import choices, execution, _preview +from facefusion import _preview, choices, execution, state_manager, vision from facefusion.processors import choices as processors_choices from facefusion.state_manager import get_state from facefusion.typing import ExecutionDevice @@ -47,7 +47,7 @@ async def read_execution_providers() -> Any: @stream('/execution/devices') -async def stream_execution_devices(socket : WebSocket) -> None: +async def stream_execution_devices(socket : WebSocket[Any, Any, Any]) -> None: await socket.accept() while True: @@ -66,7 +66,7 @@ async def read_static_execution_devices() -> List[ExecutionDevice]: @stream('/state') -async def stream_state(socket : WebSocket) -> None: +async def stream_state(socket : WebSocket[Any, Any, Any]) -> None: await socket.accept() while True: @@ -74,12 +74,30 @@ async def stream_state(socket : WebSocket) -> None: await asyncio.sleep(0.5) -@read('/preview', media_type = 'image/png') -async def read_preview() -> None: - _, preview_vision_frame = cv2.imencode('.png', _preview.process_frame()) +@read('/preview', media_type = 'image/png', mode = "binary") +async def read_preview(frame_number : int) -> bytes: + _, preview_vision_frame = cv2.imencode('.png', _preview.process_frame(frame_number)) #type:ignore return preview_vision_frame.tobytes() +@websocket_listener("/preview", send_mode = "binary") +async def stream_preview(data : str) -> bytes: + frame_number = int(json.loads(data).get('frame_number')) + _, preview_vision_frame = cv2.imencode('.png', _preview.process_frame(frame_number)) #type:ignore + return preview_vision_frame.tobytes() + + +@read('/ui/preview_slider') +async def read_ui_preview_slider() -> Any: + target_path = state_manager.get_item('target_path') + video_frame_total = vision.count_video_frame_total(target_path) + + return\ + { + 'video_frame_total': video_frame_total + } + + api = Litestar( [ read_choices, @@ -88,7 +106,14 @@ api = Litestar( read_execution_devices, read_static_execution_devices, stream_state, - read_preview + read_preview, + read_ui_preview_slider, + stream_preview, + create_static_files_router( + path = '/frontend', + directories = [ 'facefusion/static' ], + html_mode = True, + ) ]) diff --git a/facefusion/core.py b/facefusion/core.py index 48d9a60d..776ce5fd 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -2,6 +2,7 @@ import itertools import shutil import signal import sys +import webbrowser from time import time import numpy @@ -61,6 +62,9 @@ def route(args : Args) -> None: if not pre_check(): return conditional_exit(2) if state_manager.get_item('command') == 'run': + if state_manager.get_item('open_browser'): + webbrowser.open('http://127.0.0.1:8000/frontend') + logger.info('http://127.0.0.1:8000/frontend', __name__) api.run() if state_manager.get_item('command') == 'headless-run': if not job_manager.init_jobs(state_manager.get_item('jobs_path')): diff --git a/facefusion/static/index.html b/facefusion/static/index.html new file mode 100644 index 00000000..30ac2568 --- /dev/null +++ b/facefusion/static/index.html @@ -0,0 +1,242 @@ + + +
+ + +Frame: 0
+ + + Use WebSocket for Preview +Video Memory
GPU Utilization