switch to multipart output for asset capture to avoid overload

facefusion/apis/endpoints/assets.py

@@ -12,9 +12,10 @@ from facefusion.apis import asset_store
 from facefusion.apis.asset_helper import capture_asset_faces, capture_asset_frames, save_asset_files
 from facefusion.apis.asset_validator import validate_asset_resolution, validate_asset_type, validate_frame_index, validate_subject
 from facefusion.apis.media_validator import validate_resolution
+from facefusion.apis.multipart_helper import create_multipart_response
 from facefusion.filesystem import remove_file
 from facefusion.types import AssetType, ImageAsset, VideoAsset
-from facefusion.vision import is_vision_frames, to_strip_buffer
+from facefusion.vision import is_vision_frames
 
 
 async def get_assets(request : Request) -> Response:
@@ -63,7 +64,7 @@ async def get_asset(request : Request) -> Response:
 					vision_frames = await asyncio.to_thread(capture_asset_faces, cast(ImageAsset | VideoAsset, asset), frame_indexes, resolution)
 
 			if is_vision_frames(vision_frames):
-				return Response(content = to_strip_buffer(vision_frames), media_type = 'image/jpeg')
+				return create_multipart_response(vision_frames)
 
 		return Response(status_code = HTTP_400_BAD_REQUEST)
 

facefusion/apis/multipart_helper.py

@@ -0,0 +1,31 @@
+import secrets
+from typing import Iterator, List
+
+from starlette.responses import StreamingResponse
+
+from facefusion.types import Buffer, VisionFrame
+from facefusion.vision import to_buffer
+
+
+def create_multipart_response(vision_frames : List[VisionFrame]) -> StreamingResponse:
+	boundary = secrets.token_hex(16)
+	content = stream_multipart(vision_frames, boundary)
+	media_type = 'multipart/mixed; boundary=' + boundary
+
+	return StreamingResponse(content, media_type = media_type)
+
+
+def stream_multipart(vision_frames : List[VisionFrame], boundary : str) -> Iterator[Buffer]:
+	eol = '\r\n'
+
+	for vision_frame in vision_frames:
+		vision_buffer = to_buffer(vision_frame)
+		content_total = len(vision_buffer)
+
+		part_header = '--' + boundary + eol
+		part_header += 'Content-Type: image/jpeg' + eol
+		part_header += 'Content-Length: ' + str(content_total) + eol + eol
+
+		yield part_header.encode() + vision_buffer + eol.encode()
+
+	yield ('--' + boundary + '--' + eol).encode()

facefusion/types.py

@@ -276,7 +276,7 @@ VideoReader = TypedDict('VideoReader',
 {
 	'id' : str,
 	'file_path' : str,
-	'process' : subprocess.Popen[bytes],
+	'process' : subprocess.Popen[Buffer],
 	'metadata' : VideoReaderMetadata,
 	'frame_index' : int
 })
@@ -285,7 +285,7 @@ VideoWriter = TypedDict('VideoWriter',
 {
 	'id' : str,
 	'file_path' : str,
-	'process' : subprocess.Popen[bytes],
+	'process' : subprocess.Popen[Buffer],
 	'metadata' : VideoWriterMetadata
 })
 VideoWriterSet : TypeAlias = Dict[str, VideoWriter]

facefusion/vision.py

@@ -334,10 +334,6 @@ def to_buffer(vision_frame : VisionFrame) -> Buffer:
 	return bytes()
 
 
-def to_strip_buffer(vision_frames : List[VisionFrame]) -> Buffer:
-	return to_buffer(cv2.hconcat(vision_frames))
-
-
 def create_tile_frames(vision_frame : VisionFrame, size : Size) -> Tuple[List[VisionFrame], int, int]:
 	tile_width = size[0] - 2 * size[2]
 	pad_size_top = size[1] + size[2]

tests/test_api_assets.py

@@ -287,6 +287,22 @@ def test_get_asset(test_client : TestClient) -> None:
 
 	assert get_response.status_code == 200
 
+	get_response = test_client.get('/assets/' + asset_ids[0], headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	}, params =
+	{
+		'subject': 'frame',
+		'action': 'capture',
+		'resolution': '1920x1080',
+		'frame_index': [ '0', '1' ]
+	})
+
+	assert get_response.headers.get('content-type').startswith('multipart/mixed; boundary=')
+	assert get_response.content.count('Content-Type: image/jpeg'.encode()) == 1
+
+	assert get_response.status_code == 200
+
 
 def test_delete_assets(test_client : TestClient) -> None:
 	create_session_response = test_client.post('/session', json =