diff --git a/facefusion/apis/assets.py b/facefusion/apis/assets.py new file mode 100644 index 00000000..926656b5 --- /dev/null +++ b/facefusion/apis/assets.py @@ -0,0 +1,186 @@ +import os +import tempfile + +from starlette.requests import Request +from starlette.responses import FileResponse, JSONResponse +from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_400_BAD_REQUEST, HTTP_404_NOT_FOUND + +from facefusion import asset_store, filesystem, logger +from facefusion.vision import count_video_frame_total, detect_video_fps, detect_video_resolution + + +async def upload_asset(request : Request) -> JSONResponse: + asset_type = request.query_params.get('type') + + if not asset_type: + return JSONResponse({'message': 'Missing required query parameter: type'}, status_code = HTTP_400_BAD_REQUEST) + + if asset_type not in ['source', 'target']: + return JSONResponse({'message': 'Invalid type. Must be "source" or "target"'}, status_code = HTTP_400_BAD_REQUEST) + + form = await request.form() + + if asset_type == 'source': + files = form.getlist('file') + + if not files: + return JSONResponse({'message': 'No file provided'}, status_code = HTTP_400_BAD_REQUEST) + + asset_ids = [] + + for file in files: + filename = file.filename if hasattr(file, 'filename') else 'source.jpg' + file_extension = os.path.splitext(filename)[1] if filename else '.jpg' + + with tempfile.NamedTemporaryFile(suffix=file_extension, delete=False) as temp_file: + content = await file.read() + temp_file.write(content) + file_path = temp_file.name + + if not (filesystem.is_image(file_path) or filesystem.is_video(file_path) or filesystem.is_audio(file_path)): + if os.path.exists(file_path): + os.remove(file_path) + return JSONResponse( + { + 'message': 'Unsupported file format. Allowed formats - Images: bmp, jpeg, png, tiff, webp. Videos: avi, m4v, mkv, mov, mp4, mpeg, mxf, webm, wmv.' + }, + status_code = HTTP_400_BAD_REQUEST + ) + + asset_id = asset_store.register('source', file_path, filename) + asset_ids.append(asset_id) + + logger.debug(f'Uploaded {len(asset_ids)} source(s)', __name__) + + return JSONResponse( + { + 'message': f'{len(asset_ids)} source(s) uploaded successfully', + 'asset_ids': asset_ids + }, + status_code = HTTP_201_CREATED + ) + + if asset_type == 'target': + file = form.get('file') + + if not file: + return JSONResponse({'message': 'No file provided'}, status_code = HTTP_400_BAD_REQUEST) + + if isinstance(file, str): + return JSONResponse({'message': 'Expected file upload, got string. Use /stream/target for URLs'}, status_code = HTTP_400_BAD_REQUEST) + + if not hasattr(file, 'filename'): + return JSONResponse({'message': 'Invalid file object'}, status_code = HTTP_400_BAD_REQUEST) + + filename = file.filename + file_extension = os.path.splitext(filename)[1] if filename else '.jpg' + + with tempfile.NamedTemporaryFile(suffix=file_extension, delete=False) as temp_file: + content = await file.read() + temp_file.write(content) + file_path = temp_file.name + + if not (filesystem.is_image(file_path) or filesystem.is_video(file_path) or filesystem.is_audio(file_path)): + if os.path.exists(file_path): + os.remove(file_path) + return JSONResponse( + { + 'message': 'Unsupported file format. Allowed formats - Images: bmp, jpeg, png, tiff, webp. Videos: avi, m4v, mkv, mov, mp4, mpeg, mxf, webm, wmv.' + }, + status_code = HTTP_400_BAD_REQUEST + ) + + metadata = None + + if filesystem.is_video(file_path): + frame_total = count_video_frame_total(file_path) + fps = detect_video_fps(file_path) + resolution = detect_video_resolution(file_path) + metadata =\ + { + 'frame_total': frame_total, + 'fps': fps, + 'resolution': resolution + } + logger.debug(f'Video metadata - frames: {frame_total}, fps: {fps}, resolution: {resolution}', __name__) + + asset_id = asset_store.register('target', file_path, filename, metadata) + + logger.debug(f'Target uploaded with asset_id: {asset_id}', __name__) + + return JSONResponse( + { + 'message': 'Target uploaded successfully', + 'asset_id': asset_id + }, + status_code = HTTP_201_CREATED + ) + + +async def list_all_assets(request : Request) -> JSONResponse: + asset_type = request.query_params.get('type') + media_type = request.query_params.get('media_type') + format = request.query_params.get('format') + + assets = asset_store.list_assets(asset_type) + + if media_type: + assets = [a for a in assets if a.get('media_type') == media_type] + + if format: + assets = [a for a in assets if a.get('format') == format] + + safe_assets = [] + for asset in assets: + safe_asset = {k: v for k, v in asset.items() if k != 'path'} + safe_assets.append(safe_asset) + + return JSONResponse({'assets': safe_assets, 'count': len(safe_assets)}, status_code = HTTP_200_OK) + + +async def get_asset_by_id(request : Request) -> JSONResponse | FileResponse: + from facefusion.session_context import get_session_id + + asset_id = request.path_params.get('asset_id') + action = request.query_params.get('action') + asset = asset_store.get_asset(asset_id) + + if not asset: + return JSONResponse({'message': 'Asset not found'}, status_code = HTTP_404_NOT_FOUND) + + if asset.get('session_id') != get_session_id(): + return JSONResponse({'message': 'Asset not found'}, status_code = HTTP_404_NOT_FOUND) + + if action == 'download': + file_path = asset.get('path') + + if not file_path or not os.path.exists(file_path): + return JSONResponse({'message': 'Asset file not found'}, status_code = HTTP_404_NOT_FOUND) + + filename = asset.get('filename', 'download') + + return FileResponse(file_path, filename = filename) + + safe_asset = {k: v for k, v in asset.items() if k != 'path'} + + return JSONResponse(safe_asset, status_code = HTTP_200_OK) + + +async def delete_asset_by_id(request : Request) -> JSONResponse: + from facefusion.session_context import get_session_id + + asset_id = request.path_params.get('asset_id') + asset = asset_store.get_asset(asset_id) + + if not asset: + return JSONResponse({'message': 'Asset not found'}, status_code = HTTP_404_NOT_FOUND) + + if asset.get('session_id') != get_session_id(): + return JSONResponse({'message': 'Asset not found'}, status_code = HTTP_404_NOT_FOUND) + + success = asset_store.delete_asset(asset_id) + + if not success: + return JSONResponse({'message': 'Asset not found'}, status_code = HTTP_404_NOT_FOUND) + + return JSONResponse({'message': 'Asset deleted successfully'}, status_code = HTTP_200_OK) diff --git a/facefusion/apis/choices.py b/facefusion/apis/choices.py new file mode 100644 index 00000000..1d1971d8 --- /dev/null +++ b/facefusion/apis/choices.py @@ -0,0 +1,162 @@ +from typing import Any, Dict + +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.status import HTTP_200_OK + +import facefusion.choices +from facefusion.execution import get_available_execution_providers +from facefusion.ffmpeg import get_available_encoder_set +from facefusion.processors.modules.face_debugger import choices as face_debugger_choices +from facefusion.processors.modules.face_enhancer import choices as face_enhancer_choices +from facefusion.processors.modules.face_swapper import choices as face_swapper_choices +from facefusion.processors.modules.frame_enhancer import choices as frame_enhancer_choices + + +async def get_choices(request : Request) -> JSONResponse: + available_execution_providers = get_available_execution_providers() + available_encoder_set = get_available_encoder_set() + + choices_data : Dict[str, Any] =\ + { + 'face_detector_models': facefusion.choices.face_detector_models, + 'face_detector_set': facefusion.choices.face_detector_set, + 'face_landmarker_models': facefusion.choices.face_landmarker_models, + 'face_selector_modes': facefusion.choices.face_selector_modes, + 'face_selector_orders': facefusion.choices.face_selector_orders, + 'face_selector_genders': facefusion.choices.face_selector_genders, + 'face_selector_races': facefusion.choices.face_selector_races, + 'face_occluder_models': facefusion.choices.face_occluder_models, + 'face_parser_models': facefusion.choices.face_parser_models, + 'face_mask_types': facefusion.choices.face_mask_types, + 'face_mask_areas': facefusion.choices.face_mask_areas, + 'face_mask_regions': facefusion.choices.face_mask_regions, + 'voice_extractor_models': facefusion.choices.voice_extractor_models, + 'workflows': facefusion.choices.workflows, + 'audio_formats': facefusion.choices.audio_formats, + 'image_formats': facefusion.choices.image_formats, + 'video_formats': facefusion.choices.video_formats, + 'temp_frame_formats': facefusion.choices.temp_frame_formats, + 'output_audio_encoders': available_encoder_set.get('audio'), + 'output_video_encoders': available_encoder_set.get('video'), + 'output_video_presets': facefusion.choices.output_video_presets, + 'execution_providers': available_execution_providers, + 'video_memory_strategies': facefusion.choices.video_memory_strategies, + 'log_levels': facefusion.choices.log_levels, + 'face_swapper_models': face_swapper_choices.face_swapper_models, + 'face_swapper_set': face_swapper_choices.face_swapper_set, + 'face_enhancer_models': face_enhancer_choices.face_enhancer_models, + 'frame_enhancer_models': frame_enhancer_choices.frame_enhancer_models, + 'face_debugger_items': face_debugger_choices.face_debugger_items, + 'face_detector_angles': list(facefusion.choices.face_detector_angles), + 'face_detector_score_range': + { + 'min': min(facefusion.choices.face_detector_score_range), + 'max': max(facefusion.choices.face_detector_score_range), + 'step': facefusion.choices.face_detector_score_range[1] - facefusion.choices.face_detector_score_range[0] + }, + 'face_landmarker_score_range': + { + 'min': min(facefusion.choices.face_landmarker_score_range), + 'max': max(facefusion.choices.face_landmarker_score_range), + 'step': facefusion.choices.face_landmarker_score_range[1] - facefusion.choices.face_landmarker_score_range[0] + }, + 'face_mask_blur_range': + { + 'min': min(facefusion.choices.face_mask_blur_range), + 'max': max(facefusion.choices.face_mask_blur_range), + 'step': facefusion.choices.face_mask_blur_range[1] - facefusion.choices.face_mask_blur_range[0] + }, + 'face_mask_padding_range': + { + 'min': min(facefusion.choices.face_mask_padding_range), + 'max': max(facefusion.choices.face_mask_padding_range), + 'step': 1 + }, + 'face_selector_age_range': + { + 'min': min(facefusion.choices.face_selector_age_range), + 'max': max(facefusion.choices.face_selector_age_range), + 'step': 1 + }, + 'reference_face_distance_range': + { + 'min': min(facefusion.choices.reference_face_distance_range), + 'max': max(facefusion.choices.reference_face_distance_range), + 'step': facefusion.choices.reference_face_distance_range[1] - facefusion.choices.reference_face_distance_range[0] + }, + 'output_image_quality_range': + { + 'min': min(facefusion.choices.output_image_quality_range), + 'max': max(facefusion.choices.output_image_quality_range), + 'step': 1 + }, + 'output_image_scale_range': + { + 'min': min(facefusion.choices.output_image_scale_range), + 'max': max(facefusion.choices.output_image_scale_range), + 'step': facefusion.choices.output_image_scale_range[1] - facefusion.choices.output_image_scale_range[0] + }, + 'output_audio_quality_range': + { + 'min': min(facefusion.choices.output_audio_quality_range), + 'max': max(facefusion.choices.output_audio_quality_range), + 'step': 1 + }, + 'output_audio_volume_range': + { + 'min': min(facefusion.choices.output_audio_volume_range), + 'max': max(facefusion.choices.output_audio_volume_range), + 'step': 1 + }, + 'output_video_quality_range': + { + 'min': min(facefusion.choices.output_video_quality_range), + 'max': max(facefusion.choices.output_video_quality_range), + 'step': 1 + }, + 'output_video_scale_range': + { + 'min': min(facefusion.choices.output_video_scale_range), + 'max': max(facefusion.choices.output_video_scale_range), + 'step': facefusion.choices.output_video_scale_range[1] - facefusion.choices.output_video_scale_range[0] + }, + 'execution_thread_count_range': + { + 'min': min(facefusion.choices.execution_thread_count_range), + 'max': max(facefusion.choices.execution_thread_count_range), + 'step': 1 + }, + 'face_detector_margin_range': + { + 'min': min(facefusion.choices.face_detector_margin_range), + 'max': max(facefusion.choices.face_detector_margin_range), + 'step': 1 + }, + 'face_swapper_weight_range': + { + 'min': min(face_swapper_choices.face_swapper_weight_range), + 'max': max(face_swapper_choices.face_swapper_weight_range), + 'step': face_swapper_choices.face_swapper_weight_range[1] - face_swapper_choices.face_swapper_weight_range[0] + }, + 'face_enhancer_blend_range': + { + 'min': min(face_enhancer_choices.face_enhancer_blend_range), + 'max': max(face_enhancer_choices.face_enhancer_blend_range), + 'step': 1 + }, + 'face_enhancer_weight_range': + { + 'min': min(face_enhancer_choices.face_enhancer_weight_range), + 'max': max(face_enhancer_choices.face_enhancer_weight_range), + 'step': face_enhancer_choices.face_enhancer_weight_range[1] - face_enhancer_choices.face_enhancer_weight_range[0] + }, + 'frame_enhancer_blend_range': + { + 'min': min(frame_enhancer_choices.frame_enhancer_blend_range), + 'max': max(frame_enhancer_choices.frame_enhancer_blend_range), + 'step': 1 + } + } + + return JSONResponse(choices_data, status_code = HTTP_200_OK) diff --git a/facefusion/apis/core.py b/facefusion/apis/core.py index 10541ba9..37407268 100644 --- a/facefusion/apis/core.py +++ b/facefusion/apis/core.py @@ -3,23 +3,42 @@ from starlette.middleware import Middleware from starlette.middleware.cors import CORSMiddleware from starlette.routing import Route, WebSocketRoute +from facefusion.apis.assets import delete_asset_by_id, get_asset_by_id, list_all_assets, upload_asset +from facefusion.apis.choices import get_choices +from facefusion.apis.metrics import websocket_metrics from facefusion.apis.ping import websocket_ping +from facefusion.apis.process import webrtc_offer, webrtc_stream_offer, websocket_process +from facefusion.apis.remote import remote from facefusion.apis.session import create_session, create_session_guard, destroy_session, get_session, refresh_session from facefusion.apis.state import get_state, set_state +from facefusion.apis.timeline import get_timeline +from facefusion.apis.version import create_version_guard def create_api() -> Starlette: + version_guard = Middleware(create_version_guard) session_guard = Middleware(create_session_guard) routes =\ - [ - Route('/session', create_session, methods = [ 'POST' ]), - Route('/session', get_session, methods = [ 'GET' ], middleware = [ session_guard ]), - Route('/session', refresh_session, methods = [ 'PUT' ]), - Route('/session', destroy_session, methods = [ 'DELETE' ], middleware = [ session_guard ]), - Route('/state', get_state, methods = [ 'GET' ], middleware = [ session_guard ]), - Route('/state', set_state, methods = [ 'PUT' ], middleware = [ session_guard ]), - WebSocketRoute('/ping', websocket_ping, middleware = [ session_guard ]) - ] + [ + Route('/session', create_session, methods = [ 'POST' ], middleware = [ version_guard ]), + Route('/session', get_session, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), + Route('/session', refresh_session, methods = [ 'PUT' ], middleware = [ version_guard ]), + Route('/session', destroy_session, methods = [ 'DELETE' ], middleware = [ version_guard, session_guard ]), + Route('/state', get_state, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), + Route('/state', set_state, methods = [ 'PUT' ], middleware = [ version_guard, session_guard ]), + Route('/assets', upload_asset, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), + Route('/assets', list_all_assets, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), + Route('/assets/{asset_id}', get_asset_by_id, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), + Route('/assets/{asset_id}', delete_asset_by_id, methods = [ 'DELETE' ], middleware = [ version_guard, session_guard ]), + Route('/choices', get_choices, methods=['GET'], middleware=[ version_guard, session_guard ]), + Route('/remote', remote, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), + Route('/timeline/{count:int}', get_timeline, methods = [ 'GET' ], middleware = [ version_guard, session_guard ]), + Route('/webrtc/offer', webrtc_offer, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), + Route('/stream/webrtc/offer', webrtc_stream_offer, methods = [ 'POST' ], middleware = [ version_guard, session_guard ]), + WebSocketRoute('/metrics', websocket_metrics, middleware = [ version_guard, session_guard ]), + WebSocketRoute('/ping', websocket_ping, middleware = [ version_guard, session_guard ]), + WebSocketRoute('/process', websocket_process, middleware = [ version_guard, session_guard ]) + ] api = Starlette(routes = routes) api.add_middleware(CORSMiddleware, allow_origins = [ '*' ], allow_methods = [ '*' ], allow_headers = [ '*' ]) diff --git a/facefusion/apis/locales.py b/facefusion/apis/locals.py similarity index 75% rename from facefusion/apis/locales.py rename to facefusion/apis/locals.py index 884e99b5..97be9e00 100644 --- a/facefusion/apis/locales.py +++ b/facefusion/apis/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/apis/metrics.py b/facefusion/apis/metrics.py new file mode 100644 index 00000000..abfa3442 --- /dev/null +++ b/facefusion/apis/metrics.py @@ -0,0 +1,76 @@ +import asyncio +from functools import lru_cache +from typing import Any, Dict, Optional, cast + +from starlette.datastructures import Headers +from starlette.websockets import WebSocket, WebSocketDisconnect + +from facefusion import state_manager +from facefusion.execution import detect_execution_devices +from facefusion.system import get_cpu_info, get_disk_info, get_load_average, get_network_info +from facefusion.system import get_operating_system_info, get_python_info, get_ram_info, get_temperature_info +from facefusion.types import SystemInfo + + +@lru_cache(maxsize = 1) +def get_cached_static_system_info() -> Dict[str, Any]: + return\ + { + 'operating_system': get_operating_system_info(), + 'python': get_python_info() + } + + +@lru_cache(maxsize = 1) +def get_cached_semi_static_system_info(temp_path : Optional[str]) -> Dict[str, Any]: + return\ + { + 'disk': get_disk_info(temp_path), + 'network': get_network_info() + } + + +def get_optimized_system_info(temp_path : Optional[str] = None) -> SystemInfo: + static_data = get_cached_static_system_info() + semi_static_data = get_cached_semi_static_system_info(temp_path) + dynamic_data : Dict[str, Any] =\ + { + 'cpu': get_cpu_info(), + 'ram': get_ram_info(), + 'temperatures': get_temperature_info(), + 'load_average': get_load_average() + } + + return cast(SystemInfo, {**static_data, **semi_static_data, **dynamic_data}) + + +async def websocket_metrics(websocket : WebSocket) -> None: + subprotocol = get_requested_subprotocol(websocket) + await websocket.accept(subprotocol = subprotocol) + + try: + while True: + temp_path = state_manager.get_temp_path() + execution_devices = detect_execution_devices() + system_info = get_optimized_system_info(temp_path) + metrics =\ + { + 'devices': execution_devices, + 'system': system_info + } + await websocket.send_json(metrics) + await asyncio.sleep(2) + + except (WebSocketDisconnect, Exception): + pass + + +def get_requested_subprotocol(websocket : WebSocket) -> Optional[str]: + headers = Headers(scope = websocket.scope) + protocol_header = headers.get('Sec-WebSocket-Protocol') + + if protocol_header: + protocol, _, _ = protocol_header.partition(',') + return protocol.strip() + + return None diff --git a/facefusion/apis/process.py b/facefusion/apis/process.py new file mode 100644 index 00000000..11a2645f --- /dev/null +++ b/facefusion/apis/process.py @@ -0,0 +1,830 @@ +import asyncio +import fractions +import subprocess +from functools import partial +from typing import Any, List, Optional, Set, TypeAlias + +import cv2 +import numpy +from aiortc import AudioStreamTrack, RTCPeerConnection, RTCSessionDescription, VideoStreamTrack +from aiortc.codecs import h264 +from av import AudioFrame, VideoFrame +from starlette.datastructures import Headers +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.websockets import WebSocket + +from facefusion import config, content_analyser, logger, state_manager +from facefusion.streamer import process_stream_frame +from facefusion.vision import obscure_frame + +PeerConnectionSet : TypeAlias = Set[RTCPeerConnection] +ResolutionTuple : TypeAlias = tuple[int, int] + +NSFW_LOCK = False + +pcs : PeerConnectionSet = set() + +RESOLUTION_MAP : dict[str, ResolutionTuple] =\ +{ + '480p': (640, 480), + '720p': (1280, 720), + '1080p': (1920, 1080) +} + + +def init_default_state() -> None: + if state_manager.get_item('execution_providers') is None: + state_manager.set_item('execution_providers', config.get_str_list('execution', 'execution_providers', 'cuda')) + if state_manager.get_item('execution_thread_count') is None: + state_manager.set_item('execution_thread_count', config.get_int_value('execution', 'execution_thread_count', '32')) + if state_manager.get_item('face_detector_model') is None: + state_manager.set_item('face_detector_model', config.get_str_value('face_detector', 'face_detector_model', 'yolo_face')) + if state_manager.get_item('face_detector_size') is None: + state_manager.set_item('face_detector_size', config.get_str_value('face_detector', 'face_detector_size', '640x640')) + if state_manager.get_item('face_detector_margin') is None: + state_manager.set_item('face_detector_margin', config.get_int_list('face_detector', 'face_detector_margin', '0 0 0 0')) + if state_manager.get_item('face_detector_angles') is None: + state_manager.set_item('face_detector_angles', config.get_int_list('face_detector', 'face_detector_angles', '0')) + if state_manager.get_item('face_detector_score') is None: + state_manager.set_item('face_detector_score', config.get_float_value('face_detector', 'face_detector_score', '0.5')) + if state_manager.get_item('face_landmarker_model') is None: + state_manager.set_item('face_landmarker_model', config.get_str_value('face_landmarker', 'face_landmarker_model', '2dfan4')) + if state_manager.get_item('face_landmarker_score') is None: + state_manager.set_item('face_landmarker_score', config.get_float_value('face_landmarker', 'face_landmarker_score', '0.5')) + if state_manager.get_item('face_selector_mode') is None: + state_manager.set_item('face_selector_mode', config.get_str_value('face_selector', 'face_selector_mode', 'many')) + if state_manager.get_item('face_selector_order') is None: + state_manager.set_item('face_selector_order', config.get_str_value('face_selector', 'face_selector_order', 'large-small')) + if state_manager.get_item('face_mask_types') is None: + state_manager.set_item('face_mask_types', config.get_str_list('face_masker', 'face_mask_types', 'occlusion')) + if state_manager.get_item('face_mask_blur') is None: + state_manager.set_item('face_mask_blur', config.get_float_value('face_masker', 'face_mask_blur', '0.3')) + if state_manager.get_item('face_mask_padding') is None: + state_manager.set_item('face_mask_padding', config.get_int_list('face_masker', 'face_mask_padding', '0 0 0 0')) + if state_manager.get_item('face_swapper_model') is None: + state_manager.set_item('face_swapper_model', config.get_str_value('processors', 'face_swapper_model', 'hyperswap_1a_256')) + if state_manager.get_item('face_swapper_pixel_boost') is None: + state_manager.set_item('face_swapper_pixel_boost', config.get_str_value('processors', 'face_swapper_pixel_boost', '256x256')) + if state_manager.get_item('face_swapper_weight') is None: + state_manager.set_item('face_swapper_weight', config.get_float_value('processors', 'face_swapper_weight', '0.5')) + if state_manager.get_item('face_enhancer_model') is None: + state_manager.set_item('face_enhancer_model', config.get_str_value('processors', 'face_enhancer_model', 'gfpgan_1.4')) + if state_manager.get_item('face_enhancer_blend') is None: + state_manager.set_item('face_enhancer_blend', config.get_int_value('processors', 'face_enhancer_blend', '80')) + if state_manager.get_item('frame_enhancer_model') is None: + state_manager.set_item('frame_enhancer_model', config.get_str_value('processors', 'frame_enhancer_model', 'real_esrgan_x2')) + if state_manager.get_item('frame_enhancer_blend') is None: + state_manager.set_item('frame_enhancer_blend', config.get_int_value('processors', 'frame_enhancer_blend', '80')) + if state_manager.get_item('face_debugger_items') is None: + state_manager.set_item('face_debugger_items', config.get_str_list('processors', 'face_debugger_items', 'kps')) + logger.debug(f'Initialized state - execution_providers: {state_manager.get_item("execution_providers")}', __name__) + + +def setup_bitrate_config(bitrate : int, encoder : str, mode_prefix : str = 'WebRTC') -> tuple[int, bool]: + if bitrate == 0: + bitrate_bps = 100000 + h264.DEFAULT_BITRATE = bitrate_bps + h264.MIN_BITRATE = 100000 + h264.MAX_BITRATE = 2000000 + logger.info( + f'{mode_prefix} setup: mode=auto, encoder={encoder}, ' + f'DEF={h264.DEFAULT_BITRATE / 1000} kbps, ' + f'MIN={h264.MIN_BITRATE / 1000} kbps, MAX={h264.MAX_BITRATE / 1000} kbps', + __name__ + ) + adaptive_bitrate = True + else: + bitrate_bps = bitrate * 1000 + h264.DEFAULT_BITRATE = bitrate_bps + h264.MIN_BITRATE = max(500000, bitrate_bps // 2) + h264.MAX_BITRATE = max(bitrate_bps * 2, 3000000) + logger.info( + f'{mode_prefix} setup: mode=manual, encoder={encoder}, ' + f'DEF={h264.DEFAULT_BITRATE / 1000} kbps, ' + f'MIN={h264.MIN_BITRATE / 1000} kbps, MAX={h264.MAX_BITRATE / 1000} kbps', + __name__ + ) + adaptive_bitrate = False + + return bitrate_bps, adaptive_bitrate + + +def create_video_stream_track(pc : RTCPeerConnection, bitrate_bps : int, adaptive_bitrate : bool, buffer_size : int = 30, mode_prefix : str = 'WebRTC') -> tuple[VideoStreamTrack, Any]: + logger.info(f'Creating {mode_prefix} output queue with buffer size: {buffer_size}', __name__) + + processed_track = VideoStreamTrack() + processed_track.frame_queue = asyncio.Queue(maxsize = buffer_size) + processed_track.recv = partial(recv_from_queue, processed_track.frame_queue) + processed_track.data_channel = None + processed_track.ready_sent = False + processed_track._target_bitrate = bitrate_bps + processed_track._adaptive = adaptive_bitrate + processed_track._current_bitrate = bitrate_bps + + sender = pc.addTrack(processed_track) + + return processed_track, sender + + +async def monitor_and_set_bitrate(sender : Any, bitrate_bps : int, adaptive_bitrate : bool, processed_track : VideoStreamTrack, buffer_size : int) -> None: + encoder_obj = None + for attempt in range(30): + if hasattr(sender, '_RTCRtpSender__encoder') and sender._RTCRtpSender__encoder: + encoder_obj = sender._RTCRtpSender__encoder + encoder_type = type(encoder_obj).__name__ + if hasattr(encoder_obj, 'target_bitrate'): + old_bitrate = encoder_obj.target_bitrate + encoder_obj.target_bitrate = bitrate_bps + logger.info( + f'Encoder type: {encoder_type}, updated bitrate from {old_bitrate / 1000} kbps to {bitrate_bps / 1000} kbps', + __name__ + ) + if hasattr(encoder_obj, 'codec') and encoder_obj.codec: + logger.info( + f'Encoder codec context: {encoder_obj.codec.name if hasattr(encoder_obj.codec, "name") else "unknown"}', + __name__ + ) + break + + if not encoder_obj: + logger.warn('Encoder not created after 9 seconds', __name__) + return + + if not adaptive_bitrate: + return + + stable_checks = 0 + INCREASE_STEP = 50000 + DECREASE_FACTOR = 0.9 + + while True: + await asyncio.sleep(0.5) + + if not hasattr(encoder_obj, 'target_bitrate'): + break + + queue_ratio = processed_track.frame_queue.qsize() / buffer_size + + if queue_ratio > 0.7: + new_bitrate = max(int(processed_track._current_bitrate * DECREASE_FACTOR), h264.MIN_BITRATE) + processed_track._current_bitrate = new_bitrate + encoder_obj.target_bitrate = new_bitrate + stable_checks = 0 + logger.info(f'Auto: decreased to {new_bitrate / 1000} kbps (congestion)', __name__) + elif queue_ratio < 0.3: + stable_checks += 1 + if stable_checks >= 4: + new_bitrate = min(processed_track._current_bitrate + INCREASE_STEP, h264.MAX_BITRATE) + if new_bitrate > processed_track._current_bitrate: + processed_track._current_bitrate = new_bitrate + encoder_obj.target_bitrate = new_bitrate + logger.info(f'Auto: increased to {new_bitrate / 1000} kbps (stable)', __name__) + stable_checks = 0 + else: + stable_checks = 0 + + +async def websocket_process(websocket : WebSocket) -> None: + subprotocol = get_requested_subprotocol(websocket) + await websocket.accept(subprotocol = subprotocol) + init_default_state() + + output_resolution = websocket.query_params.get('output_resolution', 'original') + + while True: + message = await websocket.receive() + + if message['type'] == 'websocket.disconnect': + logger.debug('Client disconnected', __name__) + break + + if message['type'] == 'websocket.receive': + if 'bytes' in message: + logger.debug(f'Received {len(message["bytes"])} bytes', __name__) + + target_vision_frame = cv2.imdecode(numpy.frombuffer(message['bytes'], numpy.uint8), cv2.IMREAD_COLOR) + if target_vision_frame is None: + logger.error('Failed to decode target image!', __name__) + continue + + logger.debug(f'Decoded target frame shape: {target_vision_frame.shape}', __name__) + + if output_resolution and output_resolution != 'original': + resolution_map =\ + { + '480p': (640, 480), + '720p': (1280, 720), + '1080p': (1920, 1080) + } + if output_resolution in resolution_map: + target_width, target_height = resolution_map[output_resolution] + current_height, current_width = target_vision_frame.shape[:2] + if current_width > target_width or current_height > target_height: + scale = min(target_width / current_width, target_height / current_height) + new_width = int(current_width * scale) + new_height = int(current_height * scale) + target_vision_frame = cv2.resize(target_vision_frame, (new_width, new_height), interpolation = cv2.INTER_AREA) + logger.debug(f'Downscaled target from {current_width}x{current_height} to {new_width}x{new_height}', __name__) + + temp_vision_frame = process_stream_frame(target_vision_frame) + if temp_vision_frame is None: + continue + + if content_analyser.analyse_frame(temp_vision_frame): + logger.warn('NSFW content detected in output, blurring frame', __name__) + temp_vision_frame = obscure_frame(temp_vision_frame) + + success, result_bytes = cv2.imencode('.jpg', temp_vision_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 50]) + if success: + await websocket.send_bytes(result_bytes.tobytes()) + + +async def process_incoming_video_track(track : Any, frame_queue : Any, output_track : Any = None, output_resolution : str = 'original') -> None: + from aiortc.mediastreams import MediaStreamError + + logger.debug(f'Track received: {track.kind}', __name__) + max_fps = 60 + min_frame_time = 1.0 / max_fps + last_process_time = 0.0 + frame_counter = 0 + frame_skip = 2 + last_processed_frame = None + + try: + global NSFW_LOCK + while True: + try: + frame = await track.recv() + except MediaStreamError: + logger.info('Media stream ended (connection closed)', __name__) + break + except asyncio.CancelledError: + logger.info('Video processing cancelled', __name__) + raise + except Exception as e: + logger.error(f'Error receiving frame: {e}', __name__) + break + + current_time = asyncio.get_event_loop().time() + time_since_last = current_time - last_process_time + + if time_since_last < min_frame_time: + continue + + img = frame.to_ndarray(format='bgr24') + logger.debug(f'Received frame shape: {img.shape}', __name__) + + if output_resolution and output_resolution != 'original': + if output_resolution in RESOLUTION_MAP: + target_width, target_height = RESOLUTION_MAP[output_resolution] + current_height, current_width = img.shape[:2] + if current_width > target_width or current_height > target_height: + scale = min(target_width / current_width, target_height / current_height) + new_width = int(current_width * scale) + new_height = int(current_height * scale) + img = cv2.resize(img, (new_width, new_height), interpolation = cv2.INTER_AREA) + logger.debug(f'Downscaled target from {current_width}x{current_height} to {new_width}x{new_height}', __name__) + + if content_analyser.analyse_stream(img, float(max_fps)): + NSFW_LOCK = True + + if NSFW_LOCK: + temp_vision_frame = obscure_frame(img) + else: + frame_counter += 1 + if frame_counter % frame_skip == 0: + temp_vision_frame = process_stream_frame(img) + last_processed_frame = temp_vision_frame + else: + if last_processed_frame is not None: + temp_vision_frame = last_processed_frame + else: + temp_vision_frame = process_stream_frame(img) + last_processed_frame = temp_vision_frame + + if temp_vision_frame is not None: + new_frame = VideoFrame.from_ndarray(temp_vision_frame, format='bgr24') + new_frame.pts = frame.pts + new_frame.time_base = frame.time_base + + if frame_queue.full(): + try: + frame_queue.get_nowait() + except asyncio.QueueEmpty: + pass + + try: + frame_queue.put_nowait(new_frame) + last_process_time = current_time + + if output_track and not output_track.ready_sent and frame_queue.qsize() >= int(frame_queue.maxsize * 0.5): + output_track.ready_sent = True + logger.info(f'Buffer ready ({frame_queue.qsize()}/{frame_queue.maxsize} frames), sending ready signal', __name__) + if output_track.data_channel and output_track.data_channel.readyState == 'open': + output_track.data_channel.send('ready') + logger.info('Ready signal sent to client', __name__) + + if output_track and output_track.data_channel and output_track.data_channel.readyState == 'open': + try: + output_track.data_channel.send(f'frame:{frame_counter}') + except Exception: + pass + except asyncio.QueueFull: + logger.debug('Frame queue full, frame dropped', __name__) + except Exception as e: + logger.error(f'Unexpected error in video processing: {e}', __name__) + finally: + logger.info('Video processing task completed', __name__) + + +async def recv_from_queue(frame_queue : Any) -> Any: + frame = await frame_queue.get() + return frame + + +def check_and_lock_nsfw(vision_frame : Any, fps : float) -> bool: + global NSFW_LOCK + + if NSFW_LOCK: + return True + + if content_analyser.analyse_stream(vision_frame, fps): + NSFW_LOCK = True + logger.warn('NSFW content detected, locking all future frames', __name__) + return True + + return False + + +def get_requested_subprotocol(websocket : WebSocket) -> Optional[str]: + headers = Headers(scope = websocket.scope) + protocol_header = headers.get('Sec-WebSocket-Protocol') + + if protocol_header: + protocol, _, _ = protocol_header.partition(',') + return protocol.strip() + + return None + + +async def webrtc_offer(request : Request) -> JSONResponse: + global NSFW_LOCK + init_default_state() + NSFW_LOCK = False + + params = await request.json() + offer = RTCSessionDescription(sdp=params['sdp'], type=params['type']) + bitrate = int(params.get('bitrate', 0)) + encoder = params.get('encoder', 'VP8') + buffer_size = int(params.get('stream_buffer_size', 30)) + output_resolution = params.get('output_resolution', 'original') + + bitrate_bps, adaptive_bitrate = setup_bitrate_config(bitrate, encoder, 'WebRTC') + + pc = RTCPeerConnection() + pcs.add(pc) + + processed_track, sender = create_video_stream_track(pc, bitrate_bps, adaptive_bitrate, buffer_size, 'WebRTC') + + asyncio.create_task(monitor_and_set_bitrate(sender, bitrate_bps, adaptive_bitrate, processed_track, buffer_size)) + + processing_tasks : List[Any] = [] + + @pc.on('connectionstatechange') + async def on_connectionstatechange() -> None: + logger.info(f'WebRTC connection state: {pc.connectionState}', __name__) + if pc.connectionState == 'failed' or pc.connectionState == 'closed': + logger.info('WebRTC connection closed, cleaning up', __name__) + pcs.discard(pc) + for task in processing_tasks: + task.cancel() + + @pc.on('datachannel') + def on_datachannel(channel : Any) -> None: + logger.info(f'Data channel received: {channel.label}', __name__) + processed_track.data_channel = channel + + await pc.setRemoteDescription(offer) + + for transceiver in pc.getTransceivers(): + if transceiver.receiver and transceiver.receiver.track: + track = transceiver.receiver.track + if track.kind == 'video': + logger.info('Found video track, starting processing', __name__) + video_task = asyncio.create_task(process_incoming_video_track(track, processed_track.frame_queue, processed_track, output_resolution)) + processing_tasks.append(video_task) + if track.kind == 'audio': + logger.info('Found audio track, forwarding as-is', __name__) + pc.addTrack(track) + + answer = await pc.createAnswer() + await pc.setLocalDescription(answer) + + return JSONResponse({'sdp': pc.localDescription.sdp, 'type': pc.localDescription.type}) + + +async def process_stream_from_url(stream_url : str, frame_queue : Any, output_track : Any = None, width : int = 1280, height : int = 720, target_fps : int = 30, duration : float = 0, output_resolution : str = 'original') -> None: + logger.info(f'Opening stream from URL: {stream_url[:100]}', __name__) + logger.info(f'Using metadata - Resolution: {width}x{height}, FPS: {target_fps}, Duration: {duration}s', __name__) + + frame_size = width * height * 3 + frame_interval = 1.0 / target_fps + + import threading + import time + + current_process = None + seek_position = [0.0] + stop_flag = [False] + lock = threading.Lock() + + def start_ffmpeg(start_time : float) -> Any: + ffmpeg_command =\ + [ + 'ffmpeg', + '-user_agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', + '-reconnect', '1', + '-reconnect_streamed', '1', + '-reconnect_delay_max', '5' + ] + + if start_time > 0: + ffmpeg_command.extend(['-ss', str(start_time)]) + + ffmpeg_command.extend([ + '-i', stream_url, + '-f', 'rawvideo', + '-pix_fmt', 'bgr24', + '-' + ]) + + try: + return subprocess.Popen(ffmpeg_command, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, bufsize = 10**8) + except Exception as e: + logger.error(f'Failed to start ffmpeg: {e}', __name__) + return None + + def read_stream() -> None: + nonlocal current_process + frame_count = int(seek_position[0] * target_fps) + + current_process = start_ffmpeg(seek_position[0]) + if not current_process: + return + + logger.info(f'FFmpeg stream started at position {seek_position[0]}s', __name__) + + last_frame_time = time.time() + local_process = current_process + + while not stop_flag[0]: + with lock: + if local_process != current_process: + logger.info('Process changed due to seek, switching to new process', __name__) + local_process = current_process + frame_count = int(seek_position[0] * target_fps) + if not local_process: + break + continue + + if local_process.poll() is not None: + logger.info('Stream process terminated', __name__) + break + + raw_frame = local_process.stdout.read(frame_size) + if not raw_frame or len(raw_frame) != frame_size: + with lock: + if local_process != current_process: + logger.info('Incomplete read due to seek, switching to new process', __name__) + local_process = current_process + frame_count = int(seek_position[0] * target_fps) + continue + logger.info('Stream ended or incomplete frame', __name__) + break + + frame = numpy.frombuffer(raw_frame, dtype = numpy.uint8).reshape((height, width, 3)) + + if output_resolution and output_resolution != 'original': + if output_resolution in RESOLUTION_MAP: + target_width, target_height = RESOLUTION_MAP[output_resolution] + current_height, current_width = frame.shape[:2] + if current_width > target_width or current_height > target_height: + scale = min(target_width / current_width, target_height / current_height) + new_width = int(current_width * scale) + new_height = int(current_height * scale) + frame = cv2.resize(frame, (new_width, new_height), interpolation = cv2.INTER_AREA) + + if check_and_lock_nsfw(frame, float(target_fps)): + processed_frame = obscure_frame(frame) + else: + processed_frame = process_stream_frame(frame) + + if processed_frame is not None: + current_time = time.time() + elapsed = current_time - last_frame_time + + if elapsed < frame_interval: + time.sleep(frame_interval - elapsed) + current_time = time.time() + + video_frame = VideoFrame.from_ndarray(processed_frame, format = 'bgr24') + video_frame.pts = frame_count + video_frame.time_base = fractions.Fraction(1, target_fps) + + if frame_queue.full(): + try: + frame_queue.get_nowait() + except asyncio.QueueEmpty: + pass + + try: + frame_queue.put_nowait(video_frame) + if output_track and not output_track.ready_sent and frame_queue.qsize() >= int(frame_queue.maxsize * 0.5): + output_track.ready_sent = True + logger.info(f'Buffer ready ({frame_queue.qsize()}/{frame_queue.maxsize} frames)', __name__) + if output_track.data_channel and output_track.data_channel.readyState == 'open': + output_track.data_channel.send('ready') + + if output_track and output_track.data_channel and output_track.data_channel.readyState == 'open': + try: + output_track.data_channel.send(f'frame:{frame_count}') + except Exception: + pass + except asyncio.QueueFull: + pass + + last_frame_time = current_time + frame_count += 1 + + if current_process: + current_process.terminate() + current_process.wait() + logger.info(f'Stream reading completed, {frame_count} frames processed', __name__) + + def handle_seek(new_position : float) -> None: + nonlocal current_process + with lock: + seek_position[0] = new_position + if current_process: + logger.info(f'Seeking to {new_position}s, restarting ffmpeg', __name__) + current_process.terminate() + current_process = start_ffmpeg(new_position) + + if output_track: + output_track.seek_handler = handle_seek + + stream_task = asyncio.get_event_loop().run_in_executor(None, read_stream) + + try: + await stream_task + except asyncio.CancelledError: + logger.info('Video stream task cancelled', __name__) + raise + finally: + stop_flag[0] = True + if current_process: + current_process.terminate() + logger.info('Video stream cleanup completed', __name__) + + +async def process_audio_from_url(stream_url : str, audio_queue : Any, video_output_track : Any = None) -> None: + logger.info('Opening audio stream from URL', __name__) + + import threading + + current_process = None + seek_position = [0.0] + stop_flag = [False] + lock = threading.Lock() + + sample_rate = 24000 + channels = 2 + frame_samples = 480 + bytes_per_sample = 2 + frame_size = frame_samples * channels * bytes_per_sample + + def start_ffmpeg_audio(start_time : float) -> Any: + ffmpeg_command =\ + [ + 'ffmpeg', + '-user_agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', + '-reconnect', '1', + '-reconnect_streamed', '1', + '-reconnect_delay_max', '5' + ] + + if start_time > 0: + ffmpeg_command.extend(['-ss', str(start_time)]) + + ffmpeg_command.extend([ + '-i', stream_url, + '-vn', + '-f', 's16le', + '-acodec', 'pcm_s16le', + '-ar', str(sample_rate), + '-ac', str(channels), + '-' + ]) + + try: + return subprocess.Popen(ffmpeg_command, stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, bufsize = 10**8) + except Exception as e: + logger.error(f'Failed to start ffmpeg audio: {e}', __name__) + return None + + def read_audio_stream() -> None: + nonlocal current_process + import time + + current_process = start_ffmpeg_audio(seek_position[0]) + if not current_process: + return + + logger.info(f'FFmpeg audio stream started at position {seek_position[0]}s', __name__) + + local_process = current_process + frame_count = 0 + + frame_duration = frame_samples / sample_rate + start_time = time.time() + expected_time = start_time + + logger.info(f'Starting audio read loop, expecting {frame_size} bytes per frame', __name__) + + while not stop_flag[0]: + with lock: + if local_process != current_process: + logger.info('Audio process changed due to seek, switching to new process', __name__) + local_process = current_process + frame_count = 0 + start_time = time.time() + expected_time = start_time + if not local_process: + break + continue + + if local_process.poll() is not None: + logger.info('Audio stream process terminated', __name__) + break + + raw_audio = local_process.stdout.read(frame_size) + + if frame_count == 0: + logger.info('Successfully read first audio frame', __name__) + if not raw_audio or len(raw_audio) != frame_size: + with lock: + if local_process != current_process: + logger.info('Incomplete audio read due to seek, switching to new process', __name__) + local_process = current_process + continue + logger.info('Audio stream ended or incomplete frame', __name__) + break + + audio_array = numpy.frombuffer(raw_audio, dtype = numpy.int16) + + audio_frame = AudioFrame(format = 's16', layout = 'stereo', samples = frame_samples) + audio_frame.sample_rate = sample_rate + audio_frame.pts = frame_count * frame_samples + audio_frame.time_base = fractions.Fraction(1, sample_rate) + + for plane in audio_frame.planes: + plane.update(audio_array.tobytes()) + + expected_time += frame_duration + current_time = time.time() + sleep_time = expected_time - current_time + + if sleep_time > 0: + time.sleep(sleep_time) + + if audio_queue.full(): + try: + audio_queue.get_nowait() + except asyncio.QueueEmpty: + pass + + try: + audio_queue.put_nowait(audio_frame) + frame_count += 1 + if frame_count % 50 == 0: + logger.info(f'Audio frames queued: {frame_count}, queue depth: {audio_queue.qsize()}', __name__) + except asyncio.QueueFull: + pass + + if current_process: + current_process.terminate() + logger.info('Audio stream reading completed', __name__) + + def handle_audio_seek(new_position : float) -> None: + nonlocal current_process + with lock: + seek_position[0] = new_position + if current_process: + logger.info(f'Seeking audio to {new_position}s, restarting ffmpeg', __name__) + current_process.terminate() + current_process = start_ffmpeg_audio(new_position) + + if video_output_track: + video_output_track.audio_seek_handler = handle_audio_seek + + audio_task = asyncio.get_event_loop().run_in_executor(None, read_audio_stream) + + try: + await audio_task + except asyncio.CancelledError: + logger.info('Audio stream task cancelled', __name__) + raise + finally: + stop_flag[0] = True + if current_process: + current_process.terminate() + logger.info('Audio stream cleanup completed', __name__) + + +async def webrtc_stream_offer(request : Request) -> JSONResponse: + global NSFW_LOCK + init_default_state() + NSFW_LOCK = False + + params = await request.json() + offer = RTCSessionDescription(sdp = params['sdp'], type = params['type']) + bitrate = int(params.get('bitrate', 0)) + encoder = params.get('encoder', 'VP8') + is_remote_stream = params.get('is_remote_stream', False) + stream_url = params.get('stream_url') + target_width = int(params.get('target_width', 1280)) + target_height = int(params.get('target_height', 720)) + target_fps = int(params.get('target_fps', 30)) + target_duration = float(params.get('target_duration', 0)) + target_audio_path = params.get('target_audio_path') + buffer_size = int(params.get('stream_buffer_size', 30)) + output_resolution = params.get('output_resolution', 'original') + + logger.info(f'[WebRTC Stream] stream_url: {stream_url[:50] if stream_url else None}', __name__) + logger.info(f'[WebRTC Stream] is_remote_stream: {is_remote_stream}', __name__) + + if not stream_url: + logger.error('[WebRTC Stream] No stream URL provided', __name__) + return JSONResponse({'error': 'No stream URL provided in request'}, status_code = 400) + + if not is_remote_stream: + logger.error('[WebRTC Stream] is_remote_stream=False, use /webrtc/offer for local files', __name__) + return JSONResponse({'error': 'Use /webrtc/offer endpoint for local files, not /stream/webrtc/offer'}, status_code = 400) + + bitrate_bps, adaptive_bitrate = setup_bitrate_config(bitrate, encoder, 'WebRTC stream') + + pc = RTCPeerConnection() + pcs.add(pc) + + processed_track, sender = create_video_stream_track(pc, bitrate_bps, adaptive_bitrate, buffer_size, 'WebRTC stream') + + audio_track = AudioStreamTrack() + audio_track.audio_queue = asyncio.Queue(maxsize = buffer_size) + audio_track.recv = partial(recv_from_queue, audio_track.audio_queue) + pc.addTrack(audio_track) + + asyncio.create_task(monitor_and_set_bitrate(sender, bitrate_bps, adaptive_bitrate, processed_track, buffer_size)) + + stream_tasks : List[Any] = [] + + @pc.on('connectionstatechange') + async def on_connectionstatechange() -> None: + logger.info(f'WebRTC stream connection state: {pc.connectionState}', __name__) + if pc.connectionState == 'failed' or pc.connectionState == 'closed': + logger.info('WebRTC stream connection closed, stopping ffmpeg processes', __name__) + pcs.discard(pc) + for task in stream_tasks: + task.cancel() + + @pc.on('datachannel') + def on_datachannel(channel : Any) -> None: + logger.info(f'Data channel received: {channel.label}', __name__) + processed_track.data_channel = channel + + @channel.on('message') + def on_message(message : Any) -> None: + if isinstance(message, str) and message.startswith('seek:'): + try: + seek_time = float(message.split(':', 1)[1]) + logger.info(f'Received seek command: {seek_time}s', __name__) + if hasattr(processed_track, 'seek_handler') and processed_track.seek_handler: + processed_track.seek_handler(seek_time) + if hasattr(processed_track, 'audio_seek_handler') and processed_track.audio_seek_handler: + processed_track.audio_seek_handler(seek_time) + except Exception as e: + logger.error(f'Error handling seek command: {e}', __name__) + + await pc.setRemoteDescription(offer) + + audio_url = target_audio_path or stream_url + logger.info(f'Starting stream processing from URL (video: {stream_url[:50]}..., audio: {audio_url[:50]}...)', __name__) + + video_task = asyncio.create_task(process_stream_from_url(stream_url, processed_track.frame_queue, processed_track, target_width, target_height, target_fps, target_duration, output_resolution)) + audio_task = asyncio.create_task(process_audio_from_url(audio_url, audio_track.audio_queue, processed_track)) + stream_tasks.extend([video_task, audio_task]) + + answer = await pc.createAnswer() + await pc.setLocalDescription(answer) + + return JSONResponse({'sdp': pc.localDescription.sdp, 'type': pc.localDescription.type}) diff --git a/facefusion/apis/remote.py b/facefusion/apis/remote.py new file mode 100644 index 00000000..e112b9c2 --- /dev/null +++ b/facefusion/apis/remote.py @@ -0,0 +1,383 @@ +import os +import tempfile +from typing import Any, Dict, List + +import httpx +import yt_dlp # type: ignore +from gallery_dl import config as gallery_config, extractor as gallery_extractor, job as gallery_job +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_400_BAD_REQUEST, HTTP_500_INTERNAL_SERVER_ERROR + +from facefusion import asset_store, logger +from facefusion.choices import audio_formats + + +def resolve_image_urls(url : str) -> List[str]: + gallery_config.load() + image_urls : List[str] = [] + + try: + for extractor_instance in gallery_extractor.extractors(): + if extractor_instance.pattern and extractor_instance.pattern.match(url): + logger.info(f'Detected gallery URL using extractor: {extractor_instance.__name__}', __name__) + extractor_obj = extractor_instance.from_url(url) + + if extractor_obj: + for msg in extractor_obj: + if isinstance(msg, tuple) and len(msg) >= 2: + msg_type = msg[0] + if msg_type == 5: + image_data = msg[1] + image_url = image_data.get('url') + if image_url: + image_urls.append(image_url) + break + + if not image_urls: + logger.info('Not a gallery URL, treating as direct image URL', __name__) + image_urls = [url] + + except Exception as e: + logger.error(f'Failed to extract image URLs: {e}', __name__) + logger.info('Falling back to treating as direct image URL', __name__) + image_urls = [url] + + return image_urls + + +def download_images_from_url(url : str, asset_type : str) -> List[str]: + gallery_config.load() + temp_dir = tempfile.gettempdir() + asset_ids : List[str] = [] + + is_gallery = False + for extractor_instance in gallery_extractor.extractors(): + if extractor_instance.pattern and extractor_instance.pattern.match(url): + logger.info(f'Detected gallery URL using extractor: {extractor_instance.__name__}', __name__) + is_gallery = True + + output_dir = os.path.join(temp_dir, f'facefusion_gallery_{os.urandom(8).hex()}') + os.makedirs(output_dir, exist_ok = True) + + gallery_config.set((), 'base-directory', output_dir) + gallery_config.set((), 'skip', False) + + gdl_job = gallery_job.DownloadJob(url) + gdl_job.run() + + for root, dirs, files in os.walk(output_dir): + for filename in files: + file_path = os.path.join(root, filename) + asset_id = asset_store.register(asset_type, file_path, filename) + asset_ids.append(asset_id) + logger.info(f'Registered image as asset {asset_id}', __name__) + + break + + if not is_gallery: + logger.info('Not a gallery URL, treating as direct image URL', __name__) + with httpx.stream('GET', url, timeout = 30, follow_redirects = True) as response: + response.raise_for_status() + + content_type = response.headers.get('content-type', '') + if not content_type.startswith('image/'): + raise ValueError(f'URL does not point to an image. Content-Type: {content_type}') + + file_extension = None + if 'image/jpeg' in content_type or 'image/jpg' in content_type: + file_extension = '.jpg' + if 'image/png' in content_type: + file_extension = '.png' + if 'image/gif' in content_type: + file_extension = '.gif' + if 'image/webp' in content_type: + file_extension = '.webp' + + if not file_extension: + url_path = url.split('?')[0] + if '.' in url_path: + file_extension = '.' + url_path.split('.')[-1].lower() + else: + file_extension = '.jpg' + + filename = f'facefusion_image_{os.urandom(8).hex()}{file_extension}' + file_path = os.path.join(temp_dir, filename) + + with open(file_path, 'wb') as f: + for chunk in response.iter_bytes(chunk_size = 8192): + f.write(chunk) + + asset_id = asset_store.register(asset_type, file_path, filename) + asset_ids.append(asset_id) + logger.info(f'Downloaded and registered image as asset {asset_id}', __name__) + + return asset_ids + + +def download_audio_from_url(url : str, asset_type : str) -> List[str]: + temp_dir = tempfile.gettempdir() + asset_ids : List[str] = [] + + # Extract file extension from URL + url_path = url.split('?')[0] + url_extension = os.path.splitext(url_path)[1].lstrip('.') + + # Validate extension against supported audio formats + if url_extension not in audio_formats: + raise ValueError(f'Unsupported audio format: {url_extension}. Supported formats: {", ".join(audio_formats)}') + + logger.info(f'Downloading audio from URL with extension: {url_extension}', __name__) + with httpx.stream('GET', url, timeout = 30, follow_redirects = True) as response: + response.raise_for_status() + + filename = f'facefusion_audio_{os.urandom(8).hex()}.{url_extension}' + file_path = os.path.join(temp_dir, filename) + + with open(file_path, 'wb') as f: + for chunk in response.iter_bytes(chunk_size = 8192): + f.write(chunk) + + asset_id = asset_store.register(asset_type, file_path, filename) + asset_ids.append(asset_id) + logger.info(f'Downloaded and registered audio as asset {asset_id}', __name__) + + return asset_ids + + +async def remote(request : Request) -> JSONResponse: + body = await request.json() + url = body.get('url') + action = request.query_params.get('action') + media_type = request.query_params.get('media_type', 'video') + asset_type = request.query_params.get('asset_type', 'target') + + if not action: + return JSONResponse({'message': 'No action provided. Must be "resolve" or "download"'}, status_code = HTTP_400_BAD_REQUEST) + + if action not in ['resolve', 'download']: + return JSONResponse({'message': 'Invalid action. Must be "resolve" or "download"'}, status_code = HTTP_400_BAD_REQUEST) + + if media_type not in ['image', 'video', 'audio']: + return JSONResponse({'message': 'Invalid media_type. Must be "image", "video", or "audio"'}, status_code = HTTP_400_BAD_REQUEST) + + if asset_type not in ['source', 'target']: + return JSONResponse({'message': 'Invalid asset_type. Must be "source" or "target"'}, status_code = HTTP_400_BAD_REQUEST) + + if not url: + return JSONResponse({'message': 'No URL provided'}, status_code = HTTP_400_BAD_REQUEST) + + if not isinstance(url, str): + return JSONResponse({'message': 'URL must be a string'}, status_code = HTTP_400_BAD_REQUEST) + + url = url.strip() + + if not url.startswith('http://') and not url.startswith('https://'): + return JSONResponse({'message': 'URL must start with http:// or https://'}, status_code = HTTP_400_BAD_REQUEST) + + quality = body.get('quality', '720p') + + if quality not in ['360p', '480p', '720p', '1080p']: + return JSONResponse({'message': 'Quality must be 360p, 480p, 720p, or 1080p'}, status_code = HTTP_400_BAD_REQUEST) + + if action == 'resolve': + if media_type == 'image': + image_urls = resolve_image_urls(url) + logger.info(f'Resolved {len(image_urls)} image URL(s)', __name__) + + response_data =\ + { + 'message': 'Image URL(s) resolved successfully', + 'image_urls': image_urls, + 'count': len(image_urls) + } + + return JSONResponse(response_data, status_code = HTTP_200_OK) + + quality_map =\ + { + '360p': 'bestvideo[height<=360][ext=mp4]+bestaudio[ext=m4a]/best[height<=360]', + '480p': 'bestvideo[height<=480][ext=mp4]+bestaudio[ext=m4a]/best[height<=480]', + '720p': 'bestvideo[height<=720][ext=mp4]+bestaudio[ext=m4a]/best[height<=720]', + '1080p': 'bestvideo[height<=1080][ext=mp4]+bestaudio[ext=m4a]/best[height<=1080]' + } + + ydl_opts : Dict[str, Any] =\ + { + 'format': quality_map[quality], + 'quiet': True, + 'no_warnings': True + } + + logger.info(f'Extracting stream URL from {url} at {quality}', __name__) + + try: + ydl = yt_dlp.YoutubeDL(ydl_opts) + info = ydl.extract_info(url, download = False) + except Exception as e: + logger.error(f'Failed to extract video information: {e}', __name__) + return JSONResponse({'message': f'Failed to extract video information: {str(e)}'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + if not info: + logger.error('Failed to extract video information', __name__) + return JSONResponse({'message': 'Failed to extract video information'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + stream_url = info.get('url') + + if not stream_url: + if 'requested_formats' in info and len(info['requested_formats']) > 0: + stream_url = info['requested_formats'][0].get('url') + logger.info('Using URL from requested_formats (video track)', __name__) + elif 'formats' in info and len(info['formats']) > 0: + for fmt in reversed(info['formats']): + if fmt.get('url') and fmt.get('vcodec') != 'none': + stream_url = fmt['url'] + logger.info(f'Using URL from format: {fmt.get("format_id")}', __name__) + break + + if not stream_url: + logger.error('No stream URL found in any format', __name__) + logger.debug(f'Available keys in info: {list(info.keys())}', __name__) + return JSONResponse({'message': 'No stream URL found'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + audio_url = None + if 'requested_formats' in info and len(info['requested_formats']) > 1: + audio_url = info['requested_formats'][1].get('url') + if audio_url: + logger.info('Found separate audio track URL', __name__) + + duration = info.get('duration') + fps = info.get('fps') + width = info.get('width') + height = info.get('height') + + total_frames = None + if duration and fps: + total_frames = int(duration * fps) + logger.info(f'Calculated total frames: {total_frames} ({duration}s * {fps} fps)', __name__) + + logger.info('Stream URL extracted successfully', __name__) + + response_data =\ + { + 'message': 'Stream URL resolved successfully', + 'stream_url': stream_url, + 'audio_url': audio_url, + 'duration': duration, + 'fps': fps, + 'total_frames': total_frames, + 'width': width, + 'height': height + } + + return JSONResponse(response_data, status_code = HTTP_200_OK) + + if action == 'download': + if media_type == 'image': + try: + asset_ids = download_images_from_url(url, asset_type) + except ValueError as e: + return JSONResponse({'message': str(e)}, status_code = HTTP_400_BAD_REQUEST) + except Exception as e: + logger.error(f'Failed to download image(s): {e}', __name__) + return JSONResponse({'message': f'Failed to download image(s): {str(e)}'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + response_data =\ + { + 'message': f'Downloaded and registered {len(asset_ids)} image(s)', + 'asset_ids': asset_ids, + 'count': len(asset_ids) + } + + return JSONResponse(response_data, status_code = HTTP_201_CREATED) + + if media_type == 'audio': + try: + asset_ids = download_audio_from_url(url, asset_type) + except ValueError as e: + return JSONResponse({'message': str(e)}, status_code = HTTP_400_BAD_REQUEST) + except Exception as e: + logger.error(f'Failed to download audio: {e}', __name__) + return JSONResponse({'message': f'Failed to download audio: {str(e)}'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + response_data =\ + { + 'message': f'Downloaded and registered {len(asset_ids)} audio file(s)', + 'asset_ids': asset_ids, + 'count': len(asset_ids) + } + + return JSONResponse(response_data, status_code = HTTP_201_CREATED) + + quality_map =\ + { + '360p': 'bestvideo[height<=360][ext=mp4]+bestaudio[ext=m4a]/best[height<=360][ext=mp4]/best[height<=360]', + '480p': 'bestvideo[height<=480][ext=mp4]+bestaudio[ext=m4a]/best[height<=480][ext=mp4]/best[height<=480]', + '720p': 'bestvideo[height<=720][ext=mp4]+bestaudio[ext=m4a]/best[height<=720][ext=mp4]/best[height<=720]', + '1080p': 'bestvideo[height<=1080][ext=mp4]+bestaudio[ext=m4a]/best[height<=1080][ext=mp4]/best[height<=1080]' + } + + temp_dir = tempfile.gettempdir() + output_path = os.path.join(temp_dir, 'facefusion_remote_%(id)s.%(ext)s') + + download_opts : Dict[str, Any] =\ + { + 'format': quality_map[quality], + 'outtmpl': output_path, + 'quiet': False, + 'no_warnings': False + } + + logger.info(f'Downloading video from {url} at {quality}', __name__) + + ydl = yt_dlp.YoutubeDL(download_opts) + info = ydl.extract_info(url, download = True) + + if not info: + logger.error('Failed to download video', __name__) + return JSONResponse({'message': 'Failed to download video'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + downloaded_file = ydl.prepare_filename(info) + if not os.path.exists(downloaded_file): + logger.error(f'Downloaded file not found: {downloaded_file}', __name__) + return JSONResponse({'message': 'Downloaded file not found'}, status_code = HTTP_500_INTERNAL_SERVER_ERROR) + + duration = info.get('duration') + fps = info.get('fps') + width = info.get('width') + height = info.get('height') + + total_frames = None + if duration and fps: + total_frames = int(duration * fps) + logger.info(f'Calculated total frames: {total_frames} ({duration}s * {fps} fps)', __name__) + + filename = os.path.basename(downloaded_file) + metadata =\ + { + 'frame_total': total_frames, + 'fps': fps, + 'resolution': (width, height) if width and height else None, + 'duration': duration + } + + asset_id = asset_store.register(asset_type, downloaded_file, filename, metadata) + logger.info(f'Video downloaded and registered as asset {asset_id}', __name__) + + response_data =\ + { + 'message': 'Video downloaded and registered as asset', + 'asset_id': asset_id, + 'metadata': + { + 'duration': duration, + 'fps': fps, + 'total_frames': total_frames, + 'width': width, + 'height': height + } + } + + return JSONResponse(response_data, status_code = HTTP_201_CREATED) + + return JSONResponse({'message': 'Invalid request'}, status_code = HTTP_400_BAD_REQUEST) diff --git a/facefusion/apis/session.py b/facefusion/apis/session.py index 59b24b38..2668e408 100644 --- a/facefusion/apis/session.py +++ b/facefusion/apis/session.py @@ -106,6 +106,8 @@ def create_session_guard(app : ASGIApp) -> ASGIApp: if session_id: if session_manager.validate_session(session_id): + from facefusion.session_context import set_session_id + set_session_id(session_id) return await app(scope, receive, send) response = JSONResponse( diff --git a/facefusion/apis/state.py b/facefusion/apis/state.py index c1090581..c440b901 100644 --- a/facefusion/apis/state.py +++ b/facefusion/apis/state.py @@ -1,8 +1,8 @@ from starlette.requests import Request from starlette.responses import JSONResponse -from starlette.status import HTTP_200_OK +from starlette.status import HTTP_200_OK, HTTP_400_BAD_REQUEST, HTTP_404_NOT_FOUND -from facefusion import args_store, state_manager +from facefusion import args_store, asset_store, logger, state_manager async def get_state(request : Request) -> JSONResponse: @@ -12,11 +12,65 @@ async def get_state(request : Request) -> JSONResponse: async def set_state(request : Request) -> JSONResponse: body = await request.json() + action = request.query_params.get('action') + + if action == 'select': + asset_type = request.query_params.get('asset_type') + + if not asset_type: + return JSONResponse({'message': 'Missing required query parameter: asset_type'}, status_code = HTTP_400_BAD_REQUEST) + + if asset_type not in ['source', 'target']: + return JSONResponse({'message': 'Invalid asset_type. Must be "source" or "target"'}, status_code = HTTP_400_BAD_REQUEST) + + if asset_type == 'source': + asset_ids = body.get('asset_ids', []) + + if not isinstance(asset_ids, list): + return JSONResponse({'message': 'asset_ids must be an array'}, status_code = HTTP_400_BAD_REQUEST) + + if not asset_ids: + return JSONResponse({'message': 'asset_ids cannot be empty'}, status_code = HTTP_400_BAD_REQUEST) + + source_paths = [] + for asset_id in asset_ids: + asset = asset_store.get_asset(asset_id) + if not asset: + return JSONResponse({'message': f'Source asset not found: {asset_id}'}, status_code = HTTP_404_NOT_FOUND) + source_paths.append(asset['path']) + + state_manager.set_item('source_paths', source_paths) + + __api_args__ = args_store.filter_api_args(state_manager.get_state()) #type:ignore[arg-type] + return JSONResponse(state_manager.collect_state(__api_args__), status_code = HTTP_200_OK) + + if asset_type == 'target': + asset_id = body.get('asset_id') + + if not asset_id: + return JSONResponse({'message': 'Missing required field: asset_id'}, status_code = HTTP_400_BAD_REQUEST) + + if not isinstance(asset_id, str): + return JSONResponse({'message': 'asset_id must be a string'}, status_code = HTTP_400_BAD_REQUEST) + + asset = asset_store.get_asset(asset_id) + if not asset: + return JSONResponse({'message': f'Target asset not found: {asset_id}'}, status_code = HTTP_404_NOT_FOUND) + + state_manager.set_item('target_path', asset['path']) + + __api_args__ = args_store.filter_api_args(state_manager.get_state()) #type:ignore[arg-type] + return JSONResponse(state_manager.collect_state(__api_args__), status_code = HTTP_200_OK) + api_args = args_store.get_api_args() + logger.info(f'[State] Normal update - body keys: {list(body.keys())}', __name__) for key, value in body.items(): if key in api_args: state_manager.set_item(key, value) + logger.debug(f'[State] Set {key} = {value}', __name__) + else: + logger.warn(f'[State] Skipped {key} (not in api_args)', __name__) __api_args__ = args_store.filter_api_args(state_manager.get_state()) #type:ignore[arg-type] return JSONResponse(state_manager.collect_state(__api_args__), status_code = HTTP_200_OK) diff --git a/facefusion/apis/timeline.py b/facefusion/apis/timeline.py new file mode 100644 index 00000000..81afca8a --- /dev/null +++ b/facefusion/apis/timeline.py @@ -0,0 +1,210 @@ +import base64 +import subprocess +from typing import List, Optional + +import cv2 +import numpy +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.status import HTTP_200_OK, HTTP_400_BAD_REQUEST + +from facefusion import logger +from facefusion.asset_store import get_asset +from facefusion.filesystem import is_video +from facefusion.video_manager import get_video_capture +from facefusion.vision import fit_contain_frame + + +def extract_frame_at_timestamp(stream_url : str, timestamp : float, width : int, height : int) -> Optional[numpy.ndarray]: + ffmpeg_command =\ + [ + 'ffmpeg', + '-user_agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', + '-ss', str(timestamp), + '-i', stream_url, + '-vf', f'scale={width}:{height}', + '-frames:v', '1', + '-f', 'rawvideo', + '-pix_fmt', 'bgr24', + '-' + ] + + try: + result = subprocess.run(ffmpeg_command, capture_output = True, timeout = 10) + if result.returncode == 0 and result.stdout: + frame_size = width * height * 3 + if len(result.stdout) >= frame_size: + frame = numpy.frombuffer(result.stdout[:frame_size], dtype = numpy.uint8).reshape((height, width, 3)) + return frame + except Exception as e: + logger.debug(f'Failed to extract frame at {timestamp}s: {e}', __name__) + + return None + + +async def get_timeline(request: Request) -> JSONResponse: + """ + Return N preview frames (as base64 JPEGs) from the target video, + resized to specified resolution for timeline preview. + + Route: /timeline/{count:int}?target_path=...&is_remote_stream=true&duration=120&fps=30&target_width=1920&target_height=1080&width=160&height=120 + """ + + # Extract and validate requested count + try: + count = int(request.path_params.get('count', 0)) + except (TypeError, ValueError): + return JSONResponse({'message': 'Invalid count parameter'}, status_code=HTTP_400_BAD_REQUEST) + + if count <= 0: + return JSONResponse({'message': 'Count must be a positive integer'}, status_code=HTTP_400_BAD_REQUEST) + + # Extract and validate preview resolution parameters + try: + preview_width = int(request.query_params.get('width', 160)) + preview_height = int(request.query_params.get('height', 120)) + except (TypeError, ValueError): + return JSONResponse({'message': 'Invalid width or height parameter'}, status_code=HTTP_400_BAD_REQUEST) + + if preview_width <= 0 or preview_height <= 0 or preview_width > 1920 or preview_height > 1080: + return JSONResponse({'message': 'Width and height must be between 1 and 1920x1080'}, status_code=HTTP_400_BAD_REQUEST) + + # Extract target_path or asset_id (one is required) + target_path = request.query_params.get('target_path') + asset_id = request.query_params.get('asset_id') + + # Extract is_remote_stream flag + is_remote_stream_param = request.query_params.get('is_remote_stream', 'false').lower() + is_remote_stream = is_remote_stream_param in ['true', '1', 'yes'] + + # Resolve asset_id to path if provided (for local files) + if asset_id and not target_path: + from facefusion.session_context import get_session_id + + asset = get_asset(asset_id) + if not asset: + return JSONResponse({'message': f'Asset not found: {asset_id}'}, status_code=HTTP_400_BAD_REQUEST) + + # Verify asset belongs to current session (security) + if asset.get('session_id') != get_session_id(): + return JSONResponse({'message': 'Asset not found'}, status_code=HTTP_400_BAD_REQUEST) + + target_path = asset.get('path') + if not target_path: + return JSONResponse({'message': 'Asset has no path'}, status_code=HTTP_400_BAD_REQUEST) + + is_remote_stream = False # Assets are always local files + logger.debug(f'Resolved asset_id {asset_id} to path for timeline preview', __name__) + + # Now check if we have a target_path + if not target_path: + return JSONResponse({'message': 'Missing required parameter: either target_path or asset_id'}, status_code=HTTP_400_BAD_REQUEST) + + # Extract video metadata (optional for local files, required for remote streams) + duration = None + fps = None + width = 1280 + height = 720 + + if request.query_params.get('duration'): + try: + duration = float(request.query_params.get('duration')) + except (TypeError, ValueError): + return JSONResponse({'message': 'Invalid duration parameter'}, status_code=HTTP_400_BAD_REQUEST) + + if request.query_params.get('fps'): + try: + fps = float(request.query_params.get('fps')) + except (TypeError, ValueError): + return JSONResponse({'message': 'Invalid fps parameter'}, status_code=HTTP_400_BAD_REQUEST) + + if request.query_params.get('target_width'): + try: + width = int(request.query_params.get('target_width')) + except (TypeError, ValueError): + return JSONResponse({'message': 'Invalid target_width parameter'}, status_code=HTTP_400_BAD_REQUEST) + + if request.query_params.get('target_height'): + try: + height = int(request.query_params.get('target_height')) + except (TypeError, ValueError): + return JSONResponse({'message': 'Invalid target_height parameter'}, status_code=HTTP_400_BAD_REQUEST) + + previews: List[str] = [] + + if is_remote_stream: + if not duration or duration <= 0: + return JSONResponse({'message': 'Duration not available for remote stream'}, status_code=HTTP_400_BAD_REQUEST) + + frame_total = 0 + if duration and fps: + try: + frame_total = int(float(duration) * float(fps)) + except Exception: + frame_total = 0 + + sample_count = min(count, frame_total) if frame_total > 0 else count + timestamps = list(numpy.linspace(0, float(duration), num=sample_count, endpoint=False)) + + logger.info(f'Extracting {sample_count} frames from remote stream using ffmpeg', __name__) + + for timestamp in timestamps: + frame = extract_frame_at_timestamp(target_path, timestamp, width, height) + if frame is None: + logger.warn(f'Failed to extract frame at {timestamp}s', __name__) + continue + + thumb_bgr = fit_contain_frame(frame, (preview_width, preview_height)) + if thumb_bgr.shape[1] != preview_width or thumb_bgr.shape[0] != preview_height: + thumb_bgr = cv2.resize(thumb_bgr, (preview_width, preview_height)) + + ok_enc, buf = cv2.imencode('.jpg', thumb_bgr, [cv2.IMWRITE_JPEG_QUALITY, 50]) + if not ok_enc: + logger.warn(f'JPEG encode failed for timestamp {timestamp}s', __name__) + continue + + b64 = base64.b64encode(buf.tobytes()).decode('ascii') + previews.append(b64) + else: + video_capture = get_video_capture(target_path) + if not video_capture or not video_capture.isOpened(): + logger.error(f'Unable to open video capture for target: {target_path}', __name__) + return JSONResponse({'message': 'Unable to open target video'}, status_code=HTTP_400_BAD_REQUEST) + + frame_total = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + + if frame_total <= 0 and is_video(target_path): + return JSONResponse({'message': 'Could not determine frame count for target video'}, status_code=HTTP_400_BAD_REQUEST) + + sample_count = min(count, frame_total) + indices: List[int] = list(numpy.linspace(1, frame_total, num=sample_count, endpoint=True, dtype=int)) + + for frame_number in indices: + video_capture.set(cv2.CAP_PROP_POS_FRAMES, max(0, frame_number - 1)) + ok_read, frame = video_capture.read() + if not ok_read or frame is None: + logger.warn(f'Failed reading frame {frame_number}', __name__) + continue + + thumb_bgr = fit_contain_frame(frame, (preview_width, preview_height)) + if thumb_bgr.shape[1] != preview_width or thumb_bgr.shape[0] != preview_height: + thumb_bgr = cv2.resize(thumb_bgr, (preview_width, preview_height)) + + ok_enc, buf = cv2.imencode('.jpg', thumb_bgr, [cv2.IMWRITE_JPEG_QUALITY, 50]) + if not ok_enc: + logger.warn(f'JPEG encode failed for frame {frame_number}', __name__) + continue + b64 = base64.b64encode(buf.tobytes()).decode('ascii') + previews.append(b64) + + logger.info(f'Returned {len(previews)}/{sample_count} timeline frames at {preview_width}x{preview_height}', __name__) + + return JSONResponse({ + 'message': 'ok', + 'count': len(previews), + 'requested': count, + 'width': preview_width, + 'height': preview_height, + 'format': 'jpeg', + 'frames': previews + }, status_code=HTTP_200_OK) diff --git a/facefusion/apis/version.py b/facefusion/apis/version.py new file mode 100644 index 00000000..20a6ae83 --- /dev/null +++ b/facefusion/apis/version.py @@ -0,0 +1,93 @@ +import subprocess +from functools import lru_cache +from typing import Optional + +from starlette.datastructures import Headers +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send +from starlette.websockets import WebSocket + + +@lru_cache(maxsize = 1) +def get_api_version() -> str: + try: + result = subprocess.run(['git', 'rev-parse', 'HEAD'], capture_output = True, text = True, check = True) + return result.stdout.strip() + except Exception: + return 'unknown' + + +def check_version_match(request : Request) -> Optional[JSONResponse]: + client_version = request.headers.get('X-API-Version') + server_version = get_api_version() + + if not client_version: + return JSONResponse({'error': 'Missing X-API-Version header', 'server_version': server_version}, status_code = 400) + + if client_version != server_version: + return JSONResponse({'error': 'Version mismatch', 'client_version': client_version, 'server_version': server_version}, status_code = 409) + + return None + + +def check_version_match_websocket(websocket : WebSocket) -> Optional[str]: + client_version = websocket.headers.get('X-API-Version') + server_version = get_api_version() + + if not client_version: + return f'Missing X-API-Version header, server version: {server_version}' + + if client_version != server_version: + return f'Version mismatch: client={client_version}, server={server_version}' + + return None + + +async def version_guard_middleware(scope : Scope, receive : Receive, send : Send, app : ASGIApp) -> None: + if scope['type'] == 'http': + headers = Headers(scope = scope) + client_version = headers.get('X-API-Version') + server_version = get_api_version() + + if not client_version: + response = JSONResponse({'error': 'Missing X-API-Version header', 'server_version': server_version}, status_code = 400) + await response(scope, receive, send) + return + + if client_version != server_version: + response = JSONResponse({'error': 'Version mismatch', 'client_version': client_version, 'server_version': server_version}, status_code = 409) + await response(scope, receive, send) + return + + if scope['type'] == 'websocket': + headers = Headers(scope = scope) + client_version = headers.get('X-API-Version') + + # For WebSocket connections, also check subprotocols since browsers can't set custom headers + if not client_version: + protocol_header = headers.get('Sec-WebSocket-Protocol') + if protocol_header: + # Parse subprotocols to find api_version + protocols = [p.strip() for p in protocol_header.split(',')] + for protocol in protocols: + if protocol.startswith('api_version.'): + client_version = protocol.split('.', 1)[1] + break + + server_version = get_api_version() + + if not client_version or client_version != server_version: + websocket = WebSocket(scope, receive = receive, send = send) + reason = f'Missing X-API-Version header, server version: {server_version}' if not client_version else f'Version mismatch: client={client_version}, server={server_version}' + await websocket.close(code = 1008, reason = reason) + return + + await app(scope, receive, send) + + +def create_version_guard(app : ASGIApp) -> ASGIApp: + async def version_guard_app(scope : Scope, receive : Receive, send : Send) -> None: + await version_guard_middleware(scope, receive, send, app) + + return version_guard_app diff --git a/facefusion/asset_store.py b/facefusion/asset_store.py new file mode 100644 index 00000000..26f3db0f --- /dev/null +++ b/facefusion/asset_store.py @@ -0,0 +1,111 @@ +import os +import uuid +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, TypeAlias + +from facefusion import filesystem, state_manager +from facefusion.session_context import get_session_id + +AssetRegistry : TypeAlias = Dict[str, Dict[str, Any]] + + +def get_asset_registry() -> AssetRegistry: + registry = state_manager.get_item('asset_registry') + if not registry: + registry = {} + state_manager.set_item('asset_registry', registry) + return registry + + +def register(asset_type : str, file_path : str, filename : str = None, metadata : Optional[Dict[str, Any]] = None) -> str: + if asset_type not in ['source', 'target', 'output']: + raise ValueError(f"Invalid asset_type: {asset_type}. Must be 'source', 'target', or 'output'") + + asset_id = str(uuid.uuid4()) + session_id = get_session_id() + + if not session_id: + raise ValueError("No active session - cannot register asset without session_id") + + if not filename: + filename = os.path.basename(file_path) + + file_size = os.path.getsize(file_path) + file_format = filesystem.get_file_format(file_path) + media_type = None + + if filesystem.is_image(file_path): + media_type = 'image' + if filesystem.is_video(file_path): + media_type = 'video' + if filesystem.is_audio(file_path): + media_type = 'audio' + + asset_data =\ + { + 'id': asset_id, + 'session_id': session_id, + 'type': asset_type, + 'media_type': media_type, + 'format': file_format, + 'path': file_path, + 'filename': filename, + 'size': file_size, + 'created_at': datetime.now(timezone.utc).isoformat() + } + + if metadata: + asset_data['metadata'] = metadata + + registry = get_asset_registry() + registry[asset_id] = asset_data + state_manager.set_item('asset_registry', registry) + + return asset_id + + +def get_asset(asset_id : str) -> Optional[Dict[str, Any]]: + registry = get_asset_registry() + return registry.get(asset_id) + + +def list_assets(asset_type : Optional[str] = None) -> List[Dict[str, Any]]: + registry = get_asset_registry() + session_id = get_session_id() + + assets = list(registry.values()) + + if session_id: + assets = [a for a in assets if a.get('session_id') == session_id] + + if asset_type: + if asset_type not in ['source', 'target', 'output']: + raise ValueError(f"Invalid asset_type: {asset_type}") + assets = [a for a in assets if a.get('type') == asset_type] + + return assets + + +def delete_asset(asset_id : str) -> bool: + registry = get_asset_registry() + asset = registry.get(asset_id) + + if not asset: + return False + + file_path = asset.get('path') + if file_path and os.path.exists(file_path): + os.remove(file_path) + + del registry[asset_id] + state_manager.set_item('asset_registry', registry) + + return True + + +def cleanup_session_assets(session_id : str) -> None: + registry = get_asset_registry() + assets_to_delete = [aid for aid, asset in registry.items() if asset.get('session_id') == session_id] + + for asset_id in assets_to_delete: + delete_asset(asset_id) diff --git a/facefusion/choices.py b/facefusion/choices.py index cefd2d03..041ca20e 100755 --- a/facefusion/choices.py +++ b/facefusion/choices.py @@ -45,7 +45,7 @@ face_mask_regions : List[FaceMaskRegion] = list(face_mask_region_set.keys()) voice_extractor_models : List[VoiceExtractorModel] = [ 'kim_vocal_1', 'kim_vocal_2', 'uvr_mdxnet' ] -workflows : List[WorkFlow] = [ 'auto', 'audio-to-image:frames', 'audio-to-image:video', 'image-to-image', 'image-to-video', 'image-to-video:frames' ] +workflows : List[WorkFlow] = [ 'auto', 'audio-to-image', 'image-to-image', 'image-to-video' ] audio_type_set : AudioTypeSet =\ { diff --git a/facefusion/core.py b/facefusion/core.py index 8d0801ff..10d8127b 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -12,15 +12,15 @@ from facefusion.apis.core import create_api from facefusion.args_helper import apply_args from facefusion.download import conditional_download_hashes, conditional_download_sources from facefusion.exit_helper import hard_exit, signal_exit -from facefusion.filesystem import get_file_extension, has_audio, has_image, has_video -from facefusion.filesystem import get_file_name, resolve_file_paths, resolve_file_pattern +from facefusion.filesystem import get_file_extension, get_file_name, resolve_file_paths, resolve_file_pattern +from facefusion.filesystem import has_audio, has_image, has_video from facefusion.jobs import job_helper, job_manager, job_runner from facefusion.jobs.job_list import compose_job_list from facefusion.processors.core import get_processors_modules from facefusion.program import create_program from facefusion.program_helper import validate_args from facefusion.types import Args, ErrorCode, WorkFlow -from facefusion.workflows import audio_to_image, audio_to_image_as_frames, image_to_image, image_to_video, image_to_video_as_frames +from facefusion.workflows import audio_to_image, image_to_image, image_to_video def cli() -> None: @@ -336,29 +336,21 @@ def conditional_process() -> ErrorCode: if not processor_module.pre_process('output'): return 2 - if state_manager.get_item('workflow') == 'audio-to-image:video': + if state_manager.get_item('workflow') == 'audio-to-image': return audio_to_image.process(start_time) - if state_manager.get_item('workflow') == 'audio-to-image:frames': - return audio_to_image_as_frames.process(start_time) if state_manager.get_item('workflow') == 'image-to-image': return image_to_image.process(start_time) if state_manager.get_item('workflow') == 'image-to-video': return image_to_video.process(start_time) - if state_manager.get_item('workflow') == 'image-to-video:frames': - return image_to_video_as_frames.process(start_time) return 0 def detect_workflow() -> WorkFlow: if has_video([ state_manager.get_item('target_path') ]): - if get_file_extension(state_manager.get_item('output_path')): - return 'image-to-video' - return 'image-to-video:frames' + return 'image-to-video' if has_audio(state_manager.get_item('source_paths')) and has_image([ state_manager.get_item('target_path') ]): - if get_file_extension(state_manager.get_item('output_path')): - return 'audio-to-image:video' - return 'audio-to-image:frames' + return 'audio-to-image' return 'image-to-image' diff --git a/facefusion/curl_builder.py b/facefusion/curl_builder.py index b7252ed5..41302dfc 100644 --- a/facefusion/curl_builder.py +++ b/facefusion/curl_builder.py @@ -16,7 +16,7 @@ def chain(*commands : List[Command]) -> List[Command]: return list(itertools.chain(*commands)) -def ping(url : str) -> List[Command]: +def head(url : str) -> List[Command]: return [ '-I', url ] @@ -26,7 +26,3 @@ def download(url : str, download_file_path : str) -> List[Command]: def set_timeout(timeout : int) -> List[Command]: return [ '--connect-timeout', str(timeout) ] - - -def set_retry(retry : int) -> List[Command]: - return [ '--retry', str(retry) ] diff --git a/facefusion/download.py b/facefusion/download.py index 8ba0c1e7..6878453c 100644 --- a/facefusion/download.py +++ b/facefusion/download.py @@ -29,8 +29,7 @@ def conditional_download(download_directory_path : str, urls : List[str]) -> Non with tqdm(total = download_size, initial = initial_size, desc = translator.get('downloading'), unit = 'B', unit_scale = True, unit_divisor = 1024, ascii = ' =', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: commands = curl_builder.chain( curl_builder.download(url, download_file_path), - curl_builder.set_timeout(5), - curl_builder.set_retry(5) + curl_builder.set_timeout(5) ) open_curl(commands) current_size = initial_size @@ -45,7 +44,7 @@ def conditional_download(download_directory_path : str, urls : List[str]) -> Non @lru_cache(maxsize = 64) def get_static_download_size(url : str) -> int: commands = curl_builder.chain( - curl_builder.ping(url), + curl_builder.head(url), curl_builder.set_timeout(5) ) process = open_curl(commands) @@ -63,7 +62,7 @@ def get_static_download_size(url : str) -> int: @lru_cache(maxsize = 64) def ping_static_url(url : str) -> bool: commands = curl_builder.chain( - curl_builder.ping(url), + curl_builder.head(url), curl_builder.set_timeout(5) ) process = open_curl(commands) diff --git a/facefusion/filesystem.py b/facefusion/filesystem.py index 22d51439..42bfe727 100644 --- a/facefusion/filesystem.py +++ b/facefusion/filesystem.py @@ -170,13 +170,6 @@ def create_directory(directory_path : str) -> bool: return False -def move_directory(directory_path : str, move_path : str) -> bool: - if is_directory(directory_path): - shutil.move(directory_path, move_path) - return is_directory(move_path) - return False - - def remove_directory(directory_path : str) -> bool: if is_directory(directory_path): shutil.rmtree(directory_path, ignore_errors = True) diff --git a/facefusion/installer.py b/facefusion/installer.py index 4deebe2f..edbbd355 100644 --- a/facefusion/installer.py +++ b/facefusion/installer.py @@ -10,10 +10,9 @@ from types import FrameType from facefusion import metadata from facefusion.common_helper import is_linux, is_windows -LOCALES =\ +LOCALS =\ { 'install_dependency': 'install the {dependency} package', - 'force_reinstall': 'force reinstall of packages', 'skip_conda': 'skip the conda environment check', 'conda_not_activated': 'conda is not activated' } @@ -27,16 +26,14 @@ if is_windows() or is_linux(): if is_windows(): ONNXRUNTIME_SET['directml'] = ('onnxruntime-directml', '1.23.0') if is_linux(): - ONNXRUNTIME_SET['migraphx'] = ('onnxruntime-migraphx', '1.23.0') - ONNXRUNTIME_SET['rocm'] = ('onnxruntime_rocm', '1.22.1', '7.0.2') #type:ignore[assignment] + ONNXRUNTIME_SET['rocm'] = ('onnxruntime-rocm', '1.21.0') def cli() -> None: signal.signal(signal.SIGINT, signal_exit) program = ArgumentParser(formatter_class = partial(HelpFormatter, max_help_position = 50)) - program.add_argument('--onnxruntime', help = LOCALES.get('install_dependency').format(dependency = 'onnxruntime'), choices = ONNXRUNTIME_SET.keys(), required = True) - program.add_argument('--force-reinstall', help = LOCALES.get('force_reinstall'), action = 'store_true') - program.add_argument('--skip-conda', help = LOCALES.get('skip_conda'), action = 'store_true') + program.add_argument('--onnxruntime', help = LOCALS.get('install_dependency').format(dependency = 'onnxruntime'), choices = ONNXRUNTIME_SET.keys(), required = True) + program.add_argument('--skip-conda', help = LOCALS.get('skip_conda'), action = 'store_true') program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version') run(program) @@ -48,13 +45,10 @@ def signal_exit(signum : int, frame : FrameType) -> None: def run(program : ArgumentParser) -> None: args = program.parse_args() has_conda = 'CONDA_PREFIX' in os.environ - commands = [ shutil.which('pip'), 'install' ] - - if args.force_reinstall: - commands.append('--force-reinstall') + onnxruntime_name, onnxruntime_version = ONNXRUNTIME_SET.get(args.onnxruntime) if not args.skip_conda and not has_conda: - sys.stdout.write(LOCALES.get('conda_not_activated') + os.linesep) + sys.stdout.write(LOCALS.get('conda_not_activated') + os.linesep) sys.exit(1) with open('requirements.txt') as file: @@ -62,21 +56,17 @@ def run(program : ArgumentParser) -> None: for line in file.readlines(): __line__ = line.strip() if not __line__.startswith('onnxruntime'): - commands.append(__line__) + subprocess.call([ shutil.which('pip'), 'install', line, '--force-reinstall' ]) if args.onnxruntime == 'rocm': - onnxruntime_name, onnxruntime_version, rocm_version = ONNXRUNTIME_SET.get(args.onnxruntime) #type:ignore[misc] python_id = 'cp' + str(sys.version_info.major) + str(sys.version_info.minor) if python_id in [ 'cp310', 'cp312' ]: - wheel_name = onnxruntime_name + '-' + onnxruntime_version + '-' + python_id + '-' + python_id + '-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl' - wheel_url = 'https://repo.radeon.com/rocm/manylinux/rocm-rel-' + rocm_version + '/' + wheel_name - commands.append(wheel_url) + wheel_name = 'onnxruntime_rocm-' + onnxruntime_version + '-' + python_id + '-' + python_id + '-linux_x86_64.whl' + wheel_url = 'https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/' + wheel_name + subprocess.call([ shutil.which('pip'), 'install', wheel_url, '--force-reinstall' ]) else: - onnxruntime_name, onnxruntime_version = ONNXRUNTIME_SET.get(args.onnxruntime) - commands.append(onnxruntime_name + '==' + onnxruntime_version) - - subprocess.call(commands) + subprocess.call([ shutil.which('pip'), 'install', onnxruntime_name + '==' + onnxruntime_version, '--force-reinstall' ]) if args.onnxruntime == 'cuda' and has_conda: library_paths = [] @@ -107,3 +97,4 @@ def run(program : ArgumentParser) -> None: library_paths = list(dict.fromkeys([ library_path for library_path in library_paths if os.path.exists(library_path) ])) subprocess.call([ shutil.which('conda'), 'env', 'config', 'vars', 'set', 'PATH=' + os.pathsep.join(library_paths) ]) + diff --git a/facefusion/jobs/job_helper.py b/facefusion/jobs/job_helper.py index 71b1c00c..d7e90218 100644 --- a/facefusion/jobs/job_helper.py +++ b/facefusion/jobs/job_helper.py @@ -13,8 +13,6 @@ def get_step_output_path(job_id : str, step_index : int, output_path : str) -> O if output_file_name and output_file_extension: return os.path.join(output_directory_path, output_file_name + '-' + job_id + '-' + str(step_index) + output_file_extension) - if output_file_path and output_directory_path: - return os.path.join(output_directory_path, output_file_path + '-' + job_id + '-' + str(step_index)) return None diff --git a/facefusion/jobs/job_runner.py b/facefusion/jobs/job_runner.py index 8a05d38a..23a0e38b 100644 --- a/facefusion/jobs/job_runner.py +++ b/facefusion/jobs/job_runner.py @@ -1,7 +1,5 @@ -import os - from facefusion.ffmpeg import concat_video -from facefusion.filesystem import are_images, are_videos, copy_file, create_directory, is_directory, is_file, move_directory, move_file, remove_directory, remove_file, resolve_file_paths +from facefusion.filesystem import are_images, are_videos, move_file, remove_file from facefusion.jobs import job_helper, job_manager from facefusion.types import JobOutputSet, JobStep, ProcessStep @@ -61,8 +59,6 @@ def run_step(job_id : str, step_index : int, step : JobStep, process_step : Proc output_path = step_args.get('output_path') step_output_path = job_helper.get_step_output_path(job_id, step_index, output_path) - if is_directory(output_path): - return move_directory(output_path, step_output_path) and job_manager.set_step_status(job_id, step_index, 'completed') return move_file(output_path, step_output_path) and job_manager.set_step_status(job_id, step_index, 'completed') job_manager.set_step_status(job_id, step_index, 'failed') return False @@ -83,26 +79,13 @@ def finalize_steps(job_id : str) -> bool: output_set = collect_output_set(job_id) for output_path, temp_output_paths in output_set.items(): - has_videos = are_videos(temp_output_paths) - has_images = are_images(temp_output_paths) - - if has_videos: + if are_videos(temp_output_paths): if not concat_video(output_path, temp_output_paths): return False - if not has_videos and has_images: + if are_images(temp_output_paths): for temp_output_path in temp_output_paths: if not move_file(temp_output_path, output_path): return False - if not has_videos and not has_images: - if not create_directory(output_path): - return False - - for temp_output_path in temp_output_paths: - if is_directory(temp_output_path): - temp_frame_paths = resolve_file_paths(temp_output_path) - for temp_frame_path in temp_frame_paths: - if not copy_file(temp_frame_path, os.path.join(output_path, os.path.basename(temp_frame_path))): - return False return True @@ -111,12 +94,8 @@ def clean_steps(job_id: str) -> bool: for temp_output_paths in output_set.values(): for temp_output_path in temp_output_paths: - if is_file(temp_output_path): - if not remove_file(temp_output_path): - return False - if is_directory(temp_output_path): - if not remove_directory(temp_output_path): - return False + if not remove_file(temp_output_path): + return False return True diff --git a/facefusion/locales.py b/facefusion/locals.py similarity index 98% rename from facefusion/locales.py rename to facefusion/locals.py index 9f23050c..cbb8c7b1 100644 --- a/facefusion/locales.py +++ b/facefusion/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { @@ -40,8 +40,6 @@ LOCALES : Locales =\ 'processing_stopped': 'processing stopped', 'processing_image_succeeded': 'processing to image succeeded in {seconds} seconds', 'processing_image_failed': 'processing to image failed', - 'processing_frames_succeeded': 'processing to frames succeeded in {seconds} seconds', - 'processing_frames_failed': 'processing to frames failed', 'processing_video_succeeded': 'processing to video succeeded in {seconds} seconds', 'processing_video_failed': 'processing to video failed', 'choose_image_source': 'choose an image for the source', diff --git a/facefusion/processors/modules/age_modifier/core.py b/facefusion/processors/modules/age_modifier/core.py index 68d5d2fa..d6929a3d 100755 --- a/facefusion/processors/modules/age_modifier/core.py +++ b/facefusion/processors/modules/age_modifier/core.py @@ -108,10 +108,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/age_modifier/locales.py b/facefusion/processors/modules/age_modifier/locals.py similarity index 83% rename from facefusion/processors/modules/age_modifier/locales.py rename to facefusion/processors/modules/age_modifier/locals.py index c72bbbb4..90cf9e06 100644 --- a/facefusion/processors/modules/age_modifier/locales.py +++ b/facefusion/processors/modules/age_modifier/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/background_remover/core.py b/facefusion/processors/modules/background_remover/core.py index dda62415..aa829b12 100644 --- a/facefusion/processors/modules/background_remover/core.py +++ b/facefusion/processors/modules/background_remover/core.py @@ -421,7 +421,7 @@ def register_args(program : ArgumentParser) -> None: group_processors = find_argument_group(program, 'processors') if group_processors: group_processors.add_argument('--background-remover-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'background_remover_model', 'rmbg_2.0'), choices = background_remover_choices.background_remover_models) - group_processors.add_argument('--background-remover-color', help = translator.get('help.color', __package__), type = partial(sanitize_int_range, int_range = background_remover_choices.background_remover_color_range), default = config.get_int_list('processors', 'background_remover_color', '0 0 0 0'), nargs = '+') + group_processors.add_argument('--background-remover-color', help = translator.get('help.color', __package__), type = partial(sanitize_int_range, int_range = background_remover_choices.background_remover_color_range), default = config.get_int_list('processors', 'background_remover_color', '0 0 0 0'), nargs ='+') facefusion.args_store.register_args([ 'background_remover_model', 'background_remover_color' ], scopes = [ 'api', 'cli' ]) @@ -441,10 +441,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/background_remover/locales.py b/facefusion/processors/modules/background_remover/locals.py similarity index 88% rename from facefusion/processors/modules/background_remover/locales.py rename to facefusion/processors/modules/background_remover/locals.py index e65082c0..854f589f 100644 --- a/facefusion/processors/modules/background_remover/locales.py +++ b/facefusion/processors/modules/background_remover/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/deep_swapper/core.py b/facefusion/processors/modules/deep_swapper/core.py index 577df558..2d9ce0a4 100755 --- a/facefusion/processors/modules/deep_swapper/core.py +++ b/facefusion/processors/modules/deep_swapper/core.py @@ -299,10 +299,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/deep_swapper/locales.py b/facefusion/processors/modules/deep_swapper/locals.py similarity index 81% rename from facefusion/processors/modules/deep_swapper/locales.py rename to facefusion/processors/modules/deep_swapper/locals.py index 4c42c951..112de49d 100644 --- a/facefusion/processors/modules/deep_swapper/locales.py +++ b/facefusion/processors/modules/deep_swapper/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/expression_restorer/core.py b/facefusion/processors/modules/expression_restorer/core.py index 6a60419b..b3b14e2b 100755 --- a/facefusion/processors/modules/expression_restorer/core.py +++ b/facefusion/processors/modules/expression_restorer/core.py @@ -101,7 +101,7 @@ def register_args(program : ArgumentParser) -> None: if group_processors: group_processors.add_argument('--expression-restorer-model', help = translator.get('help.model', __package__), default = config.get_str_value('processors', 'expression_restorer_model', 'live_portrait'), choices = expression_restorer_choices.expression_restorer_models) group_processors.add_argument('--expression-restorer-factor', help = translator.get('help.factor', __package__), type = int, default = config.get_int_value('processors', 'expression_restorer_factor', '80'), choices = expression_restorer_choices.expression_restorer_factor_range, metavar = create_int_metavar(expression_restorer_choices.expression_restorer_factor_range)) - group_processors.add_argument('--expression-restorer-areas', help = translator.get('help.areas', __package__).format(choices = ', '.join(expression_restorer_choices.expression_restorer_areas)), default = config.get_str_list('processors', 'expression_restorer_areas', ' '.join(expression_restorer_choices.expression_restorer_areas)), choices = expression_restorer_choices.expression_restorer_areas, nargs = '+', metavar = 'EXPRESSION_RESTORER_AREAS') + group_processors.add_argument('--expression-restorer-areas', help = translator.get('help.areas', __package__).format(choices = ', '.join(expression_restorer_choices.expression_restorer_areas)), default = config.get_str_list('processors', 'expression_restorer_areas', ' '.join(expression_restorer_choices.expression_restorer_areas)), choices = expression_restorer_choices.expression_restorer_areas, nargs ='+', metavar ='EXPRESSION_RESTORER_AREAS') facefusion.args_store.register_args([ 'expression_restorer_model', 'expression_restorer_factor', 'expression_restorer_areas' ], scopes = [ 'api', 'cli' ]) @@ -125,10 +125,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/expression_restorer/locales.py b/facefusion/processors/modules/expression_restorer/locals.py similarity index 88% rename from facefusion/processors/modules/expression_restorer/locales.py rename to facefusion/processors/modules/expression_restorer/locals.py index 2e1c6da8..0fd4c7d0 100644 --- a/facefusion/processors/modules/expression_restorer/locales.py +++ b/facefusion/processors/modules/expression_restorer/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/face_debugger/core.py b/facefusion/processors/modules/face_debugger/core.py index 7ece4416..fe7f123b 100755 --- a/facefusion/processors/modules/face_debugger/core.py +++ b/facefusion/processors/modules/face_debugger/core.py @@ -46,10 +46,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/face_debugger/locales.py b/facefusion/processors/modules/face_debugger/locals.py similarity index 75% rename from facefusion/processors/modules/face_debugger/locales.py rename to facefusion/processors/modules/face_debugger/locals.py index db33a82e..412a68e9 100644 --- a/facefusion/processors/modules/face_debugger/locales.py +++ b/facefusion/processors/modules/face_debugger/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/face_editor/core.py b/facefusion/processors/modules/face_editor/core.py index 41380bdb..56430edb 100755 --- a/facefusion/processors/modules/face_editor/core.py +++ b/facefusion/processors/modules/face_editor/core.py @@ -176,10 +176,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/face_editor/locales.py b/facefusion/processors/modules/face_editor/locals.py similarity index 96% rename from facefusion/processors/modules/face_editor/locales.py rename to facefusion/processors/modules/face_editor/locals.py index c400c23e..86a134b3 100644 --- a/facefusion/processors/modules/face_editor/locales.py +++ b/facefusion/processors/modules/face_editor/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/face_enhancer/core.py b/facefusion/processors/modules/face_enhancer/core.py index 0b5a64fa..97df33fd 100755 --- a/facefusion/processors/modules/face_enhancer/core.py +++ b/facefusion/processors/modules/face_enhancer/core.py @@ -315,10 +315,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/face_enhancer/locales.py b/facefusion/processors/modules/face_enhancer/locals.py similarity index 86% rename from facefusion/processors/modules/face_enhancer/locales.py rename to facefusion/processors/modules/face_enhancer/locals.py index 5cd314d1..9e3d78d1 100644 --- a/facefusion/processors/modules/face_enhancer/locales.py +++ b/facefusion/processors/modules/face_enhancer/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/face_swapper/core.py b/facefusion/processors/modules/face_swapper/core.py index b48cf9b7..af330c10 100755 --- a/facefusion/processors/modules/face_swapper/core.py +++ b/facefusion/processors/modules/face_swapper/core.py @@ -551,10 +551,9 @@ def pre_process(mode : ProcessMode) -> bool: logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True @@ -576,10 +575,12 @@ def post_process() -> None: def swap_face(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: + logger.debug('swap_face - starting face swap', __name__) model_template = get_model_options().get('template') model_size = get_model_options().get('size') pixel_boost_size = unpack_resolution(state_manager.get_item('face_swapper_pixel_boost')) pixel_boost_total = pixel_boost_size[0] // model_size[0] + logger.debug(f'swap_face - model_template: {model_template}, model_size: {model_size}, pixel_boost: {pixel_boost_size}', __name__) crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, target_face.landmark_set.get('5/68'), model_template, pixel_boost_size) temp_vision_frames = [] crop_masks = [] @@ -760,12 +761,19 @@ def process_frame(inputs : FaceSwapperInputs) -> ProcessorOutputs: target_vision_frame = inputs.get('target_vision_frame') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') + logger.debug(f'process_frame - source_vision_frames count: {len(source_vision_frames) if source_vision_frames else 0}', __name__) source_face = extract_source_face(source_vision_frames) + logger.debug(f'process_frame - source_face extracted: {source_face is not None}', __name__) target_faces = select_faces(reference_vision_frame, target_vision_frame) + logger.debug(f'process_frame - target_faces count: {len(target_faces) if target_faces else 0}', __name__) if source_face and target_faces: + logger.debug(f'process_frame - swapping {len(target_faces)} faces', __name__) for target_face in target_faces: target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) temp_vision_frame = swap_face(source_face, target_face, temp_vision_frame) + logger.debug('process_frame - swap completed', __name__) + else: + logger.debug(f'process_frame - skipping swap (source_face={source_face is not None}, target_faces={target_faces is not None})', __name__) return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/face_swapper/locales.py b/facefusion/processors/modules/face_swapper/locals.py similarity index 87% rename from facefusion/processors/modules/face_swapper/locales.py rename to facefusion/processors/modules/face_swapper/locals.py index ac48146a..a1680a54 100644 --- a/facefusion/processors/modules/face_swapper/locales.py +++ b/facefusion/processors/modules/face_swapper/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/frame_colorizer/core.py b/facefusion/processors/modules/frame_colorizer/core.py index f81ad1d8..8a3f3bc6 100644 --- a/facefusion/processors/modules/frame_colorizer/core.py +++ b/facefusion/processors/modules/frame_colorizer/core.py @@ -207,10 +207,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/frame_colorizer/locales.py b/facefusion/processors/modules/frame_colorizer/locals.py similarity index 86% rename from facefusion/processors/modules/frame_colorizer/locales.py rename to facefusion/processors/modules/frame_colorizer/locals.py index 8c439944..cad3678e 100644 --- a/facefusion/processors/modules/frame_colorizer/locales.py +++ b/facefusion/processors/modules/frame_colorizer/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/frame_enhancer/core.py b/facefusion/processors/modules/frame_enhancer/core.py index 7b72e53b..ea97ecf9 100644 --- a/facefusion/processors/modules/frame_enhancer/core.py +++ b/facefusion/processors/modules/frame_enhancer/core.py @@ -594,10 +594,9 @@ def pre_process(mode : ProcessMode) -> bool: if mode in [ 'output', 'preview' ] and not is_image(state_manager.get_item('target_path')) and not is_video(state_manager.get_item('target_path')): logger.error(translator.get('choose_image_or_video_target') + translator.get('exclamation_mark'), __name__) return False - if state_manager.get_item('workflow') in [ 'audio-to-image:video', 'image-to-image', 'image-to-video' ]: - if mode == 'output' and not in_directory(state_manager.get_item('output_path')): - logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) - return False + if mode == 'output' and not in_directory(state_manager.get_item('output_path')): + logger.error(translator.get('specify_image_or_video_output') + translator.get('exclamation_mark'), __name__) + return False return True diff --git a/facefusion/processors/modules/frame_enhancer/locales.py b/facefusion/processors/modules/frame_enhancer/locals.py similarity index 82% rename from facefusion/processors/modules/frame_enhancer/locales.py rename to facefusion/processors/modules/frame_enhancer/locals.py index 2df9ff91..4b707a0f 100644 --- a/facefusion/processors/modules/frame_enhancer/locales.py +++ b/facefusion/processors/modules/frame_enhancer/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/processors/modules/lip_syncer/locales.py b/facefusion/processors/modules/lip_syncer/locals.py similarity index 81% rename from facefusion/processors/modules/lip_syncer/locales.py rename to facefusion/processors/modules/lip_syncer/locals.py index 62bc259f..bf53c31a 100644 --- a/facefusion/processors/modules/lip_syncer/locales.py +++ b/facefusion/processors/modules/lip_syncer/locals.py @@ -1,6 +1,6 @@ -from facefusion.types import Locales +from facefusion.types import Locals -LOCALES : Locales =\ +LOCALS : Locals =\ { 'en': { diff --git a/facefusion/program.py b/facefusion/program.py index 5ab48d02..0c728ab8 100755 --- a/facefusion/program.py +++ b/facefusion/program.py @@ -241,7 +241,7 @@ def create_execution_program() -> ArgumentParser: group_execution.add_argument('--execution-device-ids', help = translator.get('help.execution_device_ids'), type = int, default = config.get_str_list('execution', 'execution_device_ids', '0'), nargs = '+', metavar = 'EXECUTION_DEVICE_IDS') group_execution.add_argument('--execution-providers', help = translator.get('help.execution_providers').format(choices = ', '.join(available_execution_providers)), default = config.get_str_list('execution', 'execution_providers', get_first(available_execution_providers)), choices = available_execution_providers, nargs = '+', metavar = 'EXECUTION_PROVIDERS') group_execution.add_argument('--execution-thread-count', help = translator.get('help.execution_thread_count'), type = int, default = config.get_int_value('execution', 'execution_thread_count', '8'), choices = facefusion.choices.execution_thread_count_range, metavar = create_int_metavar(facefusion.choices.execution_thread_count_range)) - args_store.register_args([ 'execution_device_ids', 'execution_providers', 'execution_thread_count' ], scopes = [ 'cli', 'sys' ]) + args_store.register_args([ 'execution_device_ids', 'execution_providers', 'execution_thread_count' ], scopes = [ 'api', 'cli', 'sys' ]) return program diff --git a/facefusion/session_manager.py b/facefusion/session_manager.py index 06284147..f43ea4fa 100644 --- a/facefusion/session_manager.py +++ b/facefusion/session_manager.py @@ -14,7 +14,7 @@ def create_session() -> Session: 'access_token': secrets.token_urlsafe(64), 'refresh_token': secrets.token_urlsafe(64), 'created_at': datetime.now(), - 'expires_at': datetime.now() + timedelta(minutes = 10) + 'expires_at': datetime.now() + timedelta(minutes = 0.5) } return session diff --git a/facefusion/streamer.py b/facefusion/streamer.py index 66e523ee..0b3374ef 100644 --- a/facefusion/streamer.py +++ b/facefusion/streamer.py @@ -52,11 +52,10 @@ def process_stream_frame(target_vision_frame : VisionFrame) -> VisionFrame: temp_vision_mask = extract_vision_mask(temp_vision_frame) for processor_module in get_processors_modules(state_manager.get_item('processors')): - logger.disable() if processor_module.pre_process('stream'): - logger.enable() temp_vision_frame, temp_vision_mask = processor_module.process_frame( { + 'reference_vision_frame': target_vision_frame, 'source_vision_frames': source_vision_frames, 'source_audio_frame': source_audio_frame, 'source_voice_frame': source_voice_frame, @@ -64,7 +63,6 @@ def process_stream_frame(target_vision_frame : VisionFrame) -> VisionFrame: 'temp_vision_frame': temp_vision_frame, 'temp_vision_mask': temp_vision_mask }) - logger.enable() return temp_vision_frame diff --git a/facefusion/system.py b/facefusion/system.py new file mode 100644 index 00000000..e6a48127 --- /dev/null +++ b/facefusion/system.py @@ -0,0 +1,207 @@ +import os +import platform +from datetime import datetime +from functools import lru_cache +from typing import Optional + +import psutil + +from facefusion.types import CpuInfo, DiskInfo, LoadAverage, NetworkInfo, OperatingSystemInfo, PythonInfo, RamInfo, SystemInfo, TemperatureInfo + + +@lru_cache() +def detect_static_system_info() -> SystemInfo: + return detect_system_info() + + +def detect_system_info(temp_path : Optional[str] = None) -> SystemInfo: + return\ + { + 'operating_system': get_operating_system_info(), + 'python': get_python_info(), + 'cpu': get_cpu_info(), + 'ram': get_ram_info(), + 'disk': get_disk_info(temp_path), + 'temperatures': get_temperature_info(), + 'network': get_network_info(), + 'load_average': get_load_average() + } + + +def get_operating_system_info() -> OperatingSystemInfo: + boot_timestamp = psutil.boot_time() + boot_time = datetime.fromtimestamp(boot_timestamp) + uptime_seconds = int((datetime.now() - boot_time).total_seconds()) + + return\ + { + 'name': platform.system(), + 'architecture': platform.machine(), + 'platform': platform.platform(), + 'boot_time': boot_time.isoformat(), + 'uptime_seconds': uptime_seconds + } + + +def get_python_info() -> PythonInfo: + return\ + { + 'version': platform.python_version(), + 'implementation': platform.python_implementation() + } + + +def get_cpu_info() -> CpuInfo: + cpu_freq = psutil.cpu_freq() + cpu_percent = psutil.cpu_percent(interval = 0) + + cpu_info : CpuInfo =\ + { + 'model': get_cpu_model(), + 'physical_cores': psutil.cpu_count(logical = False), + 'logical_cores': psutil.cpu_count(logical = True), + 'usage_percent': cpu_percent + } + + if cpu_freq: + cpu_info['frequency'] =\ + { + 'current': cpu_freq.current, + 'min': cpu_freq.min, + 'max': cpu_freq.max + } + + return cpu_info + + +def get_cpu_model() -> Optional[str]: + if platform.system() == 'Linux': + try: + with open('/proc/cpuinfo', 'r') as f: + for line in f: + if line.startswith('model name'): + return line.split(':', 1)[1].strip() + except Exception: + pass + if platform.system() == 'Darwin': + try: + import subprocess + result = subprocess.run(['sysctl', '-n', 'machdep.cpu.brand_string'], capture_output = True, text = True) + if result.returncode == 0: + return result.stdout.strip() + except Exception: + pass + if platform.system() == 'Windows': + try: + import subprocess + result = subprocess.run(['wmic', 'cpu', 'get', 'name'], capture_output = True, text = True) + if result.returncode == 0: + lines = result.stdout.strip().split('\n') + if len(lines) > 1: + return lines[1].strip() + except Exception: + pass + return None + + +def get_ram_info() -> RamInfo: + virtual_memory = psutil.virtual_memory() + swap_memory = psutil.swap_memory() + + return\ + { + 'total': virtual_memory.total, + 'available': virtual_memory.available, + 'used': virtual_memory.used, + 'free': virtual_memory.free, + 'percent': virtual_memory.percent, + 'swap_total': swap_memory.total, + 'swap_used': swap_memory.used, + 'swap_free': swap_memory.free, + 'swap_percent': swap_memory.percent + } + + +def get_disk_info(temp_path : Optional[str] = None) -> Optional[DiskInfo]: + if temp_path is None: + temp_path = os.getcwd() + + target_mountpoint = None + target_mountpoint_len = 0 + + for partition in psutil.disk_partitions(): + if temp_path.startswith(partition.mountpoint): + if len(partition.mountpoint) > target_mountpoint_len: + target_mountpoint = partition.mountpoint + target_mountpoint_len = len(partition.mountpoint) + + if target_mountpoint: + try: + usage = psutil.disk_usage(target_mountpoint) + return\ + { + 'filesystem': next((p.fstype for p in psutil.disk_partitions() if p.mountpoint == target_mountpoint), 'unknown'), + 'total': usage.total, + 'used': usage.used, + 'free': usage.free, + 'percent': usage.percent + } + except PermissionError: + pass + + return None + + +def get_temperature_info() -> Optional[TemperatureInfo]: + if not hasattr(psutil, 'sensors_temperatures'): + return None + + try: + temps = psutil.sensors_temperatures() + if not temps: + return None + + temp_info : TemperatureInfo = {} + + for name, entries in temps.items(): + for entry in entries: + sensor_key = f'{name}_{entry.label}' if entry.label else name + temp_info[sensor_key] =\ + { + 'current': entry.current, + 'high': entry.high, + 'critical': entry.critical + } + + return temp_info + except Exception: + return None + + +def get_network_info() -> NetworkInfo: + net_io = psutil.net_io_counters() + + return\ + { + 'bytes_sent': net_io.bytes_sent, + 'bytes_recv': net_io.bytes_recv, + 'packets_sent': net_io.packets_sent, + 'packets_recv': net_io.packets_recv, + 'errin': net_io.errin, + 'errout': net_io.errout, + 'dropin': net_io.dropin, + 'dropout': net_io.dropout, + 'interfaces': {} + } + + +def get_load_average() -> Optional[LoadAverage]: + if hasattr(os, 'getloadavg'): + load1, load5, load15 = os.getloadavg() + return\ + { + 'load1': load1, + 'load5': load5, + 'load15': load15 + } + return None diff --git a/facefusion/translator.py b/facefusion/translator.py index 9f5fb4f1..f2f4e3ff 100644 --- a/facefusion/translator.py +++ b/facefusion/translator.py @@ -1,29 +1,29 @@ import importlib from typing import Optional -from facefusion.types import Language, LocalePoolSet, Locales +from facefusion.types import Language, LocalPoolSet, Locals -LOCALE_POOL_SET : LocalePoolSet = {} +LOCAL_POOL_SET : LocalPoolSet = {} CURRENT_LANGUAGE : Language = 'en' def __autoload__(module_name : str) -> None: try: - __locales__ = importlib.import_module(module_name + '.locales') - load(__locales__.LOCALES, module_name) + __locals__ = importlib.import_module(module_name + '.locals') + load(__locals__.LOCALS, module_name) except ImportError: pass -def load(__locales__ : Locales, module_name : str) -> None: - LOCALE_POOL_SET[module_name] = __locales__ +def load(__locals__ : Locals, module_name : str) -> None: + LOCAL_POOL_SET[module_name] = __locals__ def get(notation : str, module_name : str = 'facefusion') -> Optional[str]: - if module_name not in LOCALE_POOL_SET: + if module_name not in LOCAL_POOL_SET: __autoload__(module_name) - current = LOCALE_POOL_SET.get(module_name).get(CURRENT_LANGUAGE) + current = LOCAL_POOL_SET.get(module_name).get(CURRENT_LANGUAGE) for fragment in notation.split('.'): if fragment in current: diff --git a/facefusion/types.py b/facefusion/types.py index 53a46244..429283ab 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -52,10 +52,10 @@ FaceStore = TypedDict('FaceStore', }) Language = Literal['en'] -Locales : TypeAlias = Dict[Language, Dict[str, Any]] -LocalePoolSet : TypeAlias = Dict[str, Locales] +Locals : TypeAlias = Dict[Language, Dict[str, Any]] +LocalPoolSet : TypeAlias = Dict[str, Locals] -WorkFlow = Literal['auto', 'audio-to-image:frames', 'audio-to-image:video', 'image-to-image', 'image-to-video', 'image-to-video:frames'] +WorkFlow = Literal['auto', 'audio-to-image', 'image-to-image', 'image-to-video'] VideoCaptureSet : TypeAlias = Dict[str, cv2.VideoCapture] VideoWriterSet : TypeAlias = Dict[str, cv2.VideoWriter] @@ -152,10 +152,12 @@ FaceMaskAreaSet : TypeAlias = Dict[FaceMaskArea, List[int]] VoiceExtractorModel = Literal['kim_vocal_1', 'kim_vocal_2', 'uvr_mdxnet'] +MediaType = Literal['audio', 'image', 'video'] 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'] +StreamMediaType = Literal['video', 'image'] AudioTypeSet : TypeAlias = Dict[AudioFormat, str] ImageTypeSet : TypeAlias = Dict[ImageFormat, str] VideoTypeSet : TypeAlias = Dict[VideoFormat, str] @@ -233,6 +235,90 @@ ExecutionDevice = TypedDict('ExecutionDevice', 'utilization' : ExecutionDeviceUtilization }) +OperatingSystemInfo = TypedDict('OperatingSystemInfo', +{ + 'name' : str, + 'architecture' : str, + 'platform' : str, + 'boot_time' : str, + 'uptime_seconds' : int +}) +PythonInfo = TypedDict('PythonInfo', +{ + 'version' : str, + 'implementation' : str +}) +CpuFrequency = TypedDict('CpuFrequency', +{ + 'current' : float, + 'min' : float, + 'max' : float +}) +CpuInfo = TypedDict('CpuInfo', +{ + 'model' : Optional[str], + 'physical_cores' : Optional[int], + 'logical_cores' : Optional[int], + 'usage_percent' : float, + 'frequency' : Optional[CpuFrequency] +}, total = False) +RamInfo = TypedDict('RamInfo', +{ + 'total' : int, + 'available' : int, + 'used' : int, + 'free' : int, + 'percent' : float, + 'swap_total' : int, + 'swap_used' : int, + 'swap_free' : int, + 'swap_percent' : float +}) +DiskInfo = TypedDict('DiskInfo', +{ + 'filesystem' : str, + 'total' : int, + 'used' : int, + 'free' : int, + 'percent' : float +}) +TemperatureSensor = TypedDict('TemperatureSensor', +{ + 'current' : float, + 'high' : Optional[float], + 'critical' : Optional[float] +}) +TemperatureInfo : TypeAlias = Dict[str, TemperatureSensor] +NetworkInfo = TypedDict('NetworkInfo', +{ + 'bytes_sent' : int, + 'bytes_recv' : int, + 'packets_sent' : int, + 'packets_recv' : int, + 'errin' : int, + 'errout' : int, + 'dropin' : int, + 'dropout' : int, + 'interfaces' : Dict[str, Any] +}) +LoadAverage = TypedDict('LoadAverage', +{ + 'load1' : float, + 'load5' : float, + 'load15' : float +}) +SystemInfo = TypedDict('SystemInfo', +{ + 'operating_system' : OperatingSystemInfo, + 'python' : PythonInfo, + 'cpu' : CpuInfo, + 'ram' : RamInfo, + 'disk' : Optional[DiskInfo], + 'temperatures' : Optional[TemperatureInfo], + 'network' : NetworkInfo, + 'load_average' : Optional[LoadAverage] +}) + DownloadProvider = Literal['github', 'huggingface'] DownloadProviderValue = TypedDict('DownloadProviderValue', { diff --git a/facefusion/workflows/as_frames.py b/facefusion/workflows/as_frames.py deleted file mode 100644 index 88c39099..00000000 --- a/facefusion/workflows/as_frames.py +++ /dev/null @@ -1,46 +0,0 @@ -import os - -from facefusion import ffmpeg, logger, state_manager, translator -from facefusion.audio import restrict_trim_audio_frame -from facefusion.common_helper import get_first -from facefusion.filesystem import are_images, copy_file, create_directory, filter_audio_paths, resolve_file_paths -from facefusion.temp_helper import resolve_temp_frame_paths -from facefusion.time_helper import calculate_end_time -from facefusion.types import ErrorCode -from facefusion.vision import detect_image_resolution, restrict_image_resolution, scale_resolution -from facefusion.workflows.core import is_process_stopping - - -def create_temp_frames() -> ErrorCode: - state_manager.set_item('output_video_fps', 25.0) # TODO: set default fps value - source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) - output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) - temp_image_resolution = restrict_image_resolution(state_manager.get_item('target_path'), output_image_resolution) - trim_frame_start, trim_frame_end = restrict_trim_audio_frame(source_audio_path, state_manager.get_item('output_video_fps'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) - - if ffmpeg.spawn_frames(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_image_resolution, state_manager.get_item('output_video_fps'), trim_frame_start, trim_frame_end): - logger.debug(translator.get('spawning_frames_succeeded'), __name__) - else: - if is_process_stopping(): - return 4 - logger.error(translator.get('spawning_frames_failed'), __name__) - return 1 - return 0 - - -def copy_temp_frames() -> ErrorCode: - temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) - - for temp_frame_path in temp_frame_paths: - if not create_directory(state_manager.get_item('output_path')) or not copy_file(temp_frame_path, os.path.join(state_manager.get_item('output_path'), os.path.basename(temp_frame_path))): - return 1 - return 0 - - -def finalize_frames(start_time : float) -> ErrorCode: - if are_images(resolve_file_paths(state_manager.get_item('output_path'))): - logger.info(translator.get('processing_frames_succeeded').format(seconds = calculate_end_time(start_time)), __name__) - else: - logger.error(translator.get('processing_frames_failed'), __name__) - return 1 - return 0 diff --git a/facefusion/workflows/audio_to_image.py b/facefusion/workflows/audio_to_image.py index e7d9607b..ccd7a506 100644 --- a/facefusion/workflows/audio_to_image.py +++ b/facefusion/workflows/audio_to_image.py @@ -1,10 +1,12 @@ from functools import partial -from facefusion import process_manager +from facefusion import ffmpeg, logger, process_manager, state_manager, translator +from facefusion.audio import restrict_trim_audio_frame +from facefusion.common_helper import get_first +from facefusion.filesystem import filter_audio_paths from facefusion.types import ErrorCode -from facefusion.workflows.as_frames import create_temp_frames -from facefusion.workflows.core import analyse_image, clear, process_frames, setup -from facefusion.workflows.to_video import finalize_video, merge_frames, restore_audio +from facefusion.vision import detect_image_resolution, restrict_image_resolution, scale_resolution +from facefusion.workflows.core import analyse_image, clear, finalize_video, is_process_stopping, merge_frames, process_video, restore_audio, setup def process(start_time : float) -> ErrorCode: @@ -14,7 +16,7 @@ def process(start_time : float) -> ErrorCode: clear, setup, create_temp_frames, - process_frames, + process_video, merge_frames, restore_audio, partial(finalize_video, start_time), @@ -32,3 +34,20 @@ def process(start_time : float) -> ErrorCode: process_manager.end() return 0 + + +def create_temp_frames() -> ErrorCode: + state_manager.set_item('output_video_fps', 25.0) # TODO: set default fps value + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) + temp_image_resolution = restrict_image_resolution(state_manager.get_item('target_path'), output_image_resolution) + trim_frame_start, trim_frame_end = restrict_trim_audio_frame(source_audio_path, state_manager.get_item('output_video_fps'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + + if ffmpeg.spawn_frames(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_image_resolution, state_manager.get_item('output_video_fps'), trim_frame_start, trim_frame_end): + logger.debug(translator.get('spawning_frames_succeeded'), __name__) + else: + if is_process_stopping(): + return 4 + logger.error(translator.get('spawning_frames_failed'), __name__) + return 1 + return 0 diff --git a/facefusion/workflows/audio_to_image_as_frames.py b/facefusion/workflows/audio_to_image_as_frames.py deleted file mode 100644 index d2a1381b..00000000 --- a/facefusion/workflows/audio_to_image_as_frames.py +++ /dev/null @@ -1,32 +0,0 @@ -from functools import partial - -from facefusion import process_manager -from facefusion.types import ErrorCode -from facefusion.workflows.as_frames import copy_temp_frames, create_temp_frames, finalize_frames -from facefusion.workflows.core import analyse_image, clear, process_frames, setup - - -def process(start_time : float) -> ErrorCode: - tasks =\ - [ - analyse_image, - clear, - setup, - create_temp_frames, - process_frames, - copy_temp_frames, - partial(finalize_frames, start_time), - clear - ] - - process_manager.start() - - for task in tasks: - error_code = task() #type:ignore[operator] - - if error_code > 0: - process_manager.end() - return error_code - - process_manager.end() - return 0 diff --git a/facefusion/workflows/core.py b/facefusion/workflows/core.py index 1ba61053..13a58127 100644 --- a/facefusion/workflows/core.py +++ b/facefusion/workflows/core.py @@ -3,14 +3,16 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import numpy from tqdm import tqdm -from facefusion import content_analyser, logger, process_manager, state_manager, translator +from facefusion import content_analyser, ffmpeg, logger, process_manager, state_manager, translator, video_manager from facefusion.audio import create_empty_audio_frame, get_audio_frame, get_voice_frame from facefusion.common_helper import get_first -from facefusion.filesystem import filter_audio_paths +from facefusion.filesystem import filter_audio_paths, is_video +from facefusion.media_helper import restrict_trim_frame from facefusion.processors.core import get_processors_modules -from facefusion.temp_helper import clear_temp_directory, create_temp_directory, resolve_temp_frame_paths -from facefusion.types import AudioFrame, ErrorCode, VisionFrame -from facefusion.vision import conditional_merge_vision_mask, extract_vision_mask, read_static_image, read_static_images, read_static_video_frame, restrict_video_fps, write_image +from facefusion.temp_helper import clear_temp_directory, create_temp_directory, move_temp_file, resolve_temp_frame_paths +from facefusion.time_helper import calculate_end_time +from facefusion.types import AudioFrame, ErrorCode, Fps, Resolution, VisionFrame +from facefusion.vision import conditional_merge_vision_mask, detect_image_resolution, detect_video_resolution, extract_vision_mask, pack_resolution, read_static_image, read_static_images, read_static_video_frame, restrict_video_fps, scale_resolution, write_image def is_process_stopping() -> bool: @@ -39,7 +41,7 @@ def analyse_image() -> ErrorCode: def conditional_get_source_audio_frame(frame_number : int) -> AudioFrame: - if state_manager.get_item('workflow') in [ 'audio-to-image:frames', 'audio-to-image:video', 'image-to-video' ]: + if state_manager.get_item('workflow') in [ 'audio-to-image', 'image-to-video' ]: source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) output_video_fps = state_manager.get_item('output_video_fps') @@ -54,7 +56,7 @@ def conditional_get_source_audio_frame(frame_number : int) -> AudioFrame: def conditional_get_source_voice_frame(frame_number: int) -> AudioFrame: - if state_manager.get_item('workflow') in [ 'audio-to-image:frames', 'audio-to-image:video', 'image-to-video' ]: + if state_manager.get_item('workflow') in [ 'audio-to-image', 'image-to-video' ]: source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) output_video_fps = state_manager.get_item('output_video_fps') @@ -69,11 +71,28 @@ def conditional_get_source_voice_frame(frame_number: int) -> AudioFrame: def conditional_get_reference_vision_frame() -> VisionFrame: - if state_manager.get_item('workflow') in [ 'image-to-video', 'image-to-video:frames' ]: + if state_manager.get_item('workflow') == 'image-to-video': return read_static_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number')) return read_static_image(state_manager.get_item('target_path')) +def conditional_scale_resolution() -> Resolution: + if state_manager.get_item('workflow') == 'image-to-video': + return scale_resolution(detect_video_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) + return scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) + + +def conditional_restrict_video_fps() -> Fps: + if state_manager.get_item('workflow') == 'image-to-video': + return restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + return state_manager.get_item('output_video_fps') + + +def conditional_clear_video_pool() -> None: + if state_manager.get_item('workflow') == 'image-to-video': + video_manager.clear_video_pool() + + def process_temp_frame(temp_frame_path : str, frame_number : int) -> bool: reference_vision_frame = conditional_get_reference_vision_frame() source_vision_frames = read_static_images(state_manager.get_item('source_paths')) @@ -99,7 +118,7 @@ def process_temp_frame(temp_frame_path : str, frame_number : int) -> bool: return write_image(temp_frame_path, temp_vision_frame) -def process_frames() -> ErrorCode: +def process_video() -> ErrorCode: temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) if temp_frame_paths: @@ -131,3 +150,61 @@ def process_frames() -> ErrorCode: logger.error(translator.get('temp_frames_not_found'), __name__) return 1 return 0 + + +def merge_frames() -> ErrorCode: + temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) + trim_frame_start, trim_frame_end = restrict_trim_frame(len(temp_frame_paths), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + output_video_resolution = conditional_scale_resolution() + temp_video_fps = conditional_restrict_video_fps() + + logger.info(translator.get('merging_video').format(resolution = pack_resolution(output_video_resolution), fps = state_manager.get_item('output_video_fps')), __name__) + if ffmpeg.merge_video(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_video_fps, output_video_resolution, trim_frame_start, trim_frame_end): + logger.debug(translator.get('merging_video_succeeded'), __name__) + else: + if is_process_stopping(): + return 4 + logger.error(translator.get('merging_video_failed'), __name__) + return 1 + return 0 + + +def restore_audio() -> ErrorCode: + temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) + trim_frame_start, trim_frame_end = restrict_trim_frame(len(temp_frame_paths), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + + if state_manager.get_item('output_audio_volume') == 0: + logger.info(translator.get('skipping_audio'), __name__) + move_temp_file(state_manager.get_temp_path(), state_manager.get_item('output_path')) + else: + source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) + if source_audio_path: + if ffmpeg.replace_audio(source_audio_path, state_manager.get_item('output_path')): + conditional_clear_video_pool() + logger.debug(translator.get('replacing_audio_succeeded'), __name__) + else: + conditional_clear_video_pool() + if is_process_stopping(): + return 4 + logger.warn(translator.get('replacing_audio_skipped'), __name__) + move_temp_file(state_manager.get_temp_path(), state_manager.get_item('output_path')) + else: + if ffmpeg.restore_audio(state_manager.get_item('target_path'), state_manager.get_item('output_path'), trim_frame_start, trim_frame_end): + conditional_clear_video_pool() + logger.debug(translator.get('restoring_audio_succeeded'), __name__) + else: + conditional_clear_video_pool() + if is_process_stopping(): + return 4 + logger.warn(translator.get('restoring_audio_skipped'), __name__) + move_temp_file(state_manager.get_temp_path(), state_manager.get_item('output_path')) + return 0 + + +def finalize_video(start_time : float) -> ErrorCode: + if is_video(state_manager.get_item('output_path')): + logger.info(translator.get('processing_video_succeeded').format(seconds = calculate_end_time(start_time)), __name__) + else: + logger.error(translator.get('processing_video_failed'), __name__) + return 1 + return 0 diff --git a/facefusion/workflows/image_to_image.py b/facefusion/workflows/image_to_image.py index f24cf8e9..0837e3de 100644 --- a/facefusion/workflows/image_to_image.py +++ b/facefusion/workflows/image_to_image.py @@ -1,9 +1,12 @@ from functools import partial -from facefusion import process_manager +from facefusion import ffmpeg, logger, process_manager, state_manager, translator +from facefusion.filesystem import is_image +from facefusion.temp_helper import get_temp_file_path +from facefusion.time_helper import calculate_end_time from facefusion.types import ErrorCode -from facefusion.workflows.core import analyse_image, clear, setup -from facefusion.workflows.to_image import finalize_image, prepare_image, process_image +from facefusion.vision import detect_image_resolution, pack_resolution, restrict_image_resolution, scale_resolution +from facefusion.workflows.core import analyse_image, clear, is_process_stopping, process_temp_frame, setup def process(start_time : float) -> ErrorCode: @@ -29,3 +32,43 @@ def process(start_time : float) -> ErrorCode: process_manager.end() return 0 + + +def prepare_image() -> ErrorCode: + output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) + temp_image_resolution = restrict_image_resolution(state_manager.get_item('target_path'), output_image_resolution) + + logger.info(translator.get('copying_image').format(resolution = pack_resolution(temp_image_resolution)), __name__) + if ffmpeg.copy_image(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_image_resolution): + logger.debug(translator.get('copying_image_succeeded'), __name__) + else: + logger.error(translator.get('copying_image_failed'), __name__) + process_manager.end() + return 1 + return 0 + + +def process_image() -> ErrorCode: + temp_image_path = get_temp_file_path(state_manager.get_temp_path(), state_manager.get_item('output_path')) + process_temp_frame(temp_image_path, 0) + + if is_process_stopping(): + return 4 + return 0 + + +def finalize_image(start_time : float) -> ErrorCode: + output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) + + logger.info(translator.get('finalizing_image').format(resolution = pack_resolution(output_image_resolution)), __name__) + if ffmpeg.finalize_image(state_manager.get_item('output_path'), output_image_resolution): + logger.debug(translator.get('finalizing_image_succeeded'), __name__) + else: + logger.warn(translator.get('finalizing_image_skipped'), __name__) + + if is_image(state_manager.get_item('output_path')): + logger.info(translator.get('processing_image_succeeded').format(seconds = calculate_end_time(start_time)), __name__) + else: + logger.error(translator.get('processing_image_failed'), __name__) + return 1 + return 0 diff --git a/facefusion/workflows/image_to_video.py b/facefusion/workflows/image_to_video.py index 359be3a8..4b982541 100644 --- a/facefusion/workflows/image_to_video.py +++ b/facefusion/workflows/image_to_video.py @@ -1,9 +1,9 @@ from functools import partial -from facefusion import process_manager +from facefusion import content_analyser, ffmpeg, logger, process_manager, state_manager, translator from facefusion.types import ErrorCode -from facefusion.workflows.core import clear, process_frames, setup -from facefusion.workflows.to_video import analyse_video, create_temp_frames, finalize_video, merge_frames, restore_audio +from facefusion.vision import detect_video_resolution, pack_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution +from facefusion.workflows.core import clear, finalize_video, is_process_stopping, merge_frames, process_video, restore_audio, setup def process(start_time : float) -> ErrorCode: @@ -13,7 +13,7 @@ def process(start_time : float) -> ErrorCode: clear, setup, create_temp_frames, - process_frames, + process_video, merge_frames, restore_audio, partial(finalize_video, start_time), @@ -31,3 +31,28 @@ def process(start_time : float) -> ErrorCode: process_manager.end() return 0 + + +def analyse_video() -> ErrorCode: + trim_frame_start, trim_frame_end = restrict_trim_video_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) + + if content_analyser.analyse_video(state_manager.get_item('target_path'), trim_frame_start, trim_frame_end): + return 3 + return 0 + + +def create_temp_frames() -> ErrorCode: + trim_frame_start, trim_frame_end = restrict_trim_video_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) + temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) + logger.info(translator.get('extracting_frames').format(resolution=pack_resolution(temp_video_resolution), fps=temp_video_fps), __name__) + + if ffmpeg.extract_frames(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_video_resolution, temp_video_fps, trim_frame_start, trim_frame_end): + logger.debug(translator.get('extracting_frames_succeeded'), __name__) + else: + if is_process_stopping(): + return 4 + logger.error(translator.get('extracting_frames_failed'), __name__) + return 1 + return 0 diff --git a/facefusion/workflows/image_to_video_as_frames.py b/facefusion/workflows/image_to_video_as_frames.py deleted file mode 100644 index 70729d7d..00000000 --- a/facefusion/workflows/image_to_video_as_frames.py +++ /dev/null @@ -1,33 +0,0 @@ -from functools import partial - -from facefusion import process_manager -from facefusion.types import ErrorCode -from facefusion.workflows.as_frames import copy_temp_frames, finalize_frames -from facefusion.workflows.core import clear, process_frames, setup -from facefusion.workflows.to_video import analyse_video, create_temp_frames - - -def process(start_time : float) -> ErrorCode: - tasks =\ - [ - analyse_video, - clear, - setup, - create_temp_frames, - process_frames, - copy_temp_frames, - partial(finalize_frames, start_time), - clear - ] - - process_manager.start() - - for task in tasks: - error_code = task() #type:ignore[operator] - - if error_code > 0: - process_manager.end() - return error_code - - process_manager.end() - return 0 diff --git a/facefusion/workflows/to_image.py b/facefusion/workflows/to_image.py deleted file mode 100644 index dbd1adb2..00000000 --- a/facefusion/workflows/to_image.py +++ /dev/null @@ -1,47 +0,0 @@ -from facefusion import ffmpeg, logger, process_manager, state_manager, translator -from facefusion.filesystem import is_image -from facefusion.temp_helper import get_temp_file_path -from facefusion.time_helper import calculate_end_time -from facefusion.types import ErrorCode -from facefusion.vision import detect_image_resolution, pack_resolution, restrict_image_resolution, scale_resolution -from facefusion.workflows.core import is_process_stopping, process_temp_frame - - -def prepare_image() -> ErrorCode: - output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) - temp_image_resolution = restrict_image_resolution(state_manager.get_item('target_path'), output_image_resolution) - - logger.info(translator.get('copying_image').format(resolution = pack_resolution(temp_image_resolution)), __name__) - if ffmpeg.copy_image(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_image_resolution): - logger.debug(translator.get('copying_image_succeeded'), __name__) - else: - logger.error(translator.get('copying_image_failed'), __name__) - process_manager.end() - return 1 - return 0 - - -def process_image() -> ErrorCode: - temp_image_path = get_temp_file_path(state_manager.get_temp_path(), state_manager.get_item('output_path')) - process_temp_frame(temp_image_path, 0) - - if is_process_stopping(): - return 4 - return 0 - - -def finalize_image(start_time : float) -> ErrorCode: - output_image_resolution = scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_image_scale')) - - logger.info(translator.get('finalizing_image').format(resolution = pack_resolution(output_image_resolution)), __name__) - if ffmpeg.finalize_image(state_manager.get_item('output_path'), output_image_resolution): - logger.debug(translator.get('finalizing_image_succeeded'), __name__) - else: - logger.warn(translator.get('finalizing_image_skipped'), __name__) - - if is_image(state_manager.get_item('output_path')): - logger.info(translator.get('processing_image_succeeded').format(seconds = calculate_end_time(start_time)), __name__) - else: - logger.error(translator.get('processing_image_failed'), __name__) - return 1 - return 0 diff --git a/facefusion/workflows/to_video.py b/facefusion/workflows/to_video.py deleted file mode 100644 index 43366dd8..00000000 --- a/facefusion/workflows/to_video.py +++ /dev/null @@ -1,109 +0,0 @@ -from facefusion import content_analyser, ffmpeg, logger, state_manager, translator, video_manager -from facefusion.common_helper import get_first -from facefusion.filesystem import filter_audio_paths, is_video -from facefusion.media_helper import restrict_trim_frame -from facefusion.temp_helper import move_temp_file, resolve_temp_frame_paths -from facefusion.time_helper import calculate_end_time -from facefusion.types import ErrorCode, Fps, Resolution -from facefusion.vision import detect_image_resolution, detect_video_resolution, pack_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution -from facefusion.workflows.core import is_process_stopping - - -def analyse_video() -> ErrorCode: - trim_frame_start, trim_frame_end = restrict_trim_video_frame(state_manager.get_item('target_path'), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) - - if content_analyser.analyse_video(state_manager.get_item('target_path'), trim_frame_start, trim_frame_end): - return 3 - return 0 - - -def create_temp_frames() -> ErrorCode: - trim_frame_start, trim_frame_end = restrict_trim_video_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) - temp_video_fps = restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) - logger.info(translator.get('extracting_frames').format(resolution=pack_resolution(temp_video_resolution), fps=temp_video_fps), __name__) - - if ffmpeg.extract_frames(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_video_resolution, temp_video_fps, trim_frame_start, trim_frame_end): - logger.debug(translator.get('extracting_frames_succeeded'), __name__) - else: - if is_process_stopping(): - return 4 - logger.error(translator.get('extracting_frames_failed'), __name__) - return 1 - return 0 - - -def merge_frames() -> ErrorCode: - temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) - trim_frame_start, trim_frame_end = restrict_trim_frame(len(temp_frame_paths), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) - output_video_resolution = conditional_scale_resolution() - temp_video_fps = conditional_restrict_video_fps() - - logger.info(translator.get('merging_video').format(resolution = pack_resolution(output_video_resolution), fps = state_manager.get_item('output_video_fps')), __name__) - if ffmpeg.merge_video(state_manager.get_item('target_path'), state_manager.get_item('output_path'), temp_video_fps, output_video_resolution, trim_frame_start, trim_frame_end): - logger.debug(translator.get('merging_video_succeeded'), __name__) - else: - if is_process_stopping(): - return 4 - logger.error(translator.get('merging_video_failed'), __name__) - return 1 - return 0 - - -def restore_audio() -> ErrorCode: - temp_frame_paths = resolve_temp_frame_paths(state_manager.get_temp_path(), state_manager.get_item('output_path'), state_manager.get_item('temp_frame_format')) - trim_frame_start, trim_frame_end = restrict_trim_frame(len(temp_frame_paths), state_manager.get_item('trim_frame_start'), state_manager.get_item('trim_frame_end')) - - if state_manager.get_item('output_audio_volume') == 0: - logger.info(translator.get('skipping_audio'), __name__) - move_temp_file(state_manager.get_temp_path(), state_manager.get_item('output_path')) - else: - source_audio_path = get_first(filter_audio_paths(state_manager.get_item('source_paths'))) - if source_audio_path: - if ffmpeg.replace_audio(source_audio_path, state_manager.get_item('output_path')): - conditional_clear_video_pool() - logger.debug(translator.get('replacing_audio_succeeded'), __name__) - else: - conditional_clear_video_pool() - if is_process_stopping(): - return 4 - logger.warn(translator.get('replacing_audio_skipped'), __name__) - move_temp_file(state_manager.get_temp_path(), state_manager.get_item('output_path')) - else: - if ffmpeg.restore_audio(state_manager.get_item('target_path'), state_manager.get_item('output_path'), trim_frame_start, trim_frame_end): - conditional_clear_video_pool() - logger.debug(translator.get('restoring_audio_succeeded'), __name__) - else: - conditional_clear_video_pool() - if is_process_stopping(): - return 4 - logger.warn(translator.get('restoring_audio_skipped'), __name__) - move_temp_file(state_manager.get_temp_path(), state_manager.get_item('output_path')) - return 0 - - -def finalize_video(start_time : float) -> ErrorCode: - if is_video(state_manager.get_item('output_path')): - logger.info(translator.get('processing_video_succeeded').format(seconds = calculate_end_time(start_time)), __name__) - else: - logger.error(translator.get('processing_video_failed'), __name__) - return 1 - return 0 - - -def conditional_clear_video_pool() -> None: - if state_manager.get_item('workflow') == 'image-to-video': - video_manager.clear_video_pool() - - -def conditional_restrict_video_fps() -> Fps: - if state_manager.get_item('workflow') == 'image-to-video': - return restrict_video_fps(state_manager.get_item('target_path'), state_manager.get_item('output_video_fps')) - return state_manager.get_item('output_video_fps') - - -def conditional_scale_resolution() -> Resolution: - if state_manager.get_item('workflow') == 'image-to-video': - return scale_resolution(detect_video_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) - return scale_resolution(detect_image_resolution(state_manager.get_item('target_path')), state_manager.get_item('output_video_scale')) diff --git a/requirements.txt b/requirements.txt index 5981d96c..41439f34 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,14 @@ -numpy==2.2.6 +aiortc +numpy==2.3.4 onnx==1.19.1 onnxruntime==1.23.2 opencv-python==4.12.0.88 -psutil==7.1.3 +psutil==7.1.2 +python-multipart tqdm==4.67.1 scipy==1.16.3 starlette==0.50.0 uvicorn==0.34.0 websockets==15.0.1 +yt-dlp==2025.12.8 +gallery-dl==1.30.10 diff --git a/tests/helper.py b/tests/helper.py index 82c07a7e..8902643c 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -1,7 +1,7 @@ import os import tempfile -from facefusion.filesystem import are_images, create_directory, is_directory, is_file, remove_directory, resolve_file_paths +from facefusion.filesystem import create_directory, is_directory, is_file, remove_directory from facefusion.types import JobStatus @@ -26,14 +26,10 @@ def get_test_examples_directory() -> str: def is_test_output_file(file_path : str) -> bool: - return is_file(get_test_output_path(file_path)) + return is_file(get_test_output_file(file_path)) -def is_test_output_sequence(directory_path : str) -> bool: - return are_images(resolve_file_paths(directory_path)) - - -def get_test_output_path(file_path : str) -> str: +def get_test_output_file(file_path : str) -> str: return os.path.join(get_test_outputs_directory(), file_path) diff --git a/tests/test_api_assets.py b/tests/test_api_assets.py new file mode 100644 index 00000000..1133b43b --- /dev/null +++ b/tests/test_api_assets.py @@ -0,0 +1,303 @@ +import io +from typing import Iterator + +import pytest +from starlette.testclient import TestClient + +from facefusion import metadata, session_manager, state_manager +from facefusion.apis.core import create_api + + +@pytest.fixture(scope = 'module') +def test_client() -> Iterator[TestClient]: + with TestClient(create_api()) as test_client: + yield test_client + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + session_manager.SESSIONS.clear() + state_manager.clear_item('asset_registry') + + +@pytest.fixture(scope = 'function') +def auth_token(test_client : TestClient) -> str: + create_session_response = test_client.post('/session', json = + { + 'client_version': metadata.get('version') + }) + return create_session_response.json().get('access_token') + + +def test_upload_source_single(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + test_image.name = 'test_face.jpg' + + response = test_client.post('/assets?type=source', + files = {'file': ('test_face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 201 + body = response.json() + assert body.get('message') == '1 source(s) uploaded successfully' + assert len(body.get('asset_ids')) == 1 + assert isinstance(body.get('asset_ids')[0], str) + + +def test_upload_source_multiple(test_client : TestClient, auth_token : str) -> None: + test_image1 = io.BytesIO(b'fake image data 1') + test_image1.name = 'face1.jpg' + test_image2 = io.BytesIO(b'fake image data 2') + test_image2.name = 'face2.jpg' + + response = test_client.post('/assets?type=source', + files = [ + ('file', ('face1.jpg', test_image1, 'image/jpeg')), + ('file', ('face2.jpg', test_image2, 'image/jpeg')) + ], + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 201 + body = response.json() + assert body.get('message') == '2 source(s) uploaded successfully' + assert len(body.get('asset_ids')) == 2 + + +def test_upload_target_image(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + test_image.name = 'target.jpg' + + response = test_client.post('/assets?type=target', + files = {'file': ('target.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 201 + body = response.json() + assert body.get('message') == 'Target uploaded successfully' + assert isinstance(body.get('asset_id'), str) + + +def test_upload_missing_type_param(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + + response = test_client.post('/assets', + files = {'file': ('test.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 400 + assert response.json().get('message') == 'Missing required query parameter: type' + + +def test_upload_invalid_type_param(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + + response = test_client.post('/assets?type=invalid', + files = {'file': ('test.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 400 + assert response.json().get('message') == 'Invalid type. Must be "source" or "target"' + + +def test_upload_no_file(test_client : TestClient, auth_token : str) -> None: + response = test_client.post('/assets?type=source', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 400 + assert response.json().get('message') == 'No file provided' + + +def test_list_assets_empty(test_client : TestClient, auth_token : str) -> None: + response = test_client.get('/assets', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 200 + body = response.json() + assert body.get('count') == 0 + assert body.get('assets') == [] + + +def test_list_assets_with_uploads(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + test_client.post('/assets?type=source', + files = {'file': ('face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + test_image2 = io.BytesIO(b'fake target data') + test_client.post('/assets?type=target', + files = {'file': ('target.jpg', test_image2, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + response = test_client.get('/assets', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 200 + body = response.json() + assert body.get('count') == 2 + assets = body.get('assets') + assert len(assets) == 2 + + +def test_list_assets_filter_by_type(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + test_client.post('/assets?type=source', + files = {'file': ('face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + test_image2 = io.BytesIO(b'fake target data') + test_client.post('/assets?type=target', + files = {'file': ('target.jpg', test_image2, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + response = test_client.get('/assets?type=source', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 200 + body = response.json() + assert body.get('count') == 1 + assets = body.get('assets') + assert assets[0].get('type') == 'source' + + +def test_get_asset_metadata(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + upload_response = test_client.post('/assets?type=source', + files = {'file': ('face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + asset_id = upload_response.json().get('asset_ids')[0] + + response = test_client.get(f'/assets/{asset_id}', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 200 + body = response.json() + assert body.get('id') == asset_id + assert body.get('type') == 'source' + assert body.get('filename') == 'face.jpg' + assert body.get('size') > 0 + assert body.get('created_at') + assert 'path' not in body + + +def test_get_asset_not_found(test_client : TestClient, auth_token : str) -> None: + response = test_client.get('/assets/non-existent-id', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 404 + assert response.json().get('message') == 'Asset not found' + + +def test_download_asset(test_client : TestClient, auth_token : str) -> None: + test_data = b'fake image data for download' + test_image = io.BytesIO(test_data) + upload_response = test_client.post('/assets?type=source', + files = {'file': ('face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + asset_id = upload_response.json().get('asset_ids')[0] + + response = test_client.get(f'/assets/{asset_id}?action=download', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 200 + assert response.content == test_data + + +def test_download_asset_not_found(test_client : TestClient, auth_token : str) -> None: + response = test_client.get('/assets/non-existent-id?action=download', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 404 + assert response.json().get('message') == 'Asset not found' + + +def test_delete_asset(test_client : TestClient, auth_token : str) -> None: + test_image = io.BytesIO(b'fake image data') + upload_response = test_client.post('/assets?type=source', + files = {'file': ('face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {auth_token}'} + ) + asset_id = upload_response.json().get('asset_ids')[0] + + response = test_client.delete(f'/assets/{asset_id}', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 200 + assert response.json().get('message') == 'Asset deleted successfully' + + get_response = test_client.get(f'/assets/{asset_id}', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + assert get_response.status_code == 404 + + +def test_delete_asset_not_found(test_client : TestClient, auth_token : str) -> None: + response = test_client.delete('/assets/non-existent-id', + headers = {'Authorization': f'Bearer {auth_token}'} + ) + + assert response.status_code == 404 + assert response.json().get('message') == 'Asset not found' + + +def test_assets_require_auth(test_client : TestClient) -> None: + response = test_client.get('/assets') + assert response.status_code == 401 + + response = test_client.post('/assets?type=source') + assert response.status_code == 401 + + response = test_client.get('/assets/some-id') + assert response.status_code == 401 + + response = test_client.delete('/assets/some-id') + assert response.status_code == 401 + + +def test_assets_session_isolation(test_client : TestClient) -> None: + session1_response = test_client.post('/session', json = {'client_version': metadata.get('version')}) + session1_token = session1_response.json().get('access_token') + + session2_response = test_client.post('/session', json = {'client_version': metadata.get('version')}) + session2_token = session2_response.json().get('access_token') + + test_image = io.BytesIO(b'session 1 data') + upload_response = test_client.post('/assets?type=source', + files = {'file': ('face.jpg', test_image, 'image/jpeg')}, + headers = {'Authorization': f'Bearer {session1_token}'} + ) + asset_id = upload_response.json().get('asset_ids')[0] + + response = test_client.get('/assets', + headers = {'Authorization': f'Bearer {session1_token}'} + ) + assert response.json().get('count') == 1 + + response = test_client.get('/assets', + headers = {'Authorization': f'Bearer {session2_token}'} + ) + assert response.json().get('count') == 0 + + response = test_client.get(f'/assets/{asset_id}', + headers = {'Authorization': f'Bearer {session2_token}'} + ) + assert response.status_code == 404 diff --git a/tests/test_asset_store.py b/tests/test_asset_store.py new file mode 100644 index 00000000..60a9259b --- /dev/null +++ b/tests/test_asset_store.py @@ -0,0 +1,247 @@ +import os +import tempfile +from typing import Iterator + +import pytest + +from facefusion import asset_store, session_manager, state_manager +from facefusion.session_context import clear_session_id, set_session_id + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + session_manager.SESSIONS.clear() + state_manager.clear_item('asset_registry') + clear_session_id() + + +@pytest.fixture(scope = 'function') +def temp_file() -> Iterator[str]: + fd, path = tempfile.mkstemp(suffix = '.jpg') + os.write(fd, b'test file content') + os.close(fd) + yield path + if os.path.exists(path): + os.remove(path) + + +@pytest.fixture(scope = 'function') +def session_id() -> str: + test_session_id = 'test-session-123' + set_session_id(test_session_id) + return test_session_id + + +def test_register_source_asset(temp_file : str, session_id : str) -> None: + asset_id = asset_store.register('source', temp_file, 'test.jpg') + + assert isinstance(asset_id, str) + assert len(asset_id) == 36 + + asset = asset_store.get_asset(asset_id) + assert asset is not None + assert asset.get('id') == asset_id + assert asset.get('session_id') == session_id + assert asset.get('type') == 'source' + assert asset.get('path') == temp_file + assert asset.get('filename') == 'test.jpg' + assert asset.get('size') > 0 + assert asset.get('created_at') + + +def test_register_target_asset(temp_file : str, session_id : str) -> None: + asset_id = asset_store.register('target', temp_file, 'video.mp4') + + asset = asset_store.get_asset(asset_id) + assert asset.get('type') == 'target' + assert asset.get('filename') == 'video.mp4' + + +def test_register_output_asset(temp_file : str, session_id : str) -> None: + metadata = {'fps': 30, 'resolution': [1920, 1080]} + asset_id = asset_store.register('output', temp_file, 'output.mp4', metadata) + + asset = asset_store.get_asset(asset_id) + assert asset.get('type') == 'output' + assert asset.get('metadata') == metadata + + +def test_register_invalid_type(temp_file : str, session_id : str) -> None: + with pytest.raises(ValueError) as exc: + asset_store.register('invalid_type', temp_file, 'test.jpg') + assert "Invalid asset_type" in str(exc.value) + + +def test_register_without_session() -> None: + fd, path = tempfile.mkstemp() + os.close(fd) + + try: + with pytest.raises(ValueError) as exc: + asset_store.register('source', path, 'test.jpg') + assert "No active session" in str(exc.value) + finally: + os.remove(path) + + +def test_register_without_filename(temp_file : str, session_id : str) -> None: + asset_id = asset_store.register('source', temp_file) + + asset = asset_store.get_asset(asset_id) + assert asset.get('filename') == os.path.basename(temp_file) + + +def test_get_asset_not_found(session_id : str) -> None: + asset = asset_store.get_asset('non-existent-id') + assert asset is None + + +def test_list_assets_empty(session_id : str) -> None: + assets = asset_store.list_assets() + assert assets == [] + + +def test_list_assets_with_multiple(temp_file : str, session_id : str) -> None: + fd1, path1 = tempfile.mkstemp(suffix = '.jpg') + os.write(fd1, b'content 1') + os.close(fd1) + + fd2, path2 = tempfile.mkstemp(suffix = '.mp4') + os.write(fd2, b'content 2') + os.close(fd2) + + try: + asset_store.register('source', path1, 'source1.jpg') + asset_store.register('source', path2, 'source2.jpg') + asset_store.register('target', temp_file, 'target.mp4') + + assets = asset_store.list_assets() + assert len(assets) == 3 + finally: + os.remove(path1) + os.remove(path2) + + +def test_list_assets_filter_by_type(temp_file : str, session_id : str) -> None: + fd, path = tempfile.mkstemp(suffix = '.jpg') + os.write(fd, b'content') + os.close(fd) + + try: + asset_store.register('source', path, 'source.jpg') + asset_store.register('target', temp_file, 'target.mp4') + + source_assets = asset_store.list_assets('source') + assert len(source_assets) == 1 + assert source_assets[0].get('type') == 'source' + + target_assets = asset_store.list_assets('target') + assert len(target_assets) == 1 + assert target_assets[0].get('type') == 'target' + + output_assets = asset_store.list_assets('output') + assert len(output_assets) == 0 + finally: + os.remove(path) + + +def test_list_assets_invalid_type(session_id : str) -> None: + with pytest.raises(ValueError) as exc: + asset_store.list_assets('invalid_type') + assert "Invalid asset_type" in str(exc.value) + + +def test_list_assets_session_scoped(temp_file : str) -> None: + session1_id = 'session-1' + set_session_id(session1_id) + asset1_id = asset_store.register('source', temp_file, 'file1.jpg') + + session2_id = 'session-2' + set_session_id(session2_id) + + fd, path2 = tempfile.mkstemp(suffix = '.jpg') + os.write(fd, b'content 2') + os.close(fd) + + try: + asset2_id = asset_store.register('source', path2, 'file2.jpg') + + assets_session2 = asset_store.list_assets() + assert len(assets_session2) == 1 + assert assets_session2[0].get('id') == asset2_id + + set_session_id(session1_id) + assets_session1 = asset_store.list_assets() + assert len(assets_session1) == 1 + assert assets_session1[0].get('id') == asset1_id + finally: + os.remove(path2) + + +def test_delete_asset(temp_file : str, session_id : str) -> None: + asset_id = asset_store.register('source', temp_file, 'test.jpg') + + assert os.path.exists(temp_file) + + success = asset_store.delete_asset(asset_id) + assert success is True + + assert not os.path.exists(temp_file) + + asset = asset_store.get_asset(asset_id) + assert asset is None + + +def test_delete_asset_not_found(session_id : str) -> None: + success = asset_store.delete_asset('non-existent-id') + assert success is False + + +def test_cleanup_session_assets(session_id : str) -> None: + fd1, path1 = tempfile.mkstemp(suffix = '.jpg') + os.write(fd1, b'content 1') + os.close(fd1) + + fd2, path2 = tempfile.mkstemp(suffix = '.mp4') + os.write(fd2, b'content 2') + os.close(fd2) + + asset_id1 = asset_store.register('source', path1, 'source.jpg') + asset_id2 = asset_store.register('target', path2, 'target.mp4') + + assert os.path.exists(path1) + assert os.path.exists(path2) + + asset_store.cleanup_session_assets(session_id) + + assert not os.path.exists(path1) + assert not os.path.exists(path2) + + assert asset_store.get_asset(asset_id1) is None + assert asset_store.get_asset(asset_id2) is None + + +def test_cleanup_session_assets_only_affects_target_session(temp_file : str) -> None: + session1_id = 'session-1' + set_session_id(session1_id) + + fd, path1 = tempfile.mkstemp(suffix = '.jpg') + os.write(fd, b'content 1') + os.close(fd) + + asset1_id = asset_store.register('source', path1, 'file1.jpg') + + session2_id = 'session-2' + set_session_id(session2_id) + asset2_id = asset_store.register('source', temp_file, 'file2.jpg') + + asset_store.cleanup_session_assets(session1_id) + + assert not os.path.exists(path1) + assert os.path.exists(temp_file) + + set_session_id(session1_id) + assert asset_store.get_asset(asset1_id) is None + + set_session_id(session2_id) + assert asset_store.get_asset(asset2_id) is not None diff --git a/tests/test_cli_age_modifier.py b/tests/test_cli_age_modifier.py index 1af9675d..06dba143 100644 --- a/tests/test_cli_age_modifier.py +++ b/tests/test_cli_age_modifier.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -25,21 +25,14 @@ def before_each() -> None: def test_modify_age_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-age-face-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-age-face-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-age-face-to-image.jpg') is True def test_modify_age_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-age-face-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-age-face-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-age-face-to-video.mp4') is True - - -def test_modify_age_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'age_modifier', '--age-modifier-direction', '100', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-age-face-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-age-face-to-video-as-frames')) is True diff --git a/tests/test_cli_background_remover.py b/tests/test_cli_background_remover.py index cbd7e026..2646a178 100644 --- a/tests/test_cli_background_remover.py +++ b/tests/test_cli_background_remover.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -15,7 +15,7 @@ def before_all() -> None: 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg', 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' ]) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg') ]) + subprocess.run(['ffmpeg', '-i', get_test_example_file('target-240p.mp4'), '-vframes', '1', get_test_example_file('target-240p.jpg')]) @pytest.fixture(scope = 'function', autouse = True) @@ -26,21 +26,14 @@ def before_each() -> None: def test_remove_background_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-remove-background-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-remove-background-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-remove-background-to-image.jpg') is True def test_remove_background_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-remove-background-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-remove-background-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-remove-background-to-video.mp4') is True - - -def test_remove_background_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'background_remover', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-remove-background-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-remove-background-to-video-as-frames')) is True diff --git a/tests/test_cli_batch_runner.py b/tests/test_cli_batch_runner.py index adafde35..4c124c9f 100644 --- a/tests/test_cli_batch_runner.py +++ b/tests/test_cli_batch_runner.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -26,7 +26,7 @@ def before_each() -> None: def test_batch_run_targets() -> None: - commands = [ sys.executable, 'facefusion.py', 'batch-run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p-batch-*.jpg'), '-o', get_test_output_path('test-batch-run-targets-{index}.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'batch-run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p-batch-*.jpg'), '-o', get_test_output_file('test-batch-run-targets-{index}.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-batch-run-targets-0.jpg') is True @@ -35,7 +35,7 @@ def test_batch_run_targets() -> None: def test_batch_run_sources_to_targets() -> None: - commands = [ sys.executable, 'facefusion.py', 'batch-run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('target-240p-batch-*.jpg'), '-t', get_test_example_file('target-240p-batch-*.jpg'), '-o', get_test_output_path('test-batch-run-sources-to-targets-{index}.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'batch-run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('target-240p-batch-*.jpg'), '-t', get_test_example_file('target-240p-batch-*.jpg'), '-o', get_test_output_file('test-batch-run-sources-to-targets-{index}.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-batch-run-sources-to-targets-0.jpg') is True diff --git a/tests/test_cli_expression_restorer.py b/tests/test_cli_expression_restorer.py index a362010a..9b4aa84c 100644 --- a/tests/test_cli_expression_restorer.py +++ b/tests/test_cli_expression_restorer.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -25,21 +25,14 @@ def before_each() -> None: def test_restore_expression_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-restore-expression-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-restore-expression-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-restore-expression-to-image.jpg') is True def test_restore_expression_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-restore-expression-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-restore-expression-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-restore-expression-to-video.mp4') is True - - -def test_restore_expression_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'expression_restorer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-restore-expression-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-restore-expression-to-video-as-frames')) is True diff --git a/tests/test_cli_face_debugger.py b/tests/test_cli_face_debugger.py index cf42a08c..2669a276 100644 --- a/tests/test_cli_face_debugger.py +++ b/tests/test_cli_face_debugger.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -26,21 +26,14 @@ def before_each() -> None: def test_debug_face_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-debug-face-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-debug-face-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-debug-face-to-image.jpg') is True def test_debug_face_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-debug-face-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-debug-face-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-debug-face-to-video.mp4') is True - - -def test_debug_face_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-debug-face-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-debug-face-to-video-as-frames')) is True diff --git a/tests/test_cli_face_editor.py b/tests/test_cli_face_editor.py index 90535f77..8c72c5bf 100644 --- a/tests/test_cli_face_editor.py +++ b/tests/test_cli_face_editor.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -26,21 +26,14 @@ def before_each() -> None: def test_edit_face_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-edit-face-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-edit-face-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-edit-face-to-image.jpg') is True def test_edit_face_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-edit-face-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-edit-face-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-edit-face-to-video.mp4') is True - - -def test_edit_face_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_editor', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-edit-face-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-edit-face-to-video-as-frames')) is True diff --git a/tests/test_cli_face_enhancer.py b/tests/test_cli_face_enhancer.py index e42935c4..b2e3a8ac 100644 --- a/tests/test_cli_face_enhancer.py +++ b/tests/test_cli_face_enhancer.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -26,21 +26,14 @@ def before_each() -> None: def test_enhance_face_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-enhance-face-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-enhance-face-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-enhance-face-to-image.jpg') is True def test_enhance_face_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-enhance-face-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-enhance-face-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-enhance-face-to-video.mp4') is True - - -def test_enhance_face_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-enhance-face-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-enhance-face-to-video-as-frames')) is True diff --git a/tests/test_cli_face_swapper.py b/tests/test_cli_face_swapper.py index 2537a956..f8f067cc 100644 --- a/tests/test_cli_face_swapper.py +++ b/tests/test_cli_face_swapper.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -26,21 +26,14 @@ def before_each() -> None: def test_swap_face_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-swap-face-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-swap-face-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-swap-face-to-image.jpg') is True def test_swap_face_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-swap-face-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-swap-face-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-swap-face-to-video.mp4') is True - - -def test_swap_face_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_swapper', '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-swap-face-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-swap-face-to-video-as-frames')) is True diff --git a/tests/test_cli_frame_colorizer.py b/tests/test_cli_frame_colorizer.py index ea93a7b6..a50e6a63 100644 --- a/tests/test_cli_frame_colorizer.py +++ b/tests/test_cli_frame_colorizer.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -27,21 +27,14 @@ def before_each() -> None: def test_colorize_frame_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.jpg'), '-o', get_test_output_path('test_colorize-frame-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.jpg'), '-o', get_test_output_file('test_colorize-frame-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test_colorize-frame-to-image.jpg') is True def test_colorize_frame_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.mp4'), '-o', get_test_output_path('test-colorize-frame-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.mp4'), '-o', get_test_output_file('test-colorize-frame-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-colorize-frame-to-video.mp4') is True - - -def test_colorize_frame_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_colorizer', '-t', get_test_example_file('target-240p-0sat.mp4'), '-o', get_test_output_path('test-colorize-frame-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-colorize-frame-to-video-as-frames')) is True diff --git a/tests/test_cli_frame_enhancer.py b/tests/test_cli_frame_enhancer.py index dba68704..d0b55f5d 100644 --- a/tests/test_cli_frame_enhancer.py +++ b/tests/test_cli_frame_enhancer.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -26,21 +26,14 @@ def before_each() -> None: def test_enhance_frame_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-enhance-frame-to-image.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-enhance-frame-to-image.jpg') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-enhance-frame-to-image.jpg') is True def test_enhance_frame_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-enhance-frame-to-video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-enhance-frame-to-video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test-enhance-frame-to-video.mp4') is True - - -def test_enhance_frame_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-enhance-frame-to-video-as-frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test-enhance-frame-to-video-as-frames')) is True diff --git a/tests/test_cli_job_manager.py b/tests/test_cli_job_manager.py index 7d462d55..9bd76d8f 100644 --- a/tests/test_cli_job_manager.py +++ b/tests/test_cli_job_manager.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, count_step_total, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_job_file +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_job_file @pytest.fixture(scope = 'module', autouse = True) @@ -50,7 +50,7 @@ def test_job_submit() -> None: assert subprocess.run(commands).returncode == 1 - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] subprocess.run(commands) commands = [ sys.executable, 'facefusion.py', 'job-submit', 'test-job-submit', '--jobs-path', get_test_jobs_directory() ] @@ -73,10 +73,10 @@ def test_submit_all() -> None: assert subprocess.run(commands).returncode == 1 - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit-all-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit-all-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit-all-2', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-submit-all-2', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] subprocess.run(commands) commands = [ sys.executable, 'facefusion.py', 'job-submit-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] @@ -122,7 +122,7 @@ def test_job_delete_all() -> None: def test_job_add_step() -> None: - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert subprocess.run(commands).returncode == 1 assert count_step_total('test-job-add-step') == 0 @@ -130,14 +130,14 @@ def test_job_add_step() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-add-step', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-add-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert subprocess.run(commands).returncode == 0 assert count_step_total('test-job-add-step') == 1 def test_job_remix() -> None: - commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert subprocess.run(commands).returncode == 1 assert count_step_total('test-job-remix-step') == 0 @@ -145,23 +145,23 @@ def test_job_remix() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-remix-step', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remix-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remix-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert count_step_total('test-job-remix-step') == 1 assert subprocess.run(commands).returncode == 0 assert count_step_total('test-job-remix-step') == 2 - commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-remix-step', 'test-job-remix-step', '-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert subprocess.run(commands).returncode == 0 assert count_step_total('test-job-remix-step') == 3 def test_job_insert_step() -> None: - commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert subprocess.run(commands).returncode == 1 assert count_step_total('test-job-insert-step') == 0 @@ -169,16 +169,16 @@ def test_job_insert_step() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-insert-step', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-insert-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-insert-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '0', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert count_step_total('test-job-insert-step') == 1 assert subprocess.run(commands).returncode == 0 assert count_step_total('test-job-insert-step') == 2 - commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-insert-step', 'test-job-insert-step', '-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] assert subprocess.run(commands).returncode == 0 assert count_step_total('test-job-insert-step') == 3 @@ -192,7 +192,7 @@ def test_job_remove_step() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-remove-step', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remove-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-remix-step.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-remove-step', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '-s', get_test_example_file('source.jpg'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-remix-step.jpg') ] subprocess.run(commands) subprocess.run(commands) diff --git a/tests/test_cli_job_runner.py b/tests/test_cli_job_runner.py index e1608ed5..d84aeaf9 100644 --- a/tests/test_cli_job_runner.py +++ b/tests/test_cli_job_runner.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs, move_job_file, set_steps_status -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -33,7 +33,7 @@ def test_job_run() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-run.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-run.jpg') ] subprocess.run(commands) commands = [ sys.executable, 'facefusion.py', 'job-run', 'test-job-run', '--jobs-path', get_test_jobs_directory() ] @@ -61,13 +61,13 @@ def test_job_run_all() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-run-all-2', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-run-all-1.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-run-all-1.jpg') ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-job-run-all-2.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-run-all-2.mp4'), '--trim-frame-end', '1' ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-job-run-all-2.mp4'), '--trim-frame-start', '0', '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-run-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-run-all-2.mp4'), '--trim-frame-start', '0', '--trim-frame-end', '1' ] subprocess.run(commands) commands = [ sys.executable, 'facefusion.py', 'job-run-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] @@ -93,7 +93,7 @@ def test_job_retry() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-retry.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-retry.jpg') ] subprocess.run(commands) commands = [ sys.executable, 'facefusion.py', 'job-retry', 'test-job-retry', '--jobs-path', get_test_jobs_directory() ] @@ -121,13 +121,13 @@ def test_job_retry_all() -> None: commands = [ sys.executable, 'facefusion.py', 'job-create', 'test-job-retry-all-2', '--jobs-path', get_test_jobs_directory() ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test-job-retry-all-1.jpg') ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-1', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test-job-retry-all-1.jpg') ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-job-retry-all-2.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-retry-all-2.mp4'), '--trim-frame-end', '1' ] subprocess.run(commands) - commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test-job-retry-all-2.mp4'), '--trim-frame-start', '0', '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'job-add-step', 'test-job-retry-all-2', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'face_debugger', '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test-job-retry-all-2.mp4'), '--trim-frame-start', '0', '--trim-frame-end', '1' ] subprocess.run(commands) commands = [ sys.executable, 'facefusion.py', 'job-retry-all', '--jobs-path', get_test_jobs_directory(), '--halt-on-error' ] diff --git a/tests/test_cli_lip_syncer.py b/tests/test_cli_lip_syncer.py index a5f524c1..1d17276c 100644 --- a/tests/test_cli_lip_syncer.py +++ b/tests/test_cli_lip_syncer.py @@ -5,7 +5,7 @@ import pytest from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -27,28 +27,14 @@ def before_each() -> None: def test_sync_lip_to_image() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'audio-to-image:video', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test_sync_lip_to_image.mp4') ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'audio-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_file('test_sync_lip_to_image.mp4') ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test_sync_lip_to_image.mp4') is True -def test_sync_lip_to_image_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'audio-to-image:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.jpg'), '-o', get_test_output_path('test_sync_lip_to_image_as_frames') ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test_sync_lip_to_image_as_frames')) is True - - def test_sync_lip_to_video() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test_sync_lip_to_video.mp4'), '--trim-frame-end', '1' ] + commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_file('test_sync_lip_to_video.mp4'), '--trim-frame-end', '1' ] assert subprocess.run(commands).returncode == 0 assert is_test_output_file('test_sync_lip_to_video.mp4') is True - - -def test_sync_lip_to_video_as_frames() -> None: - commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video:frames', '--jobs-path', get_test_jobs_directory(), '--processors', 'lip_syncer', '-s', get_test_example_file('source.mp3'), '-t', get_test_example_file('target-240p.mp4'), '-o', get_test_output_path('test_sync_lip_to_video_as_frames'), '--trim-frame-end', '1' ] - - assert subprocess.run(commands).returncode == 0 - assert is_test_output_sequence(get_test_output_path('test_sync_lip_to_video_as_frames')) is True diff --git a/tests/test_cli_output_scale.py b/tests/test_cli_output_scale.py index d970f850..9d6337d1 100644 --- a/tests/test_cli_output_scale.py +++ b/tests/test_cli_output_scale.py @@ -7,7 +7,7 @@ from facefusion.download import conditional_download from facefusion.jobs.job_manager import clear_jobs, init_jobs from facefusion.types import Resolution, Scale from facefusion.vision import detect_image_resolution, detect_video_resolution -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -35,7 +35,7 @@ def before_each() -> None: (8.0, (3408, 1808)) ]) def test_output_image_scale(output_image_scale : Scale, output_image_resolution : Resolution) -> None: - output_file_path = get_test_output_path('test-output-image-scale-' + str(output_image_scale) + '.jpg') + output_file_path = get_test_output_file('test-output-image-scale-' + str(output_image_scale) + '.jpg') commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-image', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.jpg'), '-o', output_file_path, '--output-image-scale', str(output_image_scale) ] assert subprocess.run(commands).returncode == 0 @@ -50,7 +50,7 @@ def test_output_image_scale(output_image_scale : Scale, output_image_resolution (8.0, (3408, 1808)) ]) def test_output_video_scale(output_video_scale : Scale, output_video_resolution : Resolution) -> None: - output_file_path = get_test_output_path('test-output-video-scale-' + str(output_video_scale) + '.mp4') + output_file_path = get_test_output_file('test-output-video-scale-' + str(output_video_scale) + '.mp4') commands = [ sys.executable, 'facefusion.py', 'run', '--workflow', 'image-to-video', '--jobs-path', get_test_jobs_directory(), '--processors', 'frame_enhancer', '-t', get_test_example_file('target-240p.mp4'), '-o', output_file_path, '--trim-frame-end', '1', '--output-video-scale', str(output_video_scale) ] assert subprocess.run(commands).returncode == 0 diff --git a/tests/test_curl_builder.py b/tests/test_curl_builder.py index b0837284..ca0ec5b9 100644 --- a/tests/test_curl_builder.py +++ b/tests/test_curl_builder.py @@ -1,7 +1,7 @@ from shutil import which from facefusion import metadata -from facefusion.curl_builder import chain, ping, run, set_timeout +from facefusion.curl_builder import chain, head, run def test_run() -> None: @@ -11,7 +11,4 @@ def test_run() -> None: def test_chain() -> None: - assert chain( - ping(metadata.get('url')), - set_timeout(5) - ) == [ '-I', metadata.get('url'), '--connect-timeout', '5' ] + assert chain(head(metadata.get('url'))) == [ '-I', metadata.get('url') ] diff --git a/tests/test_ffmpeg.py b/tests/test_ffmpeg.py index 560a9fd3..8d0e937e 100644 --- a/tests/test_ffmpeg.py +++ b/tests/test_ffmpeg.py @@ -11,7 +11,7 @@ from facefusion.ffmpeg import concat_video, extract_frames, merge_video, read_au from facefusion.filesystem import copy_file from facefusion.temp_helper import clear_temp_directory, create_temp_directory, get_temp_file_path, resolve_temp_frame_paths from facefusion.types import EncoderSet -from .helper import get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -112,13 +112,13 @@ def test_spawn_frames() -> None: def test_merge_video() -> None: test_set =\ [ - (get_test_example_file('target-240p-16khz.avi'), get_test_output_path('test-merge-video-240p-16khz.avi')), - (get_test_example_file('target-240p-16khz.m4v'), get_test_output_path('test-merge-video-240p-16khz.m4v')), - (get_test_example_file('target-240p-16khz.mkv'), get_test_output_path('test-merge-video-240p-16khz.mkv')), - (get_test_example_file('target-240p-16khz.mp4'), get_test_output_path('test-merge-video-240p-16khz.mp4')), - (get_test_example_file('target-240p-16khz.mov'), get_test_output_path('test-merge-video-240p-16khz.mov')), - (get_test_example_file('target-240p-16khz.webm'), get_test_output_path('test-merge-video-240p-16khz.webm')), - (get_test_example_file('target-240p-16khz.wmv'), get_test_output_path('test-merge-video-240p-16khz.wmv')) + (get_test_example_file('target-240p-16khz.avi'), get_test_output_file('test-merge-video-240p-16khz.avi')), + (get_test_example_file('target-240p-16khz.m4v'), get_test_output_file('test-merge-video-240p-16khz.m4v')), + (get_test_example_file('target-240p-16khz.mkv'), get_test_output_file('test-merge-video-240p-16khz.mkv')), + (get_test_example_file('target-240p-16khz.mp4'), get_test_output_file('test-merge-video-240p-16khz.mp4')), + (get_test_example_file('target-240p-16khz.mov'), get_test_output_file('test-merge-video-240p-16khz.mov')), + (get_test_example_file('target-240p-16khz.webm'), get_test_output_file('test-merge-video-240p-16khz.webm')), + (get_test_example_file('target-240p-16khz.wmv'), get_test_output_file('test-merge-video-240p-16khz.wmv')) ] output_video_encoders = get_available_encoder_set().get('video') @@ -138,7 +138,7 @@ def test_merge_video() -> None: def test_concat_video() -> None: - output_path = get_test_output_path('test-concat-video.mp4') + output_path = get_test_output_file('test-concat-video.mp4') temp_output_paths =\ [ get_test_example_file('target-240p-16khz.mp4'), @@ -157,14 +157,14 @@ def test_read_audio_buffer() -> None: def test_restore_audio() -> None: test_set =\ [ - (get_test_example_file('target-240p-16khz.avi'), get_test_output_path('target-240p-16khz.avi')), - (get_test_example_file('target-240p-16khz.m4v'), get_test_output_path('target-240p-16khz.m4v')), - (get_test_example_file('target-240p-16khz.mkv'), get_test_output_path('target-240p-16khz.mkv')), - (get_test_example_file('target-240p-16khz.mov'), get_test_output_path('target-240p-16khz.mov')), - (get_test_example_file('target-240p-16khz.mp4'), get_test_output_path('target-240p-16khz.mp4')), - (get_test_example_file('target-240p-48khz.mp4'), get_test_output_path('target-240p-48khz.mp4')), - (get_test_example_file('target-240p-16khz.webm'), get_test_output_path('target-240p-16khz.webm')), - (get_test_example_file('target-240p-16khz.wmv'), get_test_output_path('target-240p-16khz.wmv')) + (get_test_example_file('target-240p-16khz.avi'), get_test_output_file('target-240p-16khz.avi')), + (get_test_example_file('target-240p-16khz.m4v'), get_test_output_file('target-240p-16khz.m4v')), + (get_test_example_file('target-240p-16khz.mkv'), get_test_output_file('target-240p-16khz.mkv')), + (get_test_example_file('target-240p-16khz.mov'), get_test_output_file('target-240p-16khz.mov')), + (get_test_example_file('target-240p-16khz.mp4'), get_test_output_file('target-240p-16khz.mp4')), + (get_test_example_file('target-240p-48khz.mp4'), get_test_output_file('target-240p-48khz.mp4')), + (get_test_example_file('target-240p-16khz.webm'), get_test_output_file('target-240p-16khz.webm')), + (get_test_example_file('target-240p-16khz.wmv'), get_test_output_file('target-240p-16khz.wmv')) ] output_audio_encoders = get_available_encoder_set().get('audio') @@ -185,13 +185,13 @@ def test_restore_audio() -> None: def test_replace_audio() -> None: test_set =\ [ - (get_test_example_file('target-240p-16khz.avi'), get_test_output_path('target-240p-16khz.avi')), - (get_test_example_file('target-240p-16khz.m4v'), get_test_output_path('target-240p-16khz.m4v')), - (get_test_example_file('target-240p-16khz.mkv'), get_test_output_path('target-240p-16khz.mkv')), - (get_test_example_file('target-240p-16khz.mov'), get_test_output_path('target-240p-16khz.mov')), - (get_test_example_file('target-240p-16khz.mp4'), get_test_output_path('target-240p-16khz.mp4')), - (get_test_example_file('target-240p-48khz.mp4'), get_test_output_path('target-240p-48khz.mp4')), - (get_test_example_file('target-240p-16khz.webm'), get_test_output_path('target-240p-16khz.webm')) + (get_test_example_file('target-240p-16khz.avi'), get_test_output_file('target-240p-16khz.avi')), + (get_test_example_file('target-240p-16khz.m4v'), get_test_output_file('target-240p-16khz.m4v')), + (get_test_example_file('target-240p-16khz.mkv'), get_test_output_file('target-240p-16khz.mkv')), + (get_test_example_file('target-240p-16khz.mov'), get_test_output_file('target-240p-16khz.mov')), + (get_test_example_file('target-240p-16khz.mp4'), get_test_output_file('target-240p-16khz.mp4')), + (get_test_example_file('target-240p-48khz.mp4'), get_test_output_file('target-240p-48khz.mp4')), + (get_test_example_file('target-240p-16khz.webm'), get_test_output_file('target-240p-16khz.webm')) ] output_audio_encoders = get_available_encoder_set().get('audio') diff --git a/tests/test_job_manager.py b/tests/test_job_manager.py index 2d263bcc..27164146 100644 --- a/tests/test_job_manager.py +++ b/tests/test_job_manager.py @@ -37,12 +37,6 @@ def test_submit_job() -> None: 'target_path': 'target-1.jpg', 'output_path': 'output-1.jpg' } - args_2 =\ - { - 'source_path': 'source-2.jpg', - 'target_path': 'target-2.mp4', - 'output_path': 'output-sequence-2' - } assert submit_job('job-invalid') is False @@ -51,7 +45,6 @@ def test_submit_job() -> None: assert submit_job('job-test-submit-job') is False add_step('job-test-submit-job', args_1) - add_step('job-test-submit-job', args_2) assert submit_job('job-test-submit-job') is True assert submit_job('job-test-submit-job') is False @@ -177,18 +170,6 @@ def test_add_step() -> None: 'target_path': 'target-2.jpg', 'output_path': 'output-2.jpg' } - args_3 =\ - { - 'source_path': 'source-3.jpg', - 'target_path': 'target-3.mp4', - 'output_path': 'output-sequence-1' - } - args_4 =\ - { - 'source_path': 'source-4.jpg', - 'target_path': 'target-4.mp4', - 'output_path': 'output-sequence-1' - } assert add_step('job-invalid', args_1) is False @@ -196,16 +177,12 @@ def test_add_step() -> None: assert add_step('job-test-add-step', args_1) is True assert add_step('job-test-add-step', args_2) is True - assert add_step('job-test-add-step', args_3) is True - assert add_step('job-test-add-step', args_4) is True steps = get_steps('job-test-add-step') assert steps[0].get('args') == args_1 assert steps[1].get('args') == args_2 - assert steps[2].get('args') == args_3 - assert steps[3].get('args') == args_4 - assert count_step_total('job-test-add-step') == 4 + assert count_step_total('job-test-add-step') == 2 def test_remix_step() -> None: @@ -221,40 +198,28 @@ def test_remix_step() -> None: 'target_path': 'target-2.jpg', 'output_path': 'output-2.jpg' } - args_3 =\ - { - 'source_path': 'source-3.jpg', - 'target_path': 'target-3.mp4', - 'output_path': 'output-sequence-3' - } assert remix_step('job-invalid', 0, args_1) is False create_job('job-test-remix-step') add_step('job-test-remix-step', args_1) add_step('job-test-remix-step', args_2) - add_step('job-test-remix-step', args_3) assert remix_step('job-test-remix-step', 99, args_1) is False assert remix_step('job-test-remix-step', 0, args_2) is True assert remix_step('job-test-remix-step', -1, args_2) is True - assert remix_step('job-test-remix-step', 2, args_3) is True steps = get_steps('job-test-remix-step') assert steps[0].get('args') == args_1 assert steps[1].get('args') == args_2 - assert steps[2].get('args') == args_3 + assert steps[2].get('args').get('source_path') == args_2.get('source_path') + assert steps[2].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 0, args_1.get('output_path')) + assert steps[2].get('args').get('output_path') == args_2.get('output_path') assert steps[3].get('args').get('source_path') == args_2.get('source_path') - assert steps[3].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 0, args_1.get('output_path')) + assert steps[3].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 2, args_2.get('output_path')) assert steps[3].get('args').get('output_path') == args_2.get('output_path') - assert steps[4].get('args').get('source_path') == args_2.get('source_path') - assert steps[4].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 3, args_2.get('output_path')) - assert steps[4].get('args').get('output_path') == args_2.get('output_path') - assert steps[5].get('args').get('source_path') == args_3.get('source_path') - assert steps[5].get('args').get('target_path') == get_step_output_path('job-test-remix-step', 2, args_3.get('output_path')) - assert steps[5].get('args').get('output_path') == args_3.get('output_path') - assert count_step_total('job-test-remix-step') == 6 + assert count_step_total('job-test-remix-step') == 4 def test_insert_step() -> None: @@ -276,12 +241,6 @@ def test_insert_step() -> None: 'target_path': 'target-3.jpg', 'output_path': 'output-3.jpg' } - args_4 =\ - { - 'source_path': 'source-4.jpg', - 'target_path': 'target-4.mp4', - 'output_path': 'output-sequence-4' - } assert insert_step('job-invalid', 0, args_1) is False @@ -292,16 +251,14 @@ def test_insert_step() -> None: assert insert_step('job-test-insert-step', 99, args_1) is False assert insert_step('job-test-insert-step', 0, args_2) is True assert insert_step('job-test-insert-step', -1, args_3) is True - assert insert_step('job-test-insert-step', 2, args_4) is True steps = get_steps('job-test-insert-step') assert steps[0].get('args') == args_2 assert steps[1].get('args') == args_1 - assert steps[2].get('args') == args_4 - assert steps[3].get('args') == args_3 - assert steps[4].get('args') == args_1 - assert count_step_total('job-test-insert-step') == 5 + assert steps[2].get('args') == args_3 + assert steps[3].get('args') == args_1 + assert count_step_total('job-test-insert-step') == 4 def test_remove_step() -> None: diff --git a/tests/test_job_runner.py b/tests/test_job_runner.py index 35d57060..4e0983ae 100644 --- a/tests/test_job_runner.py +++ b/tests/test_job_runner.py @@ -1,14 +1,13 @@ -import os import subprocess import pytest from facefusion.download import conditional_download -from facefusion.filesystem import copy_file, create_directory, get_file_extension +from facefusion.filesystem import copy_file from facefusion.jobs.job_manager import add_step, clear_jobs, create_job, init_jobs, move_job_file, submit_job, submit_jobs from facefusion.jobs.job_runner import collect_output_set, finalize_steps, retry_job, retry_jobs, run_job, run_jobs, run_steps from facefusion.types import Args -from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_path, is_test_output_file, is_test_output_sequence, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_jobs_directory, get_test_output_file, is_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -29,15 +28,7 @@ def before_each() -> None: def process_step(job_id : str, step_index : int, step_args : Args) -> bool: - output_path = step_args.get('output_path') - target_path = step_args.get('target_path') - - if output_path and not get_file_extension(output_path): - if create_directory(output_path): - return copy_file(target_path, os.path.join(output_path, os.path.basename(target_path))) - return False - - return copy_file(target_path, output_path) + return copy_file(step_args.get('target_path'), step_args.get('output_path')) def test_run_job() -> None: @@ -45,31 +36,19 @@ def test_run_job() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } args_2 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-2.mp4') + 'output_path': get_test_output_file('output-2.mp4') } args_3 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-3.jpg') - } - args_4 =\ - { - 'source_path': get_test_example_file('source.jpg'), - 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-4') - } - args_5 =\ - { - 'source_path': get_test_example_file('source.jpg'), - 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-4') + 'output_path': get_test_output_file('output-3.jpg') } assert run_job('job-invalid', process_step) is False @@ -79,8 +58,6 @@ def test_run_job() -> None: add_step('job-test-run-job', args_2) add_step('job-test-run-job', args_2) add_step('job-test-run-job', args_3) - add_step('job-test-run-job', args_4) - add_step('job-test-run-job', args_5) assert run_job('job-test-run-job', process_step) is False @@ -94,19 +71,19 @@ def test_run_jobs() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } args_2 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-2.mp4') + 'output_path': get_test_output_file('output-2.mp4') } args_3 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-3.jpg') + 'output_path': get_test_output_file('output-3.jpg') } halt_on_error = True @@ -131,7 +108,7 @@ def test_retry_job() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } assert retry_job('job-invalid', process_step) is False @@ -152,19 +129,19 @@ def test_retry_jobs() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } args_2 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-2.mp4') + 'output_path': get_test_output_file('output-2.mp4') } args_3 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-3.jpg') + 'output_path': get_test_output_file('output-3.jpg') } halt_on_error = True @@ -190,19 +167,19 @@ def test_run_steps() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } args_2 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-2.mp4') + 'output_path': get_test_output_file('output-2.mp4') } args_3 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-3.jpg') + 'output_path': get_test_output_file('output-3.jpg') } assert run_steps('job-invalid', process_step) is False @@ -221,31 +198,19 @@ def test_finalize_steps() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } args_2 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-2.mp4') + 'output_path': get_test_output_file('output-2.mp4') } args_3 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-3.jpg') - } - args_4 =\ - { - 'source_path': get_test_example_file('source.jpg'), - 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-4') - } - args_5 =\ - { - 'source_path': get_test_example_file('source.jpg'), - 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-4') + 'output_path': get_test_output_file('output-3.jpg') } create_job('job-test-finalize-steps') @@ -253,26 +218,16 @@ def test_finalize_steps() -> None: add_step('job-test-finalize-steps', args_1) add_step('job-test-finalize-steps', args_2) add_step('job-test-finalize-steps', args_3) - add_step('job-test-finalize-steps', args_4) - add_step('job-test-finalize-steps', args_5) - copy_file(args_1.get('target_path'), get_test_output_path('output-1-job-test-finalize-steps-0.mp4')) - copy_file(args_1.get('target_path'), get_test_output_path('output-1-job-test-finalize-steps-1.mp4')) - copy_file(args_2.get('target_path'), get_test_output_path('output-2-job-test-finalize-steps-2.mp4')) - copy_file(args_3.get('target_path'), get_test_output_path('output-3-job-test-finalize-steps-3.jpg')) - - temp_directory_1 = get_test_output_path('output-4-job-test-finalize-steps-4') - temp_directory_2 = get_test_output_path('output-4-job-test-finalize-steps-5') - create_directory(temp_directory_1) - create_directory(temp_directory_2) - copy_file(args_4.get('target_path'), os.path.join(temp_directory_1, '00000001.jpg')) - copy_file(args_5.get('target_path'), os.path.join(temp_directory_2, '00000002.jpg')) + copy_file(args_1.get('target_path'), get_test_output_file('output-1-job-test-finalize-steps-0.mp4')) + copy_file(args_1.get('target_path'), get_test_output_file('output-1-job-test-finalize-steps-1.mp4')) + copy_file(args_2.get('target_path'), get_test_output_file('output-2-job-test-finalize-steps-2.mp4')) + copy_file(args_3.get('target_path'), get_test_output_file('output-3-job-test-finalize-steps-3.jpg')) assert finalize_steps('job-test-finalize-steps') is True assert is_test_output_file('output-1.mp4') is True assert is_test_output_file('output-2.mp4') is True assert is_test_output_file('output-3.jpg') is True - assert is_test_output_sequence(get_test_output_path('output-4')) is True def test_collect_output_set() -> None: @@ -280,25 +235,19 @@ def test_collect_output_set() -> None: { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-1.mp4') + 'output_path': get_test_output_file('output-1.mp4') } args_2 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-2.mp4') + 'output_path': get_test_output_file('output-2.mp4') } args_3 =\ { 'source_path': get_test_example_file('source.jpg'), 'target_path': get_test_example_file('target-240p.jpg'), - 'output_path': get_test_output_path('output-3.jpg') - } - args_4 = \ - { - 'source_path': get_test_example_file('source.jpg'), - 'target_path': get_test_example_file('target-240p.mp4'), - 'output_path': get_test_output_path('output-4') + 'output_path': get_test_output_file('output-3.jpg') } create_job('job-test-collect-output-set') @@ -306,26 +255,21 @@ def test_collect_output_set() -> None: add_step('job-test-collect-output-set', args_1) add_step('job-test-collect-output-set', args_2) add_step('job-test-collect-output-set', args_3) - add_step('job-test-collect-output-set', args_4) output_set =\ { - get_test_output_path('output-1.mp4'): + get_test_output_file('output-1.mp4'): [ - get_test_output_path('output-1-job-test-collect-output-set-0.mp4'), - get_test_output_path('output-1-job-test-collect-output-set-1.mp4') + get_test_output_file('output-1-job-test-collect-output-set-0.mp4'), + get_test_output_file('output-1-job-test-collect-output-set-1.mp4') ], - get_test_output_path('output-2.mp4'): + get_test_output_file('output-2.mp4'): [ - get_test_output_path('output-2-job-test-collect-output-set-2.mp4') + get_test_output_file('output-2-job-test-collect-output-set-2.mp4') ], - get_test_output_path('output-3.jpg'): + get_test_output_file('output-3.jpg'): [ - get_test_output_path('output-3-job-test-collect-output-set-3.jpg') - ], - get_test_output_path('output-4'): - [ - get_test_output_path('output-4-job-test-collect-output-set-4') + get_test_output_file('output-3-job-test-collect-output-set-3.jpg') ] } diff --git a/tests/test_program_helper.py b/tests/test_program_helper.py index a2e85a10..92b64fb2 100644 --- a/tests/test_program_helper.py +++ b/tests/test_program_helper.py @@ -1,6 +1,8 @@ from argparse import ArgumentParser -from facefusion.program_helper import find_argument_group, validate_actions, validate_args +import pytest + +from facefusion.program_helper import find_argument_group, validate_actions def test_find_argument_group() -> None: @@ -10,26 +12,12 @@ def test_find_argument_group() -> None: assert find_argument_group(program, 'test-1') assert find_argument_group(program, 'test-2') - assert find_argument_group(program, 'test-3') is None + assert find_argument_group(program, 'invalid') is None +@pytest.mark.skip() def test_validate_args() -> None: - program = ArgumentParser() - program.add_argument('--test-1', default = 'test_1', choices = [ 'test_1', 'test_2' ]) - - assert validate_args(program) is True - - subparsers = program.add_subparsers() - sub_program = subparsers.add_parser('sub-command') - sub_program.add_argument('--test-2', default = 'test_2', choices = [ 'test_1', 'test_2' ]) - - assert validate_args(program) is True - - for action in sub_program._actions: - if action.dest == 'test_2': - action.default = 'test_3' - - assert validate_args(program) is False + pass def test_validate_actions() -> None: diff --git a/tests/test_translator.py b/tests/test_translator.py index 3773cb9e..dbf85280 100644 --- a/tests/test_translator.py +++ b/tests/test_translator.py @@ -1,11 +1,11 @@ from facefusion import translator -from facefusion.locales import LOCALES +from facefusion.locals import LOCALS def test_load() -> None: - translator.load(LOCALES, __name__) + translator.load(LOCALS, __name__) - assert __name__ in translator.LOCALE_POOL_SET + assert __name__ in translator.LOCAL_POOL_SET def test_get() -> None: diff --git a/tests/test_vision.py b/tests/test_vision.py index 2eb2aae6..525f0f80 100644 --- a/tests/test_vision.py +++ b/tests/test_vision.py @@ -4,7 +4,7 @@ import pytest from facefusion.download import conditional_download from facefusion.vision import calculate_histogram_difference, count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_frame, restrict_image_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, unpack_resolution, write_image -from .helper import get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory +from .helper import get_test_example_file, get_test_examples_directory, get_test_output_file, prepare_test_output_directory @pytest.fixture(scope = 'module', autouse = True) @@ -42,8 +42,8 @@ def test_read_image() -> None: def test_write_image() -> None: vision_frame = read_image(get_test_example_file('target-240p.jpg')) - assert write_image(get_test_output_path('target-240p.jpg'), vision_frame) is True - assert write_image(get_test_output_path('目标-240p.webp'), vision_frame) is True + assert write_image(get_test_output_file('target-240p.jpg'), vision_frame) is True + assert write_image(get_test_output_file('目标-240p.webp'), vision_frame) is True def test_detect_image_resolution() -> None: