Merge pull request #4 from glucauze/safetensors

v1.1.2 Safetensors

CHANGELOG.md

@@ -1,3 +1,7 @@
+# 1.1.2 :
+
++ Switch face checkpoint format from pkl to safetensors
+
 ## 1.1.1 :
 
 + Add settings for default inpainting prompts

README.md

@@ -1,4 +1,4 @@
-# FaceSwapLab 1.1 for a1111/Vlad
+# FaceSwapLab for a1111/Vlad
 
 Please read the documentation here : https://glucauze.github.io/sd-webui-faceswaplab/
 
@@ -10,6 +10,16 @@ Some key features include the ability to reuse faces via checkpoints, multiple f
 
 While FaceSwapLab is still under development, it has reached a good level of stability. This makes it a reliable tool for those who are interested in face-swapping within the Stable Diffusion environment. As with all projects of this type, it’s expected to improve and evolve over time.
 
+## Disclaimer and license
+
+In short:
+
++ **Ethical Guideline:** This extension should not be forked to create a public, easy way to circumvent NSFW filtering.
++ **License:** This software is distributed under the terms of the GNU Affero General Public License (AGPL), version 3 or later.
++ **Model License:** This software uses InsightFace's pre-trained models, which are available for non-commercial research purposes only.
+
+More on this here : https://glucauze.github.io/sd-webui-faceswaplab/ 
+
 ### Features
 
 + **Face Unit Concept**: Similar to controlNet, the program introduces the concept of a face unit. You can configure up to 10 units (3 units are the default setting) in the program settings (sd).

client_api/api_utils.py

@@ -28,7 +28,7 @@ class FaceSwapUnit(BaseModel):
     # The checkpoint file
     source_face: str = Field(
         description="face checkpoint (from models/faceswaplab/faces)",
-        examples=["my_face.pkl"],
+        examples=["my_face.safetensors"],
         default=None,
     )
     # base64 batch source images

docs/faq.markdown

@@ -112,7 +112,7 @@ A face checkpoint is a saved embedding of a face, generated from multiple images
 
 The primary advantage of face checkpoints is their size. An embedding is only around 2KB, meaning it's lightweight and can be reused later without requiring additional calculations.
 
-Face checkpoints are saved as `.pkl` files. Please be aware that exchanging `.pkl` files carries potential security risks. These files, by default, are not secure and could potentially execute malicious code when opened. Therefore, extreme caution should be exercised when sharing or receiving this type of file.
+Face checkpoints are saved as `.safetensors` files. Please be aware that exchanging `.safetensors` files carries potential security risks. These files, by default, are not secure and could potentially execute malicious code when opened. Therefore, extreme caution should be exercised when sharing or receiving this type of file.
 
 #### How is similarity determined?
 

docs/index.markdown

@@ -20,11 +20,11 @@ While FaceSwapLab is still under development, it has reached a good level of sta
 
 In short:
 
-+ **Ethical Consideration:** This extension should not be forked to create a public, easy way to circumvent NSFW filtering.
++ **Ethical Guideline:** This extension should not be forked to create a public, easy way to circumvent NSFW filtering.
 + **License:** This software is distributed under the terms of the GNU Affero General Public License (AGPL), version 3 or later.
 + **Model License:** This software uses InsightFace's pre-trained models, which are available for non-commercial research purposes only.
 
-### Ethical Perspective
+### Ethical Guideline
 
 This extension is **not intended to facilitate the creation of not safe for work (NSFW) or non-consensual deepfake content**. Its purpose is to bring consistency to image creation, making it easier to repair existing images, or bring characters back to life.
 

requirements.txt

@@ -1,5 +1,4 @@
 cython
-dill==0.3.6
 ifnude
 insightface==0.7.3
 onnx==1.14.0
@@ -7,3 +6,4 @@ onnxruntime==1.15.0
 opencv-python==4.7.0.72
 pandas
 pydantic==1.10.9
+dill==0.3.6

scripts/faceswaplab_globals.py

@@ -8,7 +8,7 @@ REFERENCE_PATH = os.path.join(
     scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
 )
 
-VERSION_FLAG: str = "v1.1.1"
+VERSION_FLAG: str = "v1.1.2"
 EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
 
 # The NSFW score threshold. If any part of the image has a score greater than this threshold, the image will be considered NSFW.

scripts/faceswaplab_ui/faceswaplab_tab.py

@@ -1,14 +1,12 @@
 import os
 from pprint import pformat, pprint
-
-import dill as pickle
+from scripts.faceswaplab_utils import face_utils
 import gradio as gr
 import modules.scripts as scripts
 import onnx
 import pandas as pd
 from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui
 from scripts.faceswaplab_ui.faceswaplab_postprocessing_ui import postprocessing_ui
-from insightface.app.common import Face
 from modules import scripts
 from PIL import Image
 from modules.shared import opts
@@ -128,10 +126,17 @@ def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> Optional[st
 
 
 def sanitize_name(name: str) -> str:
-    logger.debug(f"Sanitize name {name}")
+    """
+    Sanitize the input name by removing special characters and replacing spaces with underscores.
+
+    Parameters:
+        name (str): The input name to be sanitized.
+
+    Returns:
+        str: The sanitized name with special characters removed and spaces replaced by underscores.
+    """
     name = re.sub("[^A-Za-z0-9_. ]+", "", name)
     name = name.replace(" ", "_")
-    logger.debug(f"Sanitized name {name[:255]}")
     return name[:255]
 
 
@@ -185,25 +190,19 @@ def build_face_checkpoint_and_save(
                 ),
             )
 
