state api updates

facefusion/apis/asset_store.py

@@ -0,0 +1,23 @@
+import uuid
+from typing import Dict, Optional
+
+ASSET_STORE: Dict[str, Dict[str, str]] = {}
+
+
+def get_asset(asset_id : str) -> Optional[Dict[str, str]]:
+	return ASSET_STORE.get(asset_id)
+
+
+def register_asset(path : str) -> str:
+	asset_id = str(uuid.uuid4())
+
+	ASSET_STORE[asset_id] =\
+	{
+		'id': asset_id,
+		'path': path
+	}
+	return asset_id
+
+
+def clear() -> None:
+	ASSET_STORE.clear()

facefusion/apis/endpoints/state.py

@@ -1,8 +1,9 @@
 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_404_NOT_FOUND
 
-from facefusion import args_store, state_manager
+from facefusion import args_store, state_manager, translator
+from facefusion.apis import asset_store
 
 
 async def get_state(request : Request) -> JSONResponse:
@@ -11,6 +12,17 @@ async def get_state(request : Request) -> JSONResponse:
 
 
 async def set_state(request : Request) -> JSONResponse:
+	action = request.query_params.get('action')
+
+	if action == 'select':
+		asset_type = request.query_params.get('type')
+
+		if asset_type == 'source':
+			return await select_source(request)
+
+		if asset_type == 'target':
+			return await select_target(request)
+
 	body = await request.json()
 	api_args = args_store.get_api_args()
 
@@ -18,5 +30,48 @@ async def set_state(request : Request) -> JSONResponse:
 		if key in api_args:
 			state_manager.set_item(key, value)
 
-	__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.filter_api_args(state_manager.get_state()) #type:ignore[arg-type,assignment]
+	return JSONResponse(state_manager.collect_state(api_args), status_code = HTTP_200_OK) #type:ignore[arg-type]
+
+
+async def select_source(request : Request) -> JSONResponse:
+	body = await request.json()
+	asset_ids = body.get('asset_ids')
+
+	if isinstance(asset_ids, list):
+		source_paths = []
+
+		for asset_id in asset_ids:
+			asset = asset_store.get_asset(asset_id)
+
+			if asset:
+				source_paths.append(asset.get('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)
+
+	return JSONResponse(
+	{
+		'message': translator.get('source_asset_not_found', 'facefusion.apis')
+	}, status_code = HTTP_404_NOT_FOUND)
+
+
+async def select_target(request : Request) -> JSONResponse:
+	body = await request.json()
+	asset_id = body.get('asset_id')
+
+	if isinstance(asset_id, str):
+		asset = asset_store.get_asset(asset_id)
+
+		if asset:
+			state_manager.set_item('target_path', asset.get('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)
+
+	return JSONResponse(
+	{
+		'message': translator.get('target_asset_not_found', 'facefusion.apis')
+	}, status_code = HTTP_404_NOT_FOUND)

facefusion/apis/locales.py

@@ -7,6 +7,8 @@ LOCALES : Locales =\
 		'ok': 'ok',
 		'something_went_wrong': 'something went wrong',
 		'invalid_access_token': 'invalid access token',
-		'invalid_refresh_token': 'invalid refresh token'
+		'invalid_refresh_token': 'invalid refresh token',
+		'source_asset_not_found': 'source asset not found',
+		'target_asset_not_found': 'target asset not found'
 	}
 }