1580158
/
sd-webui-faceswaplab
mirror of https://github.com/glucauze/sd-webui-faceswaplab
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add api for face building, add tests

pull/19/head
Tran Xen 2023-08-03 15:25:44 +02:00
parent 4533750c49
commit 02d88bac91
8 changed files with 169 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
client_api/api_utils.py
View File

@ -187,7 +187,7 @@ class FaceSwapRequest(BaseModel):
default=None,
)
units: List[FaceSwapUnit]
postprocessing: Optional[PostProcessingOptions]
postprocessing: Optional[PostProcessingOptions] = None
class FaceSwapResponse(BaseModel):

22
client_api/faceswaplab_api_example.py
View File

@ -1,7 +1,9 @@
from typing import List
import requests
from api_utils import (
FaceSwapUnit,
InswappperOptions,
base64_to_safetensors,
pil_to_base64,
PostProcessingOptions,
InpaintingWhen,
@ -98,12 +100,30 @@ for img in response.pil_images:
img.show()
#############################
# Build checkpoint
source_images: List[str] = [
pil_to_base64("../references/man.png"),
pil_to_base64("../references/woman.png"),
]
result = requests.post(
url=f"{address}/faceswaplab/build",
json=source_images,
headers={"Content-Type": "application/json; charset=utf-8"},
)
base64_to_safetensors(result.json(), output_path="test.safetensors")
#############################
# FaceSwap with local safetensors
# First face unit :
unit1 = FaceSwapUnit(
source_face=safetensors_to_base64("test.safetensors"),
source_face=safetensors_to_base64(
"test.safetensors"
), # convert the checkpoint to base64
faces_index=(0,), # Replace first face
swapping_options=InswappperOptions(
face_restorer_name="CodeFormer",

BIN
client_api/test.safetensors
View File

Binary file not shown.

24
scripts/faceswaplab_api/faceswaplab_api.py
View File

@ -1,3 +1,4 @@
import tempfile
from PIL import Image
import numpy as np
from fastapi import FastAPI
@ -17,6 +18,9 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
PostProcessingOptions,
)
from client_api import api_utils
from scripts.faceswaplab_utils.face_checkpoints_utils import (
build_face_checkpoint_and_save,
)
def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str: # type: ignore
@ -135,3 +139,23 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
result_images = [encode_to_base64(img) for img in faces]
response = api_utils.FaceSwapExtractResponse(images=result_images)
return response
@app.post(
"/faceswaplab/build",
tags=["faceswaplab"],
description="Build a face checkpoint using base64 images, return base64 satetensors",
)
async def build(base64_images: List[str]) -> Optional[str]:
if len(base64_images) > 0:
pil_images = [base64_to_pil(img) for img in base64_images]
with tempfile.NamedTemporaryFile(
delete=True, suffix=".safetensors"
) as temp_file:
build_face_checkpoint_and_save(
images=pil_images,
name="api_ckpt",
overwrite=True,
path=temp_file.name,
)
return api_utils.safetensors_to_base64(temp_file.name)
return None

9
scripts/faceswaplab_swapping/swapper.py
View File

@ -468,12 +468,12 @@ def get_or_default(l: List[Any], index: int, default: Any) -> Any:
return l[index] if index < len(l) else default
def get_faces_from_img_files(files: List[str]) -> List[Optional[CV2ImgU8]]:
def get_faces_from_img_files(images: List[PILImage]) -> List[Optional[CV2ImgU8]]:
"""
Extracts faces from a list of image files.
Args:
files (list): A list of file objects representing image files.
images (list): A list of PILImage objects representing image files.
Returns:
list: A list of detected faces.
@ -482,9 +482,8 @@ def get_faces_from_img_files(files: List[str]) -> List[Optional[CV2ImgU8]]:
faces = []
if len(files) > 0:
for file in files:
img = Image.open(file) # Open the image file
if len(images) > 0:
for img in images:
face = get_or_default(
get_faces(pil_to_cv2(img)), 0, None
) # Extract faces from the image

4
scripts/faceswaplab_ui/faceswaplab_tab.py
View File

@ -153,9 +153,9 @@ def build_face_checkpoint_and_save(
if not batch_files:
logger.error("No face found")
return None
filenames = [x.name for x in batch_files]
images = [Image.open(file.name) for file in batch_files]
preview_image = face_checkpoints_utils.build_face_checkpoint_and_save(
filenames, name, overwrite=overwrite
images, name, overwrite=overwrite
)
except Exception as e:
logger.error("Failed to build checkpoint %s", e)

11
scripts/faceswaplab_utils/face_checkpoints_utils.py
View File

@ -38,7 +38,7 @@ def sanitize_name(name: str) -> str:
def build_face_checkpoint_and_save(
batch_files: List[str], name: str, overwrite: bool = False
images: List[PILImage], name: str, overwrite: bool = False, path: str = None
) -> PILImage:
"""
Builds a face checkpoint using the provided image files, performs face swapping,
@ -55,9 +55,9 @@ def build_face_checkpoint_and_save(
try:
name = sanitize_name(name)
batch_files = batch_files or []
logger.info("Build %s %s", name, [x for x in batch_files])
faces = swapper.get_faces_from_img_files(batch_files)
images = images or []
logger.info("Build %s with %s images", name, len(images))
faces = swapper.get_faces_from_img_files(images)
blended_face = swapper.blend_faces(faces)
preview_path = os.path.join(
scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
@ -95,6 +95,9 @@ def build_face_checkpoint_and_save(
)
preview_image = result.image
if path:
file_path = path
else:
file_path = os.path.join(get_checkpoint_path(), f"{name}.safetensors")
if not overwrite:
file_number = 1

108
tests/test_api.py
View File

@ -2,22 +2,28 @@ from typing import List
import pytest
import requests
import sys
import tempfile
import safetensors
sys.path.append(".")
import requests
from client_api.api_utils import (
FaceSwapUnit,
FaceSwapResponse,
PostProcessingOptions,
FaceSwapRequest,
base64_to_pil,
InswappperOptions,
pil_to_base64,
PostProcessingOptions,
InpaintingWhen,
FaceSwapCompareRequest,
InpaintingOptions,
FaceSwapRequest,
FaceSwapResponse,
FaceSwapExtractRequest,
FaceSwapCompareRequest,
FaceSwapExtractResponse,
compare_faces,
InpaintingOptions,
base64_to_pil,
base64_to_safetensors,
safetensors_to_base64,
)
from PIL import Image
@ -37,6 +43,13 @@ def face_swap_request() -> FaceSwapRequest:
source_img=pil_to_base64("references/woman.png"), # The face you want to use
same_gender=True,
faces_index=(0,), # Replace first woman since same gender is on
swapping_options=InswappperOptions(
face_restorer_name="CodeFormer",
upscaler_name="LDSR",
improved_mask=True,
sharpen=True,
color_corrections=True,
),
)
# Post-processing config
@ -179,3 +192,86 @@ def test_faceswap_inpainting(face_swap_request: FaceSwapRequest) -> None:
data = response.json()
assert "images" in data
assert "infos" in data
def test_faceswap_checkpoint_building() -> None:
source_images: List[str] = [
pil_to_base64("references/man.png"),
pil_to_base64("references/woman.png"),
]
response = requests.post(
url=f"{base_url}/faceswaplab/build",
json=source_images,
headers={"Content-Type": "application/json; charset=utf-8"},
)
assert response.status_code == 200
with tempfile.NamedTemporaryFile(delete=True) as temp_file:
base64_to_safetensors(response.json(), output_path=temp_file.name)
with safetensors.safe_open(temp_file.name, framework="pt") as f:
assert "age" in f.keys()
assert "gender" in f.keys()
assert "embedding" in f.keys()
def test_faceswap_checkpoint_building_and_using() -> None:
source_images: List[str] = [
pil_to_base64("references/man.png"),
]
response = requests.post(
url=f"{base_url}/faceswaplab/build",
json=source_images,
headers={"Content-Type": "application/json; charset=utf-8"},
)
assert response.status_code == 200
with tempfile.NamedTemporaryFile(delete=True) as temp_file:
base64_to_safetensors(response.json(), output_path=temp_file.name)
with safetensors.safe_open(temp_file.name, framework="pt") as f:
assert "age" in f.keys()
assert "gender" in f.keys()
assert "embedding" in f.keys()
# First face unit :
unit1 = FaceSwapUnit(
source_face=safetensors_to_base64(
temp_file.name
), # convert the checkpoint to base64
faces_index=(0,), # Replace first face
swapping_options=InswappperOptions(
face_restorer_name="CodeFormer",
upscaler_name="LDSR",
improved_mask=True,
sharpen=True,
color_corrections=True,
),
)
# Prepare the request
request = FaceSwapRequest(
image=pil_to_base64("tests/test_image.png"), units=[unit1]
)
# Face Swap
response = requests.post(
url=f"{base_url}/faceswaplab/swap_face",
data=request.json(),
headers={"Content-Type": "application/json; charset=utf-8"},
)
assert response.status_code == 200
fsr = FaceSwapResponse.parse_obj(response.json())
data = response.json()
assert "images" in data
assert "infos" in data
# First face is the man
assert (
compare_faces(
fsr.pil_images[0], Image.open("references/man.png"), base_url=base_url
)
> 0.5
)