Merge pull request #3 from jahangir091/facefusion-api

Facefusion api

extensions/sd-face-fusion-extension/facefusion/core.py

@@ -38,7 +38,7 @@ def get_argument_parser():
 	# execution
 	group_execution = program.add_argument_group('execution')
 	group_execution.add_argument('--execution-providers', help = wording.get('execution_providers_help').format(choices = 'cpu'), dest = 'execution_providers', default = [ 'cuda' ], choices = encode_execution_providers(onnxruntime.get_available_providers()), nargs = '+')
-	group_execution.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), dest = 'execution_thread_count', type = int, default = 1)
+	group_execution.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), dest = 'execution_thread_count', type = int, default = 10)
 	group_execution.add_argument('--execution-queue-count', help = wording.get('execution_queue_count_help'), dest = 'execution_queue_count', type = int, default = 1)
 	group_execution.add_argument('--max-memory', help=wording.get('max_memory_help'), dest='max_memory', type = int)
 	# face recognition
@@ -61,7 +61,7 @@ def get_argument_parser():
 	group_output = program.add_argument_group('output creation')
 	group_output.add_argument('--output-image-quality', help=wording.get('output_image_quality_help'), dest = 'output_image_quality', type = int, default = 80, choices = range(101), metavar = '[0-100]')
 	group_output.add_argument('--output-video-encoder', help = wording.get('output_video_encoder_help'), dest = 'output_video_encoder', default = 'libx264', choices = facefusion.choices.output_video_encoders)
-	group_output.add_argument('--output-video-quality', help = wording.get('output_video_quality_help'), dest = 'output_video_quality', type = int, default = 100, choices = range(101), metavar = '[0-100]')
+	group_output.add_argument('--output-video-quality', help = wording.get('output_video_quality_help'), dest = 'output_video_quality', type = int, default = 80, choices = range(101), metavar = '[0-100]')
 	group_output.add_argument('--keep-fps', help = wording.get('keep_fps_help'), dest = 'keep_fps', default = True, action = 'store_true')
 	group_output.add_argument('--skip-audio', help = wording.get('skip_audio_help'), dest = 'skip_audio', action = 'store_true')
 	# frame processors
@@ -86,7 +86,7 @@ def apply_args(program : ArgumentParser) -> None:
 	# general
 	facefusion.globals.source_path = args.source_path
 	facefusion.globals.target_path = args.target_path
-	facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_path, facefusion.globals.target_path, "/siam/output")
+	facefusion.globals.output_path = normalize_output_path(facefusion.globals.source_path, facefusion.globals.target_path, "/home/sduser/stable-diffusion-webui/output")
 	# misc
 	facefusion.globals.skip_download = args.skip_download
 	facefusion.globals.headless = args.headless

extensions/sd-face-fusion-extension/facefusion/uis/core.py

@@ -127,4 +127,4 @@ def get_theme() -> gradio.Theme:
 def get_css() -> str:
 	fixes_css_path = resolve_relative_path('uis/assets/fixes.css')
 	overrides_css_path = resolve_relative_path('uis/assets/overrides.css')
-	return open(fixes_css_path, 'r').read() + open(overrides_css_path, 'r').read()
+	return open(fixes_css_path, 'r').read() + open(overrides_css_path, 'r').read()

extensions/sd-face-fusion-extension/scripts/api.py

@@ -1,47 +1,145 @@
 from fastapi import FastAPI, Body, HTTPException
 from fastapi.responses import RedirectResponse, FileResponse
-from fastapi import File, UploadFile
+from fastapi import File, UploadFile, Form
 
-from typing import List
+from typing import List, Optional, Tuple
 import gradio as gr
 
 import facefusion.globals
 from facefusion.utilities import is_image, is_video
 from facefusion.uis.components.output import predict, start
-
+
+from facefusion.uis.typing import Update
 import os
+import shutil
+import time
 
 
 def facefusion_api(_: gr.Blocks, app: FastAPI):
