Assets ffmpeg stream upload (#1069)

* ffmpeg sanitize

* fix type

* fix type

* add config

* ChunkQueue -> UploadQueue

* revert assets.py

* move resolve methods to ffmpeg_builder.py

* Refactor ffmpeg.py

* Remove partial import

* improve test

* remove put(None)

* cleanup

* without poll() not working

* ChunkReader -> MediaChunkReader

* improve assert by replacing generic is_file

* naming and cleanup

tests/assert_helper.py

@@ -1,5 +1,7 @@
+import io
 import os
 import tempfile
+from functools import partial
 
 from facefusion.filesystem import are_images, create_directory, is_directory, is_file, remove_directory, resolve_file_paths
 from facefusion.types import JobStatus
@@ -46,3 +48,8 @@ def prepare_test_output_directory() -> bool:
 	remove_directory(test_outputs_directory)
 	create_directory(test_outputs_directory)
 	return is_directory(test_outputs_directory)
+
+
+def create_media_reader(file_path : str) -> partial[bytes]:
+	file_buffer = io.BytesIO(open(file_path, 'rb').read())
+	return partial(file_buffer.read, 1024)

tests/test_api_assets.py

@@ -289,3 +289,37 @@ def test_delete_assets(test_client : TestClient) -> None:
 	})
 
 	assert delete_response.status_code == 404
+
+
+def test_upload_asset_security_strategies(test_client : TestClient) -> None:
+	source_path = get_test_example_file('source.jpg')
+	target_path = get_test_example_file('target-240p.mp4')
+
+	for strategy in [ 'strict', 'moderate' ]:
+		state_manager.init_item('api_security_strategy', strategy)
+
+		create_session_response = test_client.post('/session', json =
+		{
+			'client_version': metadata.get('version')
+		})
+		access_token = create_session_response.json().get('access_token')
+		session_id = session_manager.find_session_id(access_token)
+
+		with open(source_path, 'rb') as source_file, open(target_path, 'rb') as target_file:
+			upload_response = test_client.post('/assets?type=source', headers =
+			{
+				'Authorization': 'Bearer ' + access_token
+			}, files =
+			[
+				('file', ('source.jpg', source_file.read(), 'image/jpeg')),
+				('file', ('target.mp4', target_file.read(), 'video/mp4'))
+			])
+
+		assert upload_response.status_code == 201
+
+		asset_ids = upload_response.json().get('asset_ids')
+
+		assert asset_store.get_asset(session_id, asset_ids[0]).get('media') == 'image'
+		assert asset_store.get_asset(session_id, asset_ids[1]).get('media') == 'video'
+
+	state_manager.init_item('api_security_strategy', 'strict')

tests/test_ffmpeg.py

@@ -7,11 +7,11 @@ import pytest
 import facefusion.ffmpeg
 from facefusion import process_manager, state_manager
 from facefusion.download import conditional_download
-from facefusion.ffmpeg import concat_video, extract_frames, merge_video, read_audio_buffer, replace_audio, restore_audio, spawn_frames
-from facefusion.filesystem import copy_file
+from facefusion.ffmpeg import concat_video, extract_frames, merge_video, read_audio_buffer, replace_audio, restore_audio, sanitize_audio, sanitize_image, sanitize_video, spawn_frames
+from facefusion.filesystem import copy_file, is_audio, is_image, is_video
 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 .assert_helper import get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory
+from .assert_helper import create_media_reader, get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory
 
 
 @pytest.fixture(scope = 'module', autouse = True)
@@ -210,3 +210,37 @@ def test_replace_audio() -> None:
 		clear_temp_directory(state_manager.get_temp_path(), output_path)
 
 	state_manager.init_item('output_audio_encoder', 'aac')
+
+
+def test_sanitize_audio() -> None:
+	audio_path = get_test_example_file('source.mp3')
+
+	output_strict_path = get_test_output_path('output-strict.mp3')
+	output_moderate_path = get_test_output_path('output-moderate.mp3')
+
+	assert sanitize_audio('mp3', create_media_reader(audio_path), output_strict_path, 'strict') is True
+	assert is_audio(output_strict_path) is True
+
+	assert sanitize_audio('mp3', create_media_reader(audio_path), output_moderate_path, 'moderate') is True
+	assert is_audio(output_moderate_path) is True
+
+
+def test_sanitize_image() -> None:
+	source_path = get_test_example_file('source.jpg')
+	output_path = get_test_output_path('output.jpg')
+
+	assert sanitize_image('jpeg', create_media_reader(source_path), output_path) is True
+	assert is_image(output_path) is True
+
+
+def test_sanitize_video() -> None:
+	video_path = get_test_example_file('target-240p.mp4')
+
+	output_strict_path = get_test_output_path('output-strict.mp4')
+	output_moderate_path = get_test_output_path('output-moderate.mp4')
+
+	assert sanitize_video('mp4', create_media_reader(video_path), output_strict_path, 'strict') is True
+	assert is_video(output_strict_path) is True
+
+	assert sanitize_video('mp4', create_media_reader(video_path), output_moderate_path, 'moderate') is True
+	assert is_video(output_moderate_path) is True