Merge pull request #96 from glucauze/v1.2.7

v1.2.7 

    remove dill
    add warnings in model install and checking
    add gender selection in build

install.py

@@ -62,5 +62,11 @@ def check_install() -> None:
 
 import timeit
 
+try:
     check_time = timeit.timeit(check_install, number=1)
     print(check_time)
+except Exception as e:
+    print("FaceswapLab install failed", e)
+    print(
+        "You can try to install dependencies manually by activating venv and installing requirements.txt or requirements-gpu.txt"
+    )

requirements-gpu.txt

@@ -1,5 +1,4 @@
 cython
-dill
 ifnude
 insightface==0.7.3
 onnx>=1.14.0

requirements.txt

@@ -1,6 +1,5 @@
 protobuf>=3.20.2
 cython
-dill
 ifnude
 insightface==0.7.3
 onnx>=1.14.0

scripts/configure.py

@@ -25,6 +25,7 @@ def check_configuration() -> None:
     model_path = os.path.join(models_dir, model_name)
 
     def download(url: str, path: str) -> None:
+        try:
             request = urllib.request.urlopen(url)
             total = int(request.headers.get("Content-Length", 0))
             with tqdm(
@@ -41,6 +42,10 @@ def check_configuration() -> None:
                         block_size
                     ),
                 )
+        except:
+            logger.error(
+                "Failed to download inswapper_128.onnx model, please download it manually and put it in the (<sdwebui>/models/faceswaplab/inswapper_128.onnx) directory"
+            )
 
     os.makedirs(models_dir, exist_ok=True)
     os.makedirs(faces_dir, exist_ok=True)

scripts/faceswaplab_globals.py

@@ -16,7 +16,7 @@ REFERENCE_PATH = os.path.join(
 )
 
 # Defining the version flag for the application
-VERSION_FLAG: str = "v1.2.5"
+VERSION_FLAG: str = "v1.2.7"
 # Defining the path for 'sd-webui-faceswaplab' inside the 'extensions' directory
 EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
 

scripts/faceswaplab_swapping/face_checkpoints.py

@@ -14,7 +14,6 @@ from scripts.faceswaplab_utils import imgutils
 from scripts.faceswaplab_utils.models_utils import get_swap_models
 import traceback
 
-import dill as pickle  # will be removed in future versions
 from scripts.faceswaplab_swapping import swapper
 from pprint import pformat
 import re
@@ -40,6 +39,7 @@ def sanitize_name(name: str) -> str:
 def build_face_checkpoint_and_save(
     images: List[PILImage],
     name: str,
+    gender: Gender = Gender.AUTO,
     overwrite: bool = False,
     path: Optional[str] = None,
 ) -> Optional[PILImage]:
@@ -65,7 +65,7 @@ def build_face_checkpoint_and_save(
             logger.error("No source faces found")
             return None
 
-        blended_face: Optional[Face] = swapper.blend_faces(faces)
+        blended_face: Optional[Face] = swapper.blend_faces(faces, gender=gender)
         preview_path = os.path.join(
             scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
         )
@@ -174,20 +174,13 @@ def load_face(name: str) -> Optional[Face]:
 
     if filename.endswith(".pkl"):
         logger.warning(
-            "Pkl files for faces are deprecated to enhance safety, they will be unsupported in future versions."
+            "Pkl files for faces are deprecated to enhance safety, you need to convert them"
         )
         logger.warning("The file will be converted to .safetensors")
         logger.warning(
             "You can also use this script https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d"
         )
-        with open(filename, "rb") as file:
-            logger.info("Load pkl")
-            face = Face(pickle.load(file))
-            logger.warning(
-                "Convert to safetensors, you can remove the pkl version once you have ensured that the safetensor is working"
-            )
-            save_face(face, filename.replace(".pkl", ".safetensors"))
-        return face
+        return None
 
     elif filename.endswith(".safetensors"):
         face = {}

scripts/faceswaplab_swapping/swapper.py

@@ -33,7 +33,7 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
 )
 from scripts.faceswaplab_utils.models_utils import get_current_swap_model
-from scripts.faceswaplab_utils.typing import CV2ImgU8, PILImage, Face
+from scripts.faceswaplab_utils.typing import CV2ImgU8, Gender, PILImage, Face
 from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion
 from modules import shared
 import onnxruntime
@@ -559,7 +559,7 @@ def get_faces_from_img_files(images: List[PILImage]) -> List[Face]:
     return faces
 
 
