CalvinBackup/sd-webui-faceswaplab
1
0
Fork 0
Code Issues 56 Pull Requests 4 Actions Packages Projects Releases 8 Wiki Activity

add pre-commit hooks configuration

Browse Source
This commit is contained in:
Tran Xen
2023-07-28 18:25:28 +02:00
parent 8577d0186d
commit 5d4a29ff1e
33 changed files with 1674 additions and 820 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/faceswaplab_utils/faceswaplab_logging.py
+8 -3
View File
@@ -4,6 +4,7 @@ import sys
from modules import shared
from PIL import Image
class ColoredFormatter(logging.Formatter):
COLORS = {
"DEBUG": "\033[0;36m", # CYAN
@@ -40,14 +41,18 @@ loglevel = getattr(logging, loglevel_string.upper(), "INFO")
logger.setLevel(loglevel)
import tempfile
if logger.getEffectiveLevel() <= logging.DEBUG :
if logger.getEffectiveLevel() <= logging.DEBUG:
DEBUG_DIR = tempfile.mkdtemp()
def save_img_debug(img: Image.Image, message: str, *opts):
if logger.getEffectiveLevel() <= logging.DEBUG:
with tempfile.NamedTemporaryFile(dir=DEBUG_DIR, delete=False, suffix=".png") as temp_file:
with tempfile.NamedTemporaryFile(
dir=DEBUG_DIR, delete=False, suffix=".png"
) as temp_file:
img_path = temp_file.name
img.save(img_path)
message_with_link = f"{message}\nImage: file://{img_path}"
logger.debug(message_with_link, *opts)
logger.debug(message_with_link, *opts)
scripts/faceswaplab_utils/imgutils.py
+40 -29
View File
@@ -1,6 +1,6 @@
import io
from typing import Optional
from PIL import Image, ImageChops, ImageOps,ImageFilter
from PIL import Image, ImageChops, ImageOps, ImageFilter
import cv2
import numpy as np
from math import isqrt, ceil
@@ -10,6 +10,7 @@ from scripts.faceswaplab_globals import NSFW_SCORE
from modules import processing
import base64
def check_against_nsfw(img):
shapes = []
chunks = detect(img)
@@ -17,6 +18,7 @@ def check_against_nsfw(img):
shapes.append(chunk["score"] > NSFW_SCORE)
return any(shapes)
def pil_to_cv2(pil_img):
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
@@ -24,6 +26,7 @@ def pil_to_cv2(pil_img):
def cv2_to_pil(cv2_img):
return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB))
def torch_to_pil(images):
"""
Convert a numpy image or a batch of images to a PIL image.
@@ -49,35 +52,38 @@ def pil_to_torch(pil_images):
torch_image = torch.from_numpy(numpy_image).permute(2, 0, 1)
return torch_image
from collections import Counter
def create_square_image(image_list):
"""
Creates a square image by combining multiple images in a grid pattern.
Args:
image_list (list): List of PIL Image objects to be combined.
Returns:
PIL Image object: The resulting square image.
None: If the image_list is empty or contains only one image.
"""
# Count the occurrences of each image size in the image_list
size_counter = Counter(image.size for image in image_list)
# Get the most common image size (size with the highest count)
common_size = size_counter.most_common(1)[0][0]
# Filter the image_list to include only images with the common size
image_list = [image for image in image_list if image.size == common_size]
# Get the dimensions (width and height) of the common size
size = common_size
# If there are more than one image in the image_list
if len(image_list) > 1:
num_images = len(image_list)
# Calculate the number of rows and columns for the grid
rows = isqrt(num_images)
cols = ceil(num_images / rows)
@@ -97,10 +103,11 @@ def create_square_image(image_list):
# Return the resulting square image
return square_image
# Return None if there are no images or only one image in the image_list
return None
def create_mask(image, box_coords):
width, height = image.size
mask = Image.new("L", (width, height), 255)
@@ -113,43 +120,47 @@ def create_mask(image, box_coords):
mask.putpixel((x, y), 0)
return mask
def apply_mask(img : Image.Image,p : processing.StableDiffusionProcessing, batch_index : int) -> Image.Image :
def apply_mask(
img: Image.Image, p: processing.StableDiffusionProcessing, batch_index: int
) -> Image.Image:
"""
Apply mask overlay and color correction to an image if enabled
Args:
img: PIL Image objects.
p : The processing object
batch_index : the batch index
Returns:
PIL Image object
"""
if isinstance(p, processing.StableDiffusionProcessingImg2Img) :
if p.inpaint_full_res :
if isinstance(p, processing.StableDiffusionProcessingImg2Img):
if p.inpaint_full_res:
overlays = p.overlay_images
if overlays is None or batch_index >= len(overlays):
return img
overlay : Image.Image = overlays[batch_index]
overlay = overlay.resize((img.size), resample= Image.Resampling.LANCZOS)
overlay: Image.Image = overlays[batch_index]
overlay = overlay.resize((img.size), resample=Image.Resampling.LANCZOS)
img = img.copy()
img.paste(overlay, (0, 0), overlay)
return img
return img
img = processing.apply_overlay(img, p.paste_to, batch_index, p.overlay_images)
if p.color_corrections is not None and batch_index < len(p.color_corrections):
img = processing.apply_color_correction(p.color_corrections[batch_index], img)
img = processing.apply_color_correction(
p.color_corrections[batch_index], img
)
return img
def prepare_mask(
mask: Image.Image, p: processing.StableDiffusionProcessing
) -> Image.Image:
"""
Prepare an image mask for the inpainting process. (This comes from controlnet)
This function takes as input a PIL Image object and an instance of the
This function takes as input a PIL Image object and an instance of the
StableDiffusionProcessing class, and performs the following steps to prepare the mask:
1. Convert the mask to grayscale (mode "L").
@@ -160,26 +171,26 @@ def prepare_mask(
Args:
mask (Image.Image): The input mask as a PIL Image object.
p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class
p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class
containing the processing parameters.
Returns:
mask (Image.Image): The prepared mask as a PIL Image object.
"""
mask = mask.convert("L")
#FIXME : Properly fix blur
# FIXME : Properly fix blur
# if getattr(p, "mask_blur", 0) > 0:
# mask = mask.filter(ImageFilter.GaussianBlur(p.mask_blur))
return mask
def base64_to_pil(base64str : Optional[str]) -> Optional[Image.Image] :
if base64str is None :
def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]:
if base64str is None:
return None
if 'base64,' in base64str: # check if the base64 string has a data URL scheme
base64_data = base64str.split('base64,')[-1]
if "base64," in base64str: # check if the base64 string has a data URL scheme
base64_data = base64str.split("base64,")[-1]
img_bytes = base64.b64decode(base64_data)
else:
# if no data URL scheme, just decode
img_bytes = base64.b64decode(base64str)
return Image.open(io.BytesIO(img_bytes))
scripts/faceswaplab_utils/models_utils.py
+11 -5
View File
@@ -1,4 +1,3 @@
import glob
import os
import modules.scripts as scripts
@@ -7,6 +6,7 @@ from scripts.faceswaplab_globals import EXTENSION_PATH
from modules.shared import opts
from scripts.faceswaplab_utils.faceswaplab_logging import logger
def get_models():
"""
Retrieve a list of swap model files.
@@ -29,17 +29,21 @@ def get_models():
return models
def get_current_model() -> str :
def get_current_model() -> str:
model = opts.data.get("faceswaplab_model", None)
if model is None :
if model is None:
models = get_models()
model = models[0] if len(models) else None
logger.info("Try to use model : %s", model)
if not os.path.isfile(model):
logger.error("The model %s cannot be found or loaded", model)
raise FileNotFoundError("No faceswap model found. Please add it to the faceswaplab directory.")
raise FileNotFoundError(
"No faceswap model found. Please add it to the faceswaplab directory."
)
return model
def get_face_checkpoints():
"""
Retrieve a list of face checkpoint paths.
@@ -50,6 +54,8 @@ def get_face_checkpoints():
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_path = os.path.join(
scripts.basedir(), "models", "faceswaplab", "faces", "*.pkl"
)
faces = glob.glob(faces_path)
return ["None"] + faces