auto convert pkl

CHANGELOG.md

@@ -1,12 +1,6 @@
 # 1.1.2 :
 
-+ BREAKING CHANGE : enforce face checkpoint format from pkl to safetensors
-
-Using pkl files to store faces is dangerous from a security point of view. For the same reason that models are now stored in safetensors, We are switching to safetensors for the storage format.
-
-A script with instructions for converting existing pkl files can be found here:
-https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d
-
++ Switch face checkpoint format from pkl to safetensors
 
 ## 1.1.1 :
 

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).

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.
 

preload.py

@@ -8,14 +8,3 @@ def preload(parser: ArgumentParser) -> None:
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Set the log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)",
     )
-
-    print("FACESWAPLAB================================================================")
-    print("BREAKING CHANGE: enforce face checkpoint format from pkl to safetensors\n")
-    print("Using pkl files to store faces is dangerous from a security point of view.")
-    print("For the same reason that models are now stored in safetensors,")
-    print("We are switching to safetensors for the storage format.")
-    print(
-        "A script with instructions for converting existing pkl files can be found here:"
-    )
-    print("https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d")
-    print("==========================================================================")

requirements.txt

@@ -6,3 +6,4 @@ onnxruntime==1.15.0
 opencv-python==4.7.0.72
 pandas
 pydantic==1.10.9
+dill==0.3.6

scripts/faceswaplab_utils/face_utils.py

@@ -8,6 +8,7 @@ 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:
@@ -20,15 +21,32 @@ def save_face(face: Face, filename: str) -> None:
 
 
 def load_face(filename: str) -> Face:
-    face = {}
-    logger.debug("Try to load face from %s", filename)
-    with safe_open(filename, framework="pt", device="cpu") as f:
-        logger.debug("File contains %s keys", f.keys())
-        for k in f.keys():
-            logger.debug("load key %s", k)
-            face[k] = f.get_tensor(k).numpy()
-    logger.debug("face : %s", face)
-    return Face(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]:
@@ -45,4 +63,10 @@ def get_face_checkpoints() -> List[str]:
         scripts.basedir(), "models", "faceswaplab", "faces", "*.safetensors"
     )
     faces = glob.glob(faces_path)
-    return ["None"] + faces
+
+    faces_path = os.path.join(
+        scripts.basedir(), "models", "faceswaplab", "faces", "*.pkl"
+    )
+    faces += glob.glob(faces_path)
+
+    return ["None"] + sorted(faces)