import copy import math import os import tempfile from dataclasses import dataclass from typing import List, Union, Dict, Set, Tuple import cv2 import numpy as np from PIL import Image import insightface import onnxruntime from scripts.cimage import convert_to_sd from modules.face_restoration import FaceRestoration, restore_faces from modules.upscaler import Upscaler, UpscalerData from scripts.roop_logging import logger providers = ["CPUExecutionProvider"] models_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models") insightface_path = os.path.join(models_path, "insightface") insightface_models_path = os.path.join(insightface_path, "models") swapper_path = os.path.join(models_path, "roop") @dataclass class UpscaleOptions: scale: int = 1 upscaler: UpscalerData = None upscale_visibility: float = 0.5 face_restorer: FaceRestoration = None restorer_visibility: float = 0.5 FS_MODEL = None CURRENT_FS_MODEL_PATH = None def getFaceSwapModel(model_path: str): global FS_MODEL global CURRENT_FS_MODEL_PATH if CURRENT_FS_MODEL_PATH is None or CURRENT_FS_MODEL_PATH != model_path: CURRENT_FS_MODEL_PATH = model_path FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers) return FS_MODEL def upscale_image(image: Image, upscale_options: UpscaleOptions): result_image = image if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None": original_image = result_image.copy() logger.info( "Upscale with %s scale = %s", upscale_options.upscaler.name, upscale_options.scale, ) result_image = upscale_options.upscaler.scaler.upscale( image, upscale_options.scale, upscale_options.upscaler.data_path ) if upscale_options.scale == 1: result_image = Image.blend( original_image, result_image, upscale_options.upscale_visibility ) if upscale_options.face_restorer is not None: original_image = result_image.copy() logger.info("Restore face with %s", upscale_options.face_restorer.name()) numpy_image = np.array(result_image) numpy_image = upscale_options.face_restorer.restore(numpy_image) restored_image = Image.fromarray(numpy_image) result_image = Image.blend( original_image, restored_image, upscale_options.restorer_visibility ) return result_image def order_largest_smallest(faces, reverse=False): def area(face): x1, y1, x2, y2 = face.bbox return (x2 - x1) * (y2 - y1) return sorted(faces, key=lambda x: area(x), reverse=not reverse) def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640), sorter="left to right", reverse_order=False): face_analyser = insightface.app.FaceAnalysis(name="buffalo_l", providers=providers, root=insightface_path) buffalo_path = os.path.join(insightface_models_path, "buffalo_l.zip") if os.path.exists(buffalo_path): # Remove the zip to save space os.remove(buffalo_path) face_analyser.prepare(ctx_id=0, det_size=det_size) face = face_analyser.get(img_data) if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320: det_size_half = (det_size[0] // 2, det_size[1] // 2) return get_face_single(img_data, face_index=face_index, det_size=det_size_half) try: if sorter == "left to right": result = sorted(face, key=lambda x: x.bbox[0], reverse=reverse_order) elif sorter == "up to down": result = sorted(face, key=lambda x: x.bbox[1], reverse=reverse_order) elif sorter == "largest to smallest": result = order_largest_smallest(face, reverse=reverse_order) return result[face_index] except IndexError: return None @dataclass class ImageResult: path: Union[str, None] = None similarity: Union[Dict[int, float], None] = None # face, 0..1 def image(self) -> Union[Image.Image, None]: if self.path: return Image.open(self.path) return None def swap_face( source_img: Image.Image, target_img: Image.Image, model: Union[str, None] = None, faces_index: Set[int] = {0}, reference_faces_index: Set[int] = {0}, upscale_options: Union[UpscaleOptions, None] = None, face_order: str = "left to right", reverse_order: bool = False, reference_order: str = "left to right", reverse_reference_order: bool = False, ) -> ImageResult: result_image = target_img converted = convert_to_sd(target_img) scale, fn = converted[0], converted[1] if model is not None and not scale: if isinstance(source_img, str): # source_img is a base64 string import base64, io if 'base64,' in source_img: # check if the base64 string has a data URL scheme base64_data = source_img.split('base64,')[-1] img_bytes = base64.b64decode(base64_data) else: # if no data URL scheme, just decode img_bytes = base64.b64decode(source_img) source_img = Image.open(io.BytesIO(img_bytes)) source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR) target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) # # Get source faces # source_faces = [] for face_num in reference_faces_index: source_face = get_face_single(source_img, face_index=face_num, sorter=reference_order, reverse_order=reverse_reference_order) if source_face is not None: source_faces.append(source_face) else: logger.info(f"No source face found for {face_num}") logger.info(f"Found {len(source_faces)} source faces") source_face_idx = 0 if len(source_faces) > 0: result = target_img model_path = os.path.join(swapper_path, model) face_swapper = getFaceSwapModel(model_path) for face_num in faces_index: target_face = get_face_single(target_img, face_index=face_num, sorter=face_order, reverse_order=reverse_order) if target_face is not None: source_face = source_faces[source_face_idx] logger.info(f"Swapping source face {source_face_idx} onto target face {face_num}") result = face_swapper.get(result, target_face, source_face) else: logger.info(f"No target face found for {face_num}") source_face_idx = (source_face_idx + 1) % len(source_faces) result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) if upscale_options is not None: result_image = upscale_image(result_image, upscale_options) else: logger.info("No source face found") result_image.save(fn.name) return ImageResult(path=fn.name)