asset store update

facefusion/apis/api_helper.py

@@ -3,6 +3,10 @@ from typing import Optional
 from starlette.datastructures import Headers
 from starlette.types import Scope
 
+from facefusion.audio import detect_audio_duration
+from facefusion.types import AudioMetadata, ImageMetadata, VideoMetadata
+from facefusion.vision import count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution
+
 
 def get_sec_websocket_protocol(scope : Scope) -> Optional[str]:
 	protocol_header = Headers(scope = scope).get('Sec-WebSocket-Protocol')
@@ -13,3 +17,29 @@ def get_sec_websocket_protocol(scope : Scope) -> Optional[str]:
 
 	return None
 
+
+def extract_audio_metadata(file_path : str) -> AudioMetadata:
+	metadata : AudioMetadata =\
+	{
+		'duration' : detect_audio_duration(file_path)
+	}
+	return metadata
+
+
+def extract_image_metadata(file_path : str) -> ImageMetadata:
+	metadata : ImageMetadata =\
+	{
+		'resolution' : detect_image_resolution(file_path)
+	}
+	return metadata
+
+
+def extract_video_metadata(file_path : str) -> VideoMetadata:
+	metadata : VideoMetadata =\
+	{
+		'duration' : detect_video_duration(file_path),
+		'frame_total' : count_video_frame_total(file_path),
+		'fps' : detect_video_fps(file_path),
+		'resolution' : detect_video_resolution(file_path)
+	}
+	return metadata

facefusion/apis/asset_store.py

@@ -1,23 +1,82 @@
+import os
 import uuid
-from typing import Dict, Optional
+from datetime import datetime, timedelta
+from typing import Optional, cast
 
-ASSET_STORE : Dict[str, Dict[str, str]] = {}
+from facefusion.apis.api_helper import extract_audio_metadata, extract_image_metadata, extract_video_metadata
+from facefusion.filesystem import get_file_format, get_file_name, is_audio, is_image, is_video
+from facefusion.types import Asset, AssetId, AssetStore, AssetType, AudioFormat, ImageFormat, MediaType, SessionId, VideoFormat
+
+ASSET_STORE : AssetStore = {}
 
 
-def get_asset(asset_id : str) -> Optional[Dict[str, str]]:
-	return ASSET_STORE.get(asset_id)
+def get_asset(session_id : SessionId, asset_id : AssetId) -> Optional[Asset]:
+	if session_id in ASSET_STORE:
+		return ASSET_STORE[session_id].get(asset_id)
+	return None
 
 
-def register_asset(path : str) -> str:
+def create_asset(session_id : SessionId, asset_type : AssetType, file_path : str) -> Optional[Asset]:
 	asset_id = str(uuid.uuid4())
+	media_type = detect_media_type(file_path)
 
-	ASSET_STORE[asset_id] =\
-	{
-		'id': asset_id,
-		'path': path
-	}
-	return asset_id
+	if media_type:
+		file_name = get_file_name(file_path)
+		file_format = get_file_format(file_path)
+		file_size = os.path.getsize(file_path)
+		created_at = datetime.now()
+		expires_at = created_at + timedelta(hours = 24)
+		asset =\
+		{
+			'id': asset_id,
+			'created_at': created_at,
+			'expires_at': expires_at,
+			'type': asset_type,
+			'name': file_name,
+			'size': file_size,
+			'path': file_path
+		}
+
+		if media_type == 'audio':
+			asset.update(
+			{
+				'media': 'audio',
+				'format': cast(AudioFormat, file_format),
+				'metadata': extract_audio_metadata(file_path)
+			})
+		if media_type == 'image':
+			asset.update(
+			{
+				'media': 'image',
+				'format': cast(ImageFormat, file_format),
+				'metadata': extract_image_metadata(file_path)
+			})
+		if media_type == 'video':
+			asset.update(
+			{
+				'media': 'video',
+				'format': cast(VideoFormat, file_format),
+				'metadata': extract_video_metadata(file_path)
+			})
+
+		if session_id not in ASSET_STORE:
+			ASSET_STORE[session_id] = {}
+
+		ASSET_STORE[session_id][asset_id] = asset #type:ignore[assignment]
+		return asset #type:ignore[return-value]
+
+	return None
 
 
 def clear() -> None:
 	ASSET_STORE.clear()
