add missing endpoints

facefusion/apis/asset_store.py

@@ -1,11 +1,11 @@
 import os
 import uuid
 from datetime import datetime, timedelta
-from typing import Optional, cast
+from typing import List, Optional, cast
 
 from facefusion.apis.asset_helper import detect_media_type, extract_audio_metadata, extract_image_metadata, extract_video_metadata
 from facefusion.filesystem import get_file_format, get_file_name
-from facefusion.types import AssetId, AssetStore, AssetType, AudioAsset, AudioFormat, ImageAsset, ImageFormat, SessionId, VideoAsset, VideoFormat
+from facefusion.types import AssetId, AssetSet, AssetStore, AssetType, AudioAsset, AudioFormat, ImageAsset, ImageFormat, SessionId, VideoAsset, VideoFormat
 
 ASSET_STORE : AssetStore = {}
 
@@ -70,9 +70,21 @@ def create_asset(session_id : SessionId, asset_type : AssetType, asset_path : st
 	return ASSET_STORE[session_id].get(asset_id)
 
 
+def get_assets(session_id : SessionId) -> Optional[AssetSet]:
+	return ASSET_STORE.get(session_id)
+
+
 def get_asset(session_id : SessionId, asset_id : AssetId) -> Optional[AudioAsset | ImageAsset | VideoAsset]:
 	if session_id in ASSET_STORE:
-		return ASSET_STORE[session_id].get(asset_id)
+		return ASSET_STORE.get(session_id).get(asset_id)
+	return None
+
+
+def delete_assets(session_id : SessionId, asset_ids : List[AssetId]) -> None:
+	if session_id in ASSET_STORE:
+		for asset_id in asset_ids:
+			if asset_id in ASSET_STORE.get(session_id):
+				del ASSET_STORE[session_id][asset_id]
 	return None
 
 

facefusion/apis/core.py

@@ -3,7 +3,7 @@ from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
 from starlette.routing import Route, WebSocketRoute
 
-from facefusion.apis.endpoints.assets import upload_asset
+from facefusion.apis.endpoints.assets import delete_assets, get_asset, get_assets, upload_asset
 from facefusion.apis.endpoints.ping import websocket_ping
 from facefusion.apis.endpoints.session import create_session, create_session_guard, destroy_session, get_session, refresh_session
 from facefusion.apis.endpoints.state import get_state, set_state
@@ -19,7 +19,10 @@ def create_api() -> Starlette:
 			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 ]),
+			Route('/assets', get_assets, methods = [ 'GET' ], middleware = [ session_guard ]),
 			Route('/assets', upload_asset, methods = [ 'POST' ], middleware = [ session_guard ]),
+			Route('/assets/{asset_id}', get_asset, methods = [ 'GET' ], middleware = [ session_guard ]),
+			Route('/assets', delete_assets, methods = [ 'DELETE' ], middleware = [ session_guard ]),
 			WebSocketRoute('/ping', websocket_ping, middleware = [ session_guard ])
 		]
 

facefusion/apis/endpoints/assets.py

@@ -5,7 +5,7 @@ from typing import List
 from starlette.datastructures import UploadFile
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
-from starlette.status import HTTP_201_CREATED, HTTP_400_BAD_REQUEST
+from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_400_BAD_REQUEST, HTTP_404_NOT_FOUND
 
 from facefusion import ffmpeg, process_manager, session_manager, state_manager
 from facefusion.apis import asset_store
@@ -22,7 +22,7 @@ async def upload_asset(request : Request) -> Response:
 	if session_id and asset_type in [ 'source', 'target' ]:
 		form = await request.form()
 		upload_files = form.getlist('file')
-		asset_paths = await save_asset_files(upload_files)
+		asset_paths = await save_asset_files(upload_files) # type: ignore[arg-type]
 
 		if asset_paths:
 			asset_ids : List[str] = []
@@ -76,3 +76,78 @@ async def save_asset_files(upload_files : List[UploadFile]) -> List[str]:
 		remove_file(temp_file.name)
 
 	return asset_paths
+
+
+async def get_assets(request : Request) -> Response:
+	access_token = extract_access_token(request.scope)
+	session_id = session_manager.find_session_id(access_token)
+
+	if session_id:
+		asset_set = asset_store.get_assets(session_id)
+		assets = []
+
+		if asset_set:
+			for asset in asset_set.values():
+				assets.append(
+				{
+					'id': asset.get('id'),
+					'created_at': asset.get('created_at').isoformat(),
+					'expires_at': asset.get('expires_at').isoformat(),
+					'type': asset.get('type'),
+					'media': asset.get('media'),
+					'name': asset.get('name'),
+					'format': asset.get('format'),
+					'size': asset.get('size'),
+					'metadata': asset.get('metadata')
+				})
+
+		return JSONResponse(
+		{
+			'assets': assets
+		}, status_code = HTTP_200_OK)
+
+	return Response(status_code = HTTP_400_BAD_REQUEST)
+
+
+async def get_asset(request : Request) -> Response:
+	access_token = extract_access_token(request.scope)
+	session_id = session_manager.find_session_id(access_token)
+	asset_id = request.path_params.get('asset_id')
+
+	if session_id and asset_id:
+		asset = asset_store.get_asset(session_id, asset_id)
+
+		if asset:
+			return JSONResponse(
+			{
+				'id': asset.get('id'),
+				'created_at': asset.get('created_at').isoformat(),
+				'expires_at': asset.get('expires_at').isoformat(),
+				'type': asset.get('type'),
+				'media': asset.get('media'),
+				'name': asset.get('name'),
+				'format': asset.get('format'),
+				'size': asset.get('size'),
+				'metadata': asset.get('metadata')
+			}, status_code = HTTP_200_OK)
+
+	return Response(status_code = HTTP_404_NOT_FOUND)
+
+
+async def delete_assets(request : Request) -> Response:
+	access_token = extract_access_token(request.scope)
+	session_id = session_manager.find_session_id(access_token)
+	body = await request.json()
+	asset_ids = body.get('asset_ids')
+
+	if session_id and asset_ids:
+		asset_set = asset_store.get_assets(session_id)
+
+		if asset_set:
+			for asset_id in asset_ids:
+				if asset_id in asset_set:
+					remove_file(asset_set.get(asset_id).get('path'))
+			asset_store.delete_assets(session_id, asset_ids)
+			return Response(status_code = HTTP_200_OK)
+
+	return Response(status_code = HTTP_404_NOT_FOUND)

facefusion/types.py

@@ -232,7 +232,8 @@ VideoAsset = TypedDict('VideoAsset',
 })
 
 AssetMetadata : TypeAlias = AudioMetadata | ImageMetadata | VideoMetadata
-AssetStore : TypeAlias = Dict[SessionId, Dict[AssetId, AudioAsset | ImageAsset | VideoAsset]]
+AssetSet : TypeAlias = Dict[AssetId, AudioAsset | ImageAsset | VideoAsset]
+AssetStore : TypeAlias = Dict[SessionId, AssetSet]
 
 BenchmarkMode = Literal['warm', 'cold']
 BenchmarkResolution = Literal['240p', '360p', '540p', '720p', '1080p', '1440p', '2160p']