-            file_path = os.path.join(faces_path, f"{name}.pkl")
+            file_path = os.path.join(faces_path, f"{name}.safetensors")
             file_number = 1
             while os.path.exists(file_path):
-                file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl")
+                file_path = os.path.join(
+                    faces_path, f"{name}_{file_number}.safetensors"
+                )
                 file_number += 1
             result_image.save(file_path + ".png")
-            with open(file_path, "wb") as file:
-                pickle.dump(
-                    {
-                        "embedding": blended_face.embedding,
-                        "gender": blended_face.gender,
-                        "age": blended_face.age,
-                    },
-                    file,
-                )
+
+            face_utils.save_face(filename=file_path, face=blended_face)
             try:
-                with open(file_path, "rb") as file:
-                    data = Face(pickle.load(file))
-                    print(data)
+                data = face_utils.load_face(filename=file_path)
+                print(data)
             except Exception as e:
                 print(e)
             return result_image

scripts/faceswaplab_ui/faceswaplab_unit_settings.py

@@ -4,12 +4,12 @@ import base64
 import io
 from dataclasses import dataclass, fields
 from typing import Any, List, Optional, Set, Union
-import dill as pickle
 import gradio as gr
 from insightface.app.common import Face
 from PIL import Image
 from scripts.faceswaplab_utils.imgutils import pil_to_cv2
 from scripts.faceswaplab_utils.faceswaplab_logging import logger
+from scripts.faceswaplab_utils import face_utils
 
 
 @dataclass
@@ -94,8 +94,8 @@ class FaceSwapUnitSettings:
             if self.source_face and self.source_face != "None":
                 with open(self.source_face, "rb") as file:
                     try:
-                        logger.info(f"loading pickle {file.name}")
-                        face = Face(pickle.load(file))
+                        logger.info(f"loading face {file.name}")
+                        face = face_utils.load_face(file.name)
                         self._reference_face = face
                     except Exception as e:
                         logger.error("Failed to load checkpoint  : %s", e)

scripts/faceswaplab_ui/faceswaplab_unit_ui.py

@@ -1,5 +1,5 @@
 from typing import List
-from scripts.faceswaplab_utils.models_utils import get_face_checkpoints
+from scripts.faceswaplab_utils.face_utils import get_face_checkpoints
 import gradio as gr
 
 

scripts/faceswaplab_utils/face_utils.py

@@ -0,0 +1,72 @@
+import glob
+import os
+from typing import List
+from insightface.app.common import Face
+from safetensors.torch import save_file, safe_open
+import torch
+
+import modules.scripts as scripts
+from modules import scripts
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
+import dill as pickle  # will be removed in future versions
+
+
+def save_face(face: Face, filename: str) -> None:
+    tensors = {
+        "embedding": torch.tensor(face["embedding"]),
+        "gender": torch.tensor(face["gender"]),
+        "age": torch.tensor(face["age"]),
+    }
+    save_file(tensors, filename)
+
+
+def load_face(filename: str) -> Face:
+    if filename.endswith(".pkl"):
+        logger.warning(
+            "Pkl files for faces are deprecated to enhance safety, they will be unsupported in future versions."
+        )
+        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
+
+    elif filename.endswith(".safetensors"):
+        face = {}
+        with safe_open(filename, framework="pt", device="cpu") as f:
+            for k in f.keys():
+                logger.debug("load key %s", k)
+                face[k] = f.get_tensor(k).numpy()
+        return Face(face)
+
+    raise NotImplementedError("Unknown file type, face extraction not implemented")
+
+
+def get_face_checkpoints() -> List[str]:
+    """
+    Retrieve a list of face checkpoint paths.
+
+    This function searches for face files with the extension ".safetensors" in the specified directory and returns a list
+    containing the paths of those files.
+
+    Returns:
+        list: A list of face paths, including the string "None" as the first element.
+    """
+    faces_path = os.path.join(
+        scripts.basedir(), "models", "faceswaplab", "faces", "*.safetensors"
+    )
+    faces = glob.glob(faces_path)
+
+    faces_path = os.path.join(
+        scripts.basedir(), "models", "faceswaplab", "faces", "*.pkl"
+    )
+    faces += glob.glob(faces_path)
+
+    return ["None"] + sorted(faces)

scripts/faceswaplab_utils/models_utils.py

@@ -43,20 +43,3 @@ def get_current_model() -> str:
             "No faceswap model found. Please add it to the faceswaplab directory."
         )
     return model
-
-
-def get_face_checkpoints() -> List[str]:
-    """
-    Retrieve a list of face checkpoint paths.
-
-    This function searches for face files with the extension ".pkl" in the specified directory and returns a list
-    containing the paths of those files.
-
-    Returns:
-        list: A list of face paths, including the string "None" as the first element.
-    """
-    faces_path = os.path.join(
-        scripts.basedir(), "models", "faceswaplab", "faces", "*.pkl"
-    )
-    faces = glob.glob(faces_path)
-    return ["None"] + faces