+
+
+def detect_media_type(file_path : str) -> Optional[MediaType]:
+	if is_image(file_path):
+		return 'image'
+	if is_video(file_path):
+		return 'video'
+	if is_audio(file_path):
+		return 'audio'
+	return None

facefusion/apis/endpoints/state.py

@@ -2,8 +2,9 @@ from starlette.requests import Request
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_200_OK, HTTP_404_NOT_FOUND
 
-from facefusion import args_store, state_manager, translator
+from facefusion import args_store, session_manager, state_manager, translator
 from facefusion.apis import asset_store
+from facefusion.apis.endpoints.session import extract_access_token
 
 
 async def get_state(request : Request) -> JSONResponse:
@@ -35,12 +36,14 @@ async def set_state(request : Request) -> JSONResponse:
 async def select_source(request : Request) -> JSONResponse:
 	body = await request.json()
 	asset_ids = body.get('asset_ids')
+	access_token = extract_access_token(request.scope)
+	session_id = session_manager.find_session_id(access_token)
 
-	if isinstance(asset_ids, list):
+	if isinstance(asset_ids, list) and session_id:
 		source_paths = []
 
 		for asset_id in asset_ids:
-			asset = asset_store.get_asset(asset_id)
+			asset = asset_store.get_asset(session_id, asset_id)
 
 			if asset:
 				source_paths.append(asset.get('path'))
@@ -59,9 +62,11 @@ async def select_source(request : Request) -> JSONResponse:
 async def select_target(request : Request) -> JSONResponse:
 	body = await request.json()
 	asset_id = body.get('asset_id')
+	access_token = extract_access_token(request.scope)
+	session_id = session_manager.find_session_id(access_token)
 
-	if isinstance(asset_id, str):
-		asset = asset_store.get_asset(asset_id)
+	if isinstance(asset_id, str) and session_id:
+		asset = asset_store.get_asset(session_id, asset_id)
 
 		if asset:
 			state_manager.set_item('target_path', asset.get('path'))

facefusion/types.py

@@ -169,6 +169,68 @@ EncoderSet = TypedDict('EncoderSet',
 })
 VideoPreset = Literal['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow']
 
+AssetId : TypeAlias = str
+AssetType = Literal['source', 'target']
+MediaType = Literal['image', 'video', 'audio']
+AudioMetadata = TypedDict('AudioMetadata',
+{
+	'duration' : Duration
+})
+ImageMetadata = TypedDict('ImageMetadata',
+{
+	'resolution' : Resolution
+})
+VideoMetadata = TypedDict('VideoMetadata',
+{
+	'duration' : Duration,
+	'frame_total' : int,
+	'fps' : Fps,
+	'resolution' : Resolution
+})
+AudioAsset = TypedDict('AudioAsset',
+{
+	'id' : AssetId,
+	'created_at' : datetime,
+	'expires_at' : datetime,
+	'type' : AssetType,
+	'media' : Literal['audio'],
+	'name' : str,
+	'format' : AudioFormat,
+	'size' : int,
+	'path' : str,
+	'metadata' : AudioMetadata
+})
+ImageAsset = TypedDict('ImageAsset',
+{
+	'id' : AssetId,
+	'created_at' : datetime,
+	'expires_at' : datetime,
+	'type' : AssetType,
+	'media' : Literal['image'],
+	'name' : str,
+	'format' : ImageFormat,
+	'size' : int,
+	'path' : str,
+	'metadata' : ImageMetadata
+})
+VideoAsset = TypedDict('VideoAsset',
+{
+	'id' : AssetId,
+	'created_at' : datetime,
+	'expires_at' : datetime,
+	'type' : AssetType,
+	'media' : Literal['video'],
+	'name' : str,
+	'format' : VideoFormat,
+	'size' : int,
+	'path' : str,
+	'metadata' : VideoMetadata
+})
+
+Asset : TypeAlias = AudioAsset | ImageAsset | VideoAsset
+AssetMetadata : TypeAlias = AudioMetadata | ImageMetadata | VideoMetadata
+AssetStore : TypeAlias = Dict[SessionId, Dict[AssetId, Asset]]
+
 BenchmarkMode = Literal['warm', 'cold']
 BenchmarkResolution = Literal['240p', '360p', '540p', '720p', '1080p', '1440p', '2160p']
 BenchmarkSet : TypeAlias = Dict[BenchmarkResolution, str]