-def blend_faces(faces: List[Face]) -> Optional[Face]:
+def blend_faces(faces: List[Face], gender: Gender = Gender.AUTO) -> Optional[Face]:
     """
     Blends the embeddings of multiple faces into a single face.
 
@@ -587,10 +587,19 @@ def blend_faces(faces: List[Face]) -> Optional[Face]:
         # Compute the mean of all embeddings
         blended_embedding = np.mean(embeddings, axis=0)
 
+        if gender == Gender.AUTO:
+            int_gender: int = faces[0].gender  # type: ignore
+        else:
+            int_gender: int = gender.value
+
+        assert -1 < int_gender < 2, "wrong gender"
+
+        logger.info("Int Gender : %s", int_gender)
+
         # Create a new Face object using the properties of the first face in the list
         # Assign the blended embedding to the blended Face object
         blended = ISFace(
-            embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age
+            embedding=blended_embedding, gender=int_gender, age=faces[0].age
         )
 
         return blended

scripts/faceswaplab_ui/faceswaplab_tab.py

@@ -137,7 +137,7 @@ def analyse_faces(image: PILImage, det_threshold: float = 0.5) -> Optional[str]:
 
 
 def build_face_checkpoint_and_save(
-    batch_files: List[gr.File], name: str, overwrite: bool
+    batch_files: List[gr.File], name: str, str_gender: str, overwrite: bool
 ) -> PILImage:
     """
     Builds a face checkpoint using the provided image files, performs face swapping,
@@ -156,10 +156,13 @@ def build_face_checkpoint_and_save(
         if not batch_files:
             logger.error("No face found")
             return None  # type: ignore (Optional not really supported by old gradio)
+
+        gender = getattr(Gender, str_gender)
+        logger.info("Choosen gender : %s", gender)
         images: list[PILImage] = [Image.open(file.name) for file in batch_files]  # type: ignore
         preview_image: PILImage | None = (
             face_checkpoints.build_face_checkpoint_and_save(
-                images=images, name=name, overwrite=overwrite
+                images=images, name=name, overwrite=overwrite, gender=gender
             )
         )
     except Exception as e:
@@ -266,6 +269,13 @@ def tools_ui() -> None:
                 label="Name of the character",
                 elem_id="faceswaplab_build_character_name",
             )
+            build_gender = gr.Dropdown(
+                value=Gender.AUTO.name,
+                choices=[e.name for e in Gender],
+                placeholder="Gender of the character",
+                label="Gender of the character",
+                elem_id="faceswaplab_build_character_gender",
+            )
             build_overwrite = gr.Checkbox(
                 False,
                 placeholder="overwrite",
@@ -387,7 +397,7 @@ def tools_ui() -> None:
     compare_btn.click(compare, inputs=[img1, img2], outputs=[compare_result_text])
     generate_checkpoint_btn.click(
         build_face_checkpoint_and_save,
-        inputs=[build_batch_files, build_name, build_overwrite],
+        inputs=[build_batch_files, build_name, build_gender, build_overwrite],
         outputs=[preview],
     )
     extract_btn.click(

scripts/faceswaplab_utils/models_utils.py

@@ -72,10 +72,16 @@ def get_current_swap_model() -> str:
         models = get_swap_models()
         model = models[0] if len(models) else None
     logger.info("Try to use model : %s", model)
-    if not os.path.isfile(model):  # type: ignore
+    try:
+        if not model or not os.path.isfile(model):  # type: ignore
             logger.error("The model %s cannot be found or loaded", model)
             raise FileNotFoundError(
-            "No faceswap model found. Please add it to the faceswaplab directory."
+                "No faceswap model found. Please add it to the faceswaplab directory. Ensure the model is in the proper directory (<sdwebui>/models/faceswaplab/inswapper_128.onnx)"
+            )
+    except:
+        raise FileNotFoundError(
+            "No faceswap model found. Please add it to the faceswaplab directory. Ensure the model is in the proper directory (<sdwebui>/models/faceswaplab/inswapper_128.onnx)"
         )
+
     assert model is not None
     return model

scripts/faceswaplab_utils/typing.py

@@ -3,8 +3,15 @@ from numpy import uint8
 from insightface.app.common import Face as IFace
 from PIL import Image
 import numpy as np
+from enum import Enum
 
 PILImage = Image.Image
 CV2ImgU8 = np.ndarray[int, np.dtype[uint8]]
 Face = IFace
 BoxCoords = Tuple[int, int, int, int]
+
+
+class Gender(Enum):
+    AUTO = -1
+    FEMALE = 0
+    MALE = 1