From d2f6b7d29d54939bf48a2bbe859b04e48ea5d5b8 Mon Sep 17 00:00:00 2001
From: KYOOKIM <dev.kyoomin@gmail.com>
Date: Sun, 28 Apr 2024 08:15:44 +0900
Subject: [PATCH] Addbuild facemodel api

---
 scripts/reactor_api.py     | 13 ++++++++++++-
 scripts/reactor_swapper.py | 17 +++++++++++++----
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/scripts/reactor_api.py b/scripts/reactor_api.py
index 5fb9969..de766c3 100644
--- a/scripts/reactor_api.py
+++ b/scripts/reactor_api.py
@@ -20,7 +20,7 @@ from modules.api import api
 
 import gradio as gr
 
-from scripts.reactor_swapper import EnhancementOptions, swap_face, DetectionOptions
+from scripts.reactor_swapper import EnhancementOptions, blend_faces, swap_face, DetectionOptions
 from scripts.reactor_logger import logger
 from scripts.reactor_helpers import get_facemodels
 
@@ -163,6 +163,17 @@ def reactor_api(_: gr.Blocks, app: FastAPI):
         facemodels = [os.path.split(model)[1].split(".")[0] for model in get_facemodels()]
         return {"facemodels": facemodels}
 
+    @app.post("/reactor/facemodels")
+    def build_face_model(
+        source_images: list[str] = Body([""],title="Source Face Image List"),
+        name: str = Body("",title="Face Model Name"),
+        compute_method: int = Body(0,title="Compute Method (Mean, Median, Mode)"),
+        shape_check: bool = Body(False,title="Check Embedding Shape"),
+    ):
+        images = [api.decode_base64_to_image(img) for img in source_images]
+        blend_faces(images, name, compute_method, shape_check, is_api=True)
+        return {"facemodels": [os.path.split(model)[1].split(".")[0] for model in get_facemodels()]}
+
 try:
     import modules.script_callbacks as script_callbacks
 
diff --git a/scripts/reactor_swapper.py b/scripts/reactor_swapper.py
index 77b2825..eed0f7d 100644
--- a/scripts/reactor_swapper.py
+++ b/scripts/reactor_swapper.py
@@ -661,13 +661,19 @@ def build_face_model(image: Image.Image, name: str, save_model: bool = True, det
         logger.error(no_face_msg)
         return no_face_msg
 
-def blend_faces(images_list: List, name: str, compute_method: int = 0, shape_check: bool = False):
+def blend_faces(images_list: List, name: str, compute_method: int = 0, shape_check: bool = False, is_api: bool = False):
     faces = []
     embeddings = []
     images: List[Image.Image] = []
-    images, images_names = get_images_from_list(images_list)
+    if not is_api:
+        images, images_names = get_images_from_list(images_list)
+    else: 
+        images = images_list
     for i,image in enumerate(images):
-        logger.status(f"Building Face Model for {images_names[i]}...")
+        if not is_api:
+            logger.status(f"Building Face Model for {images_names[i]}...")
+        else:
+            logger.status(f"Building Face Model for Image {i+1}...")
         face = build_face_model(image,str(i),save_model=False)
         if isinstance(face, str):
             # logger.error(f"No faces found in {images_names[i]}, skipping")
@@ -676,7 +682,10 @@ def blend_faces(images_list: List, name: str, compute_method: int = 0, shape_che
             if i == 0:
                 embedding_shape = face.embedding.shape
             elif face.embedding.shape != embedding_shape:
-                logger.error(f"Embedding Shape Mismatch for {images_names[i]}, skipping")
+                if not is_api:
+                    logger.error(f"Embedding Shape Mismatch for {images_names[i]}, skipping")
+                else:
+                    logger.error(f"Embedding Shape Mismatch for Image {i+1}, skipping")
                 continue
         faces.append(face)
         embeddings.append(face.embedding)