diff --git a/scripts/faceswaplab_api/faceswaplab_api.py b/scripts/faceswaplab_api/faceswaplab_api.py index cb000f8..3a0ee9f 100644 --- a/scripts/faceswaplab_api/faceswaplab_api.py +++ b/scripts/faceswaplab_api/faceswaplab_api.py @@ -21,6 +21,7 @@ from scripts.faceswaplab_swapping.face_checkpoints import ( build_face_checkpoint_and_save, ) +from scripts.faceswaplab_utils.typing import PILImage def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str: # type: ignore @@ -99,7 +100,7 @@ async def swap_face( pp_options = None units = get_faceswap_units_settings(request.units) - swapped_images = swapper.batch_process( + swapped_images: Optional[List[PILImage]] = swapper.batch_process( [src_image], None, units=units, postprocess_options=pp_options ) diff --git a/scripts/faceswaplab_swapping/face_checkpoints.py b/scripts/faceswaplab_swapping/face_checkpoints.py index 89ea6ff..ad952d3 100644 --- a/scripts/faceswaplab_swapping/face_checkpoints.py +++ b/scripts/faceswaplab_swapping/face_checkpoints.py @@ -38,8 +38,11 @@ def sanitize_name(name: str) -> str: def build_face_checkpoint_and_save( - images: List[PILImage], name: str, overwrite: bool = False, path: str = None -) -> PILImage: + images: List[PILImage], + name: str, + overwrite: bool = False, + path: Optional[str] = None, +) -> Optional[PILImage]: """ Builds a face checkpoint using the provided image files, performs face swapping, and saves the result to a file. If a blended face is successfully obtained and the face swapping @@ -57,8 +60,12 @@ def build_face_checkpoint_and_save( name = sanitize_name(name) 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) + faces: List[Face] = swapper.get_faces_from_img_files(images=images) + if faces is None or len(faces) == 0: + logger.error("No source faces found") + return None + + blended_face: Optional[Face] = swapper.blend_faces(faces) preview_path = os.path.join( scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references" ) @@ -85,40 +92,51 @@ def build_face_checkpoint_and_save( "Failed to open reference image, cannot create preview : That should not happen unless you deleted the references folder or change the detection threshold." ) else: - result = swapper.swap_face( + result: swapper.ImageResult = swapper.swap_face( target_faces=[target_face], source_face=blended_face, target_img=reference_preview_img, model=get_swap_models()[0], - swapping_options=InswappperOptions(face_restorer_name="Codeformer"), + swapping_options=InswappperOptions( + face_restorer_name="CodeFormer", + restorer_visibility=1, + upscaler_name="Lanczos", + codeformer_weight=1, + improved_mask=True, + color_corrections=False, + sharpen=True, + ), ) 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 - while os.path.exists(file_path): - file_path = os.path.join( - get_checkpoint_path(), f"{name}_{file_number}.safetensors" - ) - file_number += 1 - save_face(filename=file_path, face=blended_face) - preview_image.save(file_path + ".png") - try: - data = load_face(file_path) - logger.debug(data) - except Exception as e: - logger.error("Error loading checkpoint, after creation %s", e) - traceback.print_exc() - - return preview_image + if path: + file_path = path + else: + file_path = os.path.join( + get_checkpoint_path(), f"{name}.safetensors" + ) + if not overwrite: + file_number = 1 + while os.path.exists(file_path): + file_path = os.path.join( + get_checkpoint_path(), + f"{name}_{file_number}.safetensors", + ) + file_number += 1 + save_face(filename=file_path, face=blended_face) + preview_image.save(file_path + ".png") + try: + data = load_face(file_path) + logger.debug(data) + except Exception as e: + logger.error("Error loading checkpoint, after creation %s", e) + traceback.print_exc() + + return preview_image else: logger.error("No face found") - return None + return None # type: ignore except Exception as e: logger.error("Failed to build checkpoint %s", e) traceback.print_exc() @@ -139,7 +157,7 @@ def save_face(face: Face, filename: str) -> None: raise e -def load_face(name: str) -> Face: +def load_face(name: str) -> Optional[Face]: if name.startswith("data:application/face;base64,"): with tempfile.NamedTemporaryFile(delete=True) as temp_file: api_utils.base64_to_safetensors(name, temp_file.name) diff --git a/scripts/faceswaplab_swapping/swapper.py b/scripts/faceswaplab_swapping/swapper.py index ceb8e56..0c58396 100644 --- a/scripts/faceswaplab_swapping/swapper.py +++ b/scripts/faceswaplab_swapping/swapper.py @@ -141,7 +141,7 @@ def batch_process( src_images: List[Union[PILImage, str]], # image or filename save_path: Optional[str], units: List[FaceSwapUnitSettings], - postprocess_options: PostProcessingOptions, + postprocess_options: Optional[PostProcessingOptions], ) -> Optional[List[PILImage]]: """ Process a batch of images, apply face swapping according to the given settings, and optionally save the resulting images to a specified path. @@ -527,7 +527,7 @@ 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(images: List[PILImage]) -> List[Optional[CV2ImgU8]]: +def get_faces_from_img_files(images: List[PILImage]) -> List[Face]: """ Extracts faces from a list of image files. @@ -539,7 +539,7 @@ def get_faces_from_img_files(images: List[PILImage]) -> List[Optional[CV2ImgU8]] """ - faces = [] + faces: List[Face] = [] if len(images) > 0: for img in images: @@ -598,7 +598,6 @@ def swap_face( target_faces: List[Face], model: str, swapping_options: Optional[InswappperOptions], - compute_similarity: bool = True, ) -> ImageResult: """ Swaps faces in the target image with the source face. @@ -680,9 +679,9 @@ def process_image_unit( model: str, unit: FaceSwapUnitSettings, image: PILImage, - info: str = None, + info: Optional[str] = None, force_blend: bool = False, -) -> List[Tuple[PILImage, str]]: +) -> List[Tuple[PILImage, Optional[str]]]: """Process one image and return a List of (image, info) (one if blended, many if not). Args: @@ -723,7 +722,9 @@ def process_image_unit( sort_by_face_size=unit.sort_by_size, ) - target_faces = filter_faces(faces, filtering_options=face_filtering_options) + target_faces: List[Face] = filter_faces( + all_faces=faces, filtering_options=face_filtering_options + ) # Apply pre-inpainting to image if unit.pre_inpainting.inpainting_denoising_strengh > 0: @@ -738,7 +739,6 @@ def process_image_unit( target_faces=target_faces, model=model, swapping_options=unit.swapping_options, - compute_similarity=unit.compute_similarity, ) # Apply post-inpainting to image if unit.post_inpainting.inpainting_denoising_strengh > 0: diff --git a/scripts/faceswaplab_ui/faceswaplab_tab.py b/scripts/faceswaplab_ui/faceswaplab_tab.py index 1e1e593..ec0e268 100644 --- a/scripts/faceswaplab_ui/faceswaplab_tab.py +++ b/scripts/faceswaplab_ui/faceswaplab_tab.py @@ -156,16 +156,18 @@ def build_face_checkpoint_and_save( if not batch_files: logger.error("No face found") return None # type: ignore (Optional not really supported by old gradio) - images = [Image.open(file.name) for file in batch_files] # type: ignore - preview_image = face_checkpoints.build_face_checkpoint_and_save( - images, name, overwrite=overwrite + images: list[PILImage] = [Image.open(file.name) for file in batch_files] # type: ignore + preview_image: PILImage | None = ( + face_checkpoints.build_face_checkpoint_and_save( + images=images, name=name, overwrite=overwrite + ) ) except Exception as e: logger.error("Failed to build checkpoint %s", e) traceback.print_exc() return None # type: ignore - return preview_image + return preview_image # type: ignore def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame: