From 31635d369f4c4211d9d7e1e83d52d64ee8517ffd Mon Sep 17 00:00:00 2001 From: Tran Xen <137925069+glucauze@users.noreply.github.com> Date: Sun, 30 Jul 2023 00:55:17 +0200 Subject: [PATCH 1/2] pkl to safetensors --- CHANGELOG.md | 10 ++++ client_api/api_utils.py | 2 +- docs/faq.markdown | 2 +- preload.py | 11 +++++ requirements.txt | 1 - scripts/faceswaplab_globals.py | 2 +- scripts/faceswaplab_ui/faceswaplab_tab.py | 37 +++++++------- .../faceswaplab_unit_settings.py | 6 +-- scripts/faceswaplab_ui/faceswaplab_unit_ui.py | 2 +- scripts/faceswaplab_utils/face_utils.py | 48 +++++++++++++++++++ scripts/faceswaplab_utils/models_utils.py | 17 ------- 11 files changed, 94 insertions(+), 44 deletions(-) create mode 100644 scripts/faceswaplab_utils/face_utils.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e5fc18..dc919ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,13 @@ +# 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 + + ## 1.1.1 : + Add settings for default inpainting prompts diff --git a/client_api/api_utils.py b/client_api/api_utils.py index 24e2159..4369ef8 100644 --- a/client_api/api_utils.py +++ b/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 diff --git a/docs/faq.markdown b/docs/faq.markdown index aab71d0..89eb744 100644 --- a/docs/faq.markdown +++ b/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? diff --git a/preload.py b/preload.py index 5c3eaf5..ba8f628 100644 --- a/preload.py +++ b/preload.py @@ -8,3 +8,14 @@ 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("==========================================================================") diff --git a/requirements.txt b/requirements.txt index 266c9dd..343be35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ cython -dill==0.3.6 ifnude insightface==0.7.3 onnx==1.14.0 diff --git a/scripts/faceswaplab_globals.py b/scripts/faceswaplab_globals.py index 40dbf12..a93249c 100644 --- a/scripts/faceswaplab_globals.py +++ b/scripts/faceswaplab_globals.py @@ -8,7 +8,7 @@ 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. diff --git a/scripts/faceswaplab_ui/faceswaplab_tab.py b/scripts/faceswaplab_ui/faceswaplab_tab.py index 15dfcbb..cea0935 100644 --- a/scripts/faceswaplab_ui/faceswaplab_tab.py +++ b/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 diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py index cd40a11..b2fb3a3 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py @@ -4,12 +4,12 @@ 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 @@ def reference_face(self) -> Optional[Face]: 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) diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py index c1cda99..9516ca1 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py +++ b/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 diff --git a/scripts/faceswaplab_utils/face_utils.py b/scripts/faceswaplab_utils/face_utils.py new file mode 100644 index 0000000..f144344 --- /dev/null +++ b/scripts/faceswaplab_utils/face_utils.py @@ -0,0 +1,48 @@ +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 + + +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: + 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) + + +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) + return ["None"] + faces diff --git a/scripts/faceswaplab_utils/models_utils.py b/scripts/faceswaplab_utils/models_utils.py index 737a173..6bde15e 100644 --- a/scripts/faceswaplab_utils/models_utils.py +++ b/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 From 7538c724ef8787db982bacb006392a7ea10f40b0 Mon Sep 17 00:00:00 2001 From: Tran Xen <137925069+glucauze@users.noreply.github.com> Date: Sun, 30 Jul 2023 13:49:24 +0200 Subject: [PATCH 2/2] auto convert pkl --- CHANGELOG.md | 8 +---- README.md | 12 ++++++- docs/index.markdown | 4 +-- preload.py | 11 ------- requirements.txt | 1 + scripts/faceswaplab_utils/face_utils.py | 44 +++++++++++++++++++------ 6 files changed, 49 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dc919ba..f1a4dd2 100644 --- a/CHANGELOG.md +++ b/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 : diff --git a/README.md b/README.md index c48bdb5..b3423f3 100644 --- a/README.md +++ b/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). diff --git a/docs/index.markdown b/docs/index.markdown index 811e521..53b97cd 100644 --- a/docs/index.markdown +++ b/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. diff --git a/preload.py b/preload.py index ba8f628..5c3eaf5 100644 --- a/preload.py +++ b/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("==========================================================================") diff --git a/requirements.txt b/requirements.txt index 343be35..e48dc85 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file diff --git a/scripts/faceswaplab_utils/face_utils.py b/scripts/faceswaplab_utils/face_utils.py index f144344..e07a6b6 100644 --- a/scripts/faceswaplab_utils/face_utils.py +++ b/scripts/faceswaplab_utils/face_utils.py @@ -8,6 +8,7 @@ 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)