predict.py

import os
import sys
import time
import uuid
import ssl
import traceback
import tempfile
import subprocess
from typing import Any

import gradio as gr
import cv2
import requests
import torch
from PIL import Image

import insightface
import onnxruntime
from insightface.app import FaceAnalysis

# Try importing GFPGAN safely so dependency issues are clearer
try:
    import gfpgan
except Exception as e:
    gfpgan = None
    gfpgan_import_error = e
else:
    gfpgan_import_error = None

from loggers import logger, request_id as _request_id

ssl._create_default_https_context = ssl._create_unverified_context

if sys.platform == 'darwin':
    cache_file_dir = '/tmp/file'
else:
    cache_file_dir = '/src/file'

os.makedirs(cache_file_dir, exist_ok=True)


def img_url_to_local_path(img_url, file_path=None):
    filename = img_url.split('/')[-1]
    max_count = 3
    count = 0

    if file_path is None:
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(filename)[-1] or ".jpg")
        temp_file_name = temp_file.name
        temp_file.close()
    else:
        temp_file_name = file_path

    while True:
        count += 1
        try:
            res = requests.get(img_url, timeout=60)
            res.raise_for_status()
            with open(temp_file_name, "wb") as f:
                f.write(res.content)
            return temp_file_name
        except Exception as e:
            logger.error(e)

        if count >= max_count:
            msg = f'request {max_count} times for url: {img_url} failed, please check'
            logger.error(msg)
            raise Exception(msg)


def delete_files_day_ago(cache_days=10):
    command = f"find {cache_file_dir} -type f -ctime +{cache_days} -exec rm -f {{}} \\;"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    if result.stdout:
        logger.info(result.stdout)
    if result.stderr:
        logger.warning(result.stderr)


def image_format_by_path(image_path):
    image = Image.open(image_path)
    image_format = image.format
    if not image_format:
        image_format = 'jpg'
    elif image_format == "JPEG":
        image_format = 'jpg'
    else:
        image_format = image_format.lower()
    return image_format


def local_file_for_url(url, cache_days=10):
    filename = url.split('/')[-1]
    file_path = os.path.join(cache_file_dir, filename)

    if not os.path.exists(file_path):
        img_url_to_local_path(url, file_path)
        logger.info(f'download file to {file_path}')
        delete_files_day_ago(cache_days)
    else:
        logger.info(f'cache file {file_path}')

    return file_path


class Predictor:
    def __init__(self):
        self.det_thresh = 0.1
        self.face_swapper = None
        self.face_enhancer = None
        self.face_analyser = None

    def setup(self):
        providers = onnxruntime.get_available_providers()

        self.face_swapper = insightface.model_zoo.get_model(
            'cache/inswapper_128.onnx',
            providers=providers
        )

        self.face_analyser = FaceAnalysis(name='buffalo_l')
        self.face_analyser.prepare(ctx_id=0, det_thresh=self.det_thresh)

        if gfpgan is None:
            raise ImportError(
                f"gfpgan import failed: {gfpgan_import_error}. "
                "This is usually caused by an incompatible torchvision/basicsr version."
            )

        self.face_enhancer = gfpgan.GFPGANer(
            model_path='cache/GFPGANv1.4.pth',
            upscale=1
        )

    def get_face(self, img_data, image_type='target'):
        try:
            if self.face_analyser is None:
                raise RuntimeError("Face analyser is not initialized. Call setup() first.")

            if image_type == 'source':
                self.face_analyser.prepare(ctx_id=0, det_thresh=self.det_thresh)

            analysed = self.face_analyser.get(img_data)
            logger.info(f'face num: {len(analysed)}')

            if len(analysed) == 0:
                msg = 'no face'
                logger.error(msg)
                raise Exception(msg)

            largest = max(
                analysed,
                key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])
            )
            return largest
        except Exception as e:
            logger.error(str(e))
            raise

    def enhance_face(self, target_face, target_frame, weight=0.5):
        if self.face_enhancer is None:
            raise RuntimeError("Face enhancer is not initialized. Call setup() first.")

        start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
        padding_x = int((end_x - start_x) * 0.5)
        padding_y = int((end_y - start_y) * 0.5)

        start_x = max(0, start_x - padding_x)
        start_y = max(0, start_y - padding_y)
        end_x = min(target_frame.shape[1], end_x + padding_x)
        end_y = min(target_frame.shape[0], end_y + padding_y)

        temp_face = target_frame[start_y:end_y, start_x:end_x]
        if temp_face.size:
            _, _, temp_face = self.face_enhancer.enhance(
                temp_face,
                paste_back=True,
                weight=weight
            )
            target_frame[start_y:end_y, start_x:end_x] = temp_face

        return target_frame

    def predict(
        self,
        source_image_path,
        target_image_path,
        enhance_face,
    ) -> Any:
        request_id = None
        det_thresh = 0.1
        weight = 0.5

        if torch.cuda.is_available():
            device = 'cuda'
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = 'mps'
        else:
            device = 'cpu'

        logger.info(f'device: {device}, det_thresh:{det_thresh}')

        try:
            self.det_thresh = det_thresh
            start_time = time.time()

            if not request_id:
                request_id = str(uuid.uuid4())
            _request_id.set(request_id)

            frame = cv2.imread(str(target_image_path))
            source_frame = cv2.imread(str(source_image_path))

            if frame is None:
                raise ValueError(f"Failed to read target image: {target_image_path}")
            if source_frame is None:
                raise ValueError(f"Failed to read source image: {source_image_path}")

            source_face = self.get_face(source_frame, image_type='source')
            target_face = self.get_face(frame)

            ext = image_format_by_path(target_image_path)
            size = os.path.getsize(target_image_path)
            logger.info(f'origin {size / 1024:.2f}k')

            result = self.face_swapper.get(frame, target_face, source_face, paste_back=True)

            if enhance_face:
                result = self.enhance_face(target_face, result, weight)

            out_dir = tempfile.mkdtemp()
            out_path = os.path.join(out_dir, f"{uuid.uuid4()}.{ext}")
            cv2.imwrite(str(out_path), result)

            out_size = os.path.getsize(out_path)
            logger.info(f'result {out_size / 1024:.2f}k')

            cost_time = time.time() - start_time
            logger.info(f'total time: {cost_time * 1000:.2f} ms')

            return Image.open(out_path)

        except Exception as e:
            logger.error(traceback.format_exc())
            logger.error(str(e))
            raise


def swap_faces(source_image_path, target_image_path, enhance_face):
    predictor = Predictor()
    predictor.setup()
    return predictor.predict(
        source_image_path,
        target_image_path,
        enhance_face
    )


if __name__ == "__main__":
    demo = gr.Interface(
        fn=swap_faces,
        inputs=[
            gr.Image(type="filepath"),
            gr.Image(type="filepath"),
            gr.Checkbox(label="Enhance Face", value=True),
        ],
        outputs=[
            # Removed unsupported `show_download_button` for the installed Gradio version
            gr.Image(type="pil")
        ],
        title="Mar's Face Swap",
        #allow_flagging="never"
    )
    demo.launch(share=True)