+    @app.post('/facefusion/video')
+    async def facefusion_video(
+        source_image: UploadFile = File(),
+        target_video: UploadFile = File(None),
+        target_video_name: str = Form("", title='target video name including extension')
+    ):
+        #  Source image must be an image
+        start_time = time.time()
+
+        curDir = os.getcwd()
+        output_dir = curDir + "/output"
+        output_video_dir = output_dir + "/videos"
+
+        # Testing only..........this two directory should be update.
+        input_dir = output_dir + "/input/"
+        template_dir = output_dir + "/template/"
+
+        source_file_location = input_dir + source_image.filename
+        save_file(source_image, source_file_location)
+        facefusion.globals.source_path = source_file_location
+
+        if target_video_name == "":
+            target_file_location = input_dir + target_video.filename
+            save_file(target_video, target_file_location)
+            facefusion.globals.target_path = target_file_location
+        else:
+            target_file_location = template_dir + target_video_name
+            if is_video(target_file_location):
+                facefusion.globals.target_path = target_file_location
+            else:
+                raise HTTPException(status_code=404, detail="Target Video not found.")
+
+        isTemplate = True 
+        if target_video_name == "":
+            isTemplate = False
+
+        video_path = process_face_fusion(output_video_dir, isTemplate)
+        if video_path == None:
+            raise HTTPException(status_code=502, detail="Couldn't process your request")
+
+        end_time = time.time()
+        server_process_time = end_time - start_time
+        return {
+            "server_process_time": server_process_time, 
+            "image_url": "file=" + video_path
+        }
+
+
+
     @app.post('/facefusion/image')
     async def facefusion_image(
-        source_image: str,
-        target_file: str
+        source_image: UploadFile = File(),
+        target_image: UploadFile = File(None),
+        target_image_name: str = Form("", title='target image name including extension')
     ):
-        if is_image(source_image):
-            facefusion.globals.source_path = source_image
-        else:
-            facefusion.globals.source_path = None
-            raise HTTPException(status_code=404, detail="Source Image not found")
-
-        if is_image(target_file) or is_video(target_file):
-            facefusion.globals.target_path = target_file
+        #  Source image must be an image
+        start_time = time.time()
+
+        curDir = os.getcwd()
+        output_dir = curDir + "/output"
+
+        # Testing only..........this two directory should be update.
+        input_dir = output_dir + "/input/"
+        template_dir = output_dir + "/template/"
+
+        source_file_location = input_dir + source_image.filename
+        save_file(source_image, source_file_location)
+        facefusion.globals.source_path = source_file_location
+
+        if target_image_name == "":
+            target_file_location = input_dir + target_image.filename
+            save_file(target_image, target_file_location)
+            facefusion.globals.target_path = target_file_location
         else:
-            facefusion.globals.target_path = None
-            raise HTTPException(status_code=404, detail="Target not found")
+            target_file_location = template_dir + target_image_name
+            if is_image(target_file_location):
+                facefusion.globals.target_path = target_file_location
+            else:
+                raise HTTPException(status_code=404, detail="Target Image not found.")
 
+        isTemplate = True 
+        if target_image_name == "":
+            isTemplate = False
+
+        image_path = process_face_fusion(output_dir, isTemplate)
+        if image_path == None:
+            raise HTTPException(status_code=502, detail="Couldn't process your request")
 
-        return start("/home/siam/output")
+        end_time = time.time()
+        server_process_time = end_time - start_time
+        return {
+            "server_process_time": server_process_time, 
+            "video_url": "file=" + image_path
+        }
+
+
+
+    def save_file(file: UploadFile, path: str):
+        with open(path, "wb+") as file_object:
+            shutil.copyfileobj(file.file, file_object)  
+
+    # isTemplate = True means this target data belongs to us. And should not be deleted
+    def process_face_fusion(output_dir: str, isTemplate: bool):
+        process_response = start(output_dir)
+
+        # extracting relative file url
+        file_path = None
+        for response in process_response:
+            if response["value"] != None:
+                file_path = response["value"] 
+                if file_path != None:
+                    file_path = os.path.relpath(file_path, os.getcwd())
+
+        # removing users data
+        os.remove(facefusion.globals.source_path)
+        facefusion.globals.source_path = None
+
+        # only delete data if its not template ie, users data
+        if not isTemplate:
+            os.remove(facefusion.globals.target_path)
+            facefusion.globals.target_path = None
 
-        # predict("/home/siam/output")
+        return file_path
 
-        # if is_image(facefusion.globals.output_path):
-        #     return {
-        #         "image": api.encode_pil_to_base64(outImg).decode("utf-8")
-        #     }
-        #     return FileResponse(facefusion.globals.output_path)
-        # else:
-        #     raise HTTPException(status_code=500, detail="Couldn't process your request")