add tests

client_api/client_utils.py

@@ -0,0 +1,179 @@
+from typing import List, Tuple
+from PIL import Image
+from pydantic import BaseModel, Field
+from enum import Enum
+import base64, io
+from io import BytesIO
+from typing import List, Tuple, Optional
+import numpy as np
+
+
+class InpaintingWhen(Enum):
+    NEVER = "Never"
+    BEFORE_UPSCALING = "Before Upscaling/all"
+    BEFORE_RESTORE_FACE = "After Upscaling/Before Restore Face"
+    AFTER_ALL = "After All"
+
+
+class FaceSwapUnit(BaseModel):
+    # The image given in reference
+    source_img: str = Field(
+        description="base64 reference image",
+        examples=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD...."],
+        default=None,
+    )
+    # The checkpoint file
+    source_face: str = Field(
+        description="face checkpoint (from models/faceswaplab/faces)",
+        examples=["my_face.pkl"],
+        default=None,
+    )
+    # base64 batch source images
+    batch_images: Tuple[str] = Field(
+        description="list of base64 batch source images",
+        examples=[
+            "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD....",
+            "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD....",
+        ],
+        default=None,
+    )
+
+    # Will blend faces if True
+    blend_faces: bool = Field(description="Will blend faces if True", default=True)
+
+    # Use same gender filtering
+    same_gender: bool = Field(description="Use same gender filtering", default=False)
+
+    # Use same gender filtering
+    sort_by_size: bool = Field(description="Sort Faces by size", default=False)
+
+    # If True, discard images with low similarity
+    check_similarity: bool = Field(
+        description="If True, discard images with low similarity", default=False
+    )
+    # if True will compute similarity and add it to the image info
+    compute_similarity: bool = Field(
+        description="If True will compute similarity and add it to the image info",
+        default=False,
+    )
+
+    # Minimum similarity against the used face (reference, batch or checkpoint)
+    min_sim: float = Field(
+        description="Minimum similarity against the used face (reference, batch or checkpoint)",
+        default=0.0,
+    )
+    # Minimum similarity against the reference (reference or checkpoint if checkpoint is given)
+    min_ref_sim: float = Field(
+        description="Minimum similarity against the reference (reference or checkpoint if checkpoint is given)",
+        default=0.0,
+    )
+
+    # The face index to use for swapping
+    faces_index: Tuple[int] = Field(
+        description="The face index to use for swapping, list of face numbers starting from 0",
+        default=(0,),
+    )
+
+    reference_face_index: int = Field(
+        description="The face index to use to extract face from reference",
+        default=0,
+    )
+
+    def get_batch_images(self) -> List[Image.Image]:
+        images = []
+        if self.batch_images:
+            for img in self.batch_images:
+                images.append(base64_to_pil(img))
+        return images
+
+
+class PostProcessingOptions(BaseModel):
+    face_restorer_name: str = Field(description="face restorer name", default=None)
+    restorer_visibility: float = Field(
+        description="face restorer visibility", default=1, le=1, ge=0
+    )
+    codeformer_weight: float = Field(
+        description="face restorer codeformer weight", default=1, le=1, ge=0
+    )
+
+    upscaler_name: str = Field(description="upscaler name", default=None)
+    scale: float = Field(description="upscaling scale", default=1, le=10, ge=0)
+    upscaler_visibility: float = Field(
+        description="upscaler visibility", default=1, le=1, ge=0
+    )
+
+    inpainting_denoising_strengh: float = Field(
+        description="Inpainting denoising strenght", default=0, lt=1, ge=0
+    )
+    inpainting_prompt: str = Field(
+        description="Inpainting denoising strenght",
+        examples=["Portrait of a [gender]"],
+        default="Portrait of a [gender]",
+    )
+    inpainting_negative_prompt: str = Field(
+        description="Inpainting denoising strenght",
+        examples=[
+            "Deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation"
+        ],
+        default="",
+    )
+    inpainting_steps: int = Field(
+        description="Inpainting steps",
+        examples=["Portrait of a [gender]"],
+        ge=1,
+        le=150,
+        default=20,
+    )
+    inpainting_sampler: str = Field(
+        description="Inpainting sampler", examples=["Euler"], default="Euler"
+    )
+    inpainting_when: InpaintingWhen = Field(
+        description="When inpainting happens",
+        examples=[e.value for e in InpaintingWhen.__members__.values()],
+        default=InpaintingWhen.NEVER,
+    )
+    inpainting_model: str = Field(
+        description="Inpainting model", examples=["Current"], default="Current"
+    )
+
+
+class FaceSwapRequest(BaseModel):
+    image: str = Field(
+        description="base64 reference image",
+        examples=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD...."],
+        default=None,
+    )
+    units: List[FaceSwapUnit]
+    postprocessing: Optional[PostProcessingOptions]
+
+
+class FaceSwapResponse(BaseModel):
+    images: List[str] = Field(description="base64 swapped image", default=None)
+    infos: List[str]
+
+    @property
+    def pil_images(self) -> Image.Image:
+        return [base64_to_pil(img) for img in self.images]
+
+
+def pil_to_base64(img: Image.Image) -> np.array:  # type:ignore
+    if isinstance(img, str):
+        img = Image.open(img)
+
+    buffer = BytesIO()
+    img.save(buffer, format="PNG")
+    img_data = buffer.getvalue()
+    base64_data = base64.b64encode(img_data)
+    return base64_data.decode("utf-8")
+
+
+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]
+        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))