278 lines
8.1 KiB
Python
278 lines
8.1 KiB
Python
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,
|
|
InswappperOptions,
|
|
pil_to_base64,
|
|
PostProcessingOptions,
|
|
InpaintingWhen,
|
|
InpaintingOptions,
|
|
FaceSwapRequest,
|
|
FaceSwapResponse,
|
|
FaceSwapExtractRequest,
|
|
FaceSwapCompareRequest,
|
|
FaceSwapExtractResponse,
|
|
compare_faces,
|
|
base64_to_pil,
|
|
base64_to_safetensors,
|
|
safetensors_to_base64,
|
|
)
|
|
from PIL import Image
|
|
|
|
base_url = "http://127.0.0.1:7860"
|
|
|
|
|
|
@pytest.fixture
|
|
def face_swap_request() -> FaceSwapRequest:
|
|
# First face unit
|
|
unit1 = FaceSwapUnit(
|
|
source_img=pil_to_base64("references/man.png"), # The face you want to use
|
|
faces_index=(0,), # Replace first face
|
|
)
|
|
|
|
# Second face unit
|
|
unit2 = FaceSwapUnit(
|
|
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
|
|
pp = PostProcessingOptions(
|
|
face_restorer_name="CodeFormer",
|
|
codeformer_weight=0.5,
|
|
restorer_visibility=1,
|
|
upscaler_name="Lanczos",
|
|
scale=4,
|
|
inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE,
|
|
inpainting_options=InpaintingOptions(
|
|
inpainting_steps=30,
|
|
inpainting_denoising_strengh=0.1,
|
|
),
|
|
)
|
|
# Prepare the request
|
|
request = FaceSwapRequest(
|
|
image=pil_to_base64("tests/test_image.png"),
|
|
units=[unit1, unit2],
|
|
postprocessing=pp,
|
|
)
|
|
|
|
return request
|
|
|
|
|
|
def test_version() -> None:
|
|
response = requests.get(f"{base_url}/faceswaplab/version")
|
|
assert response.status_code == 200
|
|
assert "version" in response.json()
|
|
|
|
|
|
def test_compare() -> None:
|
|
request = FaceSwapCompareRequest(
|
|
image1=pil_to_base64("references/man.png"),
|
|
image2=pil_to_base64("references/man.png"),
|
|
)
|
|
|
|
response = requests.post(
|
|
url=f"{base_url}/faceswaplab/compare",
|
|
data=request.json(),
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
similarity = float(response.text)
|
|
assert similarity > 0.90
|
|
|
|
|
|
def test_extract() -> None:
|
|
pp = PostProcessingOptions(
|
|
face_restorer_name="CodeFormer",
|
|
codeformer_weight=0.5,
|
|
restorer_visibility=1,
|
|
upscaler_name="Lanczos",
|
|
)
|
|
|
|
request = FaceSwapExtractRequest(
|
|
images=[pil_to_base64("tests/test_image.png")], postprocessing=pp
|
|
)
|
|
|
|
response = requests.post(
|
|
url=f"{base_url}/faceswaplab/extract",
|
|
data=request.json(),
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
res = FaceSwapExtractResponse.parse_obj(response.json())
|
|
|
|
assert len(res.pil_images) == 2
|
|
|
|
# First face is the man
|
|
assert (
|
|
compare_faces(
|
|
res.pil_images[0], Image.open("tests/test_image.png"), base_url=base_url
|
|
)
|
|
> 0.5
|
|
)
|
|
|
|
|
|
def test_faceswap(face_swap_request: FaceSwapRequest) -> None:
|
|
response = requests.post(
|
|
f"{base_url}/faceswaplab/swap_face",
|
|
data=face_swap_request.json(),
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "images" in data
|
|
assert "infos" in data
|
|
|
|
res = FaceSwapResponse.parse_obj(response.json())
|
|
images: List[Image.Image] = res.pil_images
|
|
assert len(images) == 1
|
|
image = images[0]
|
|
orig_image = base64_to_pil(face_swap_request.image)
|
|
assert image.width == orig_image.width * face_swap_request.postprocessing.scale
|
|
assert image.height == orig_image.height * face_swap_request.postprocessing.scale
|
|
|
|
# Compare the result and ensure similarity for the man (first face)
|
|
|
|
request = FaceSwapCompareRequest(
|
|
image1=pil_to_base64("references/man.png"),
|
|
image2=res.images[0],
|
|
)
|
|
|
|
response = requests.post(
|
|
url=f"{base_url}/faceswaplab/compare",
|
|
data=request.json(),
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
similarity = float(response.text)
|
|
assert similarity > 0.50
|
|
|
|
|
|
def test_faceswap_inpainting(face_swap_request: FaceSwapRequest) -> None:
|
|
face_swap_request.units[0].pre_inpainting = InpaintingOptions(
|
|
inpainting_denoising_strengh=0.4,
|
|
inpainting_prompt="Photo of a funny man",
|
|
inpainting_negative_prompt="blurry, bad art",
|
|
inpainting_steps=100,
|
|
)
|
|
|
|
face_swap_request.units[0].post_inpainting = InpaintingOptions(
|
|
inpainting_denoising_strengh=0.4,
|
|
inpainting_prompt="Photo of a funny man",
|
|
inpainting_negative_prompt="blurry, bad art",
|
|
inpainting_steps=20,
|
|
inpainting_sampler="Euler a",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{base_url}/faceswaplab/swap_face",
|
|
data=face_swap_request.json(),
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
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
|
|
)
|