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