optimized_swapper.py

import os
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

import cv2
import insightface
from insightface.app import FaceAnalysis

MODEL_PATH = os.getenv("INSWAPPER_MODEL", "inswapper_128.onnx")
GPU = os.getenv("FACE_SWAP_GPU", "1") not in {"0", "false", "False", "cpu"}
CUDA_DEVICE = int(os.getenv("CUDA_DEVICE_ID", "0"))
FFMPEG_BIN = os.getenv("FFMPEG_BIN", "ffmpeg")
FFPROBE_BIN = os.getenv("FFPROBE_BIN", "ffprobe")
GPU_ONLY = os.getenv("FACE_SWAP_GPU_ONLY", "1") in {"1", "true", "True", "yes", "on"}


class FaceSwapEngine:
    """GPU-first InsightFace engine with safer ffmpeg video handling."""

    def __init__(self):
        self._source_face = None
        self._source_key = None
        self._source_idx = None
        self._det_size = None
        self.providers = self._build_providers()
        self.gpu = GPU

        self.app = FaceAnalysis(
            name=os.getenv("INSIGHTFACE_MODEL", "buffalo_l"),
            providers=self.providers,
        )
        self.app.prepare(
            ctx_id=CUDA_DEVICE if self.gpu else -1,
            det_size=(320, 320),
        )

        self.swapper = insightface.model_zoo.get_model(
            MODEL_PATH,
            providers=self.providers,
        )

    def _build_providers(self):
        if GPU:
            cuda_provider = (
                "CUDAExecutionProvider",
                {
                    "device_id": CUDA_DEVICE,
                    "gpu_mem_limit": 80 * 1024 * 1024 * 1024,
                    "arena_extend_strategy": "kNextPowerOfTwo",
                    "cudnn_conv_algo_search": "EXHAUSTIVE",
                    "cudnn_conv_use_max_workspace": "1",
                    "do_copy_in_default_stream": True,
                },
            )
            if GPU_ONLY:
                return [cuda_provider]
            return [cuda_provider, "CPUExecutionProvider"]
        return ["CPUExecutionProvider"]

    @staticmethod
    def _img_key(img):
        return (id(img), img.shape, img.dtype.str)

    @staticmethod
    def _sort_faces(faces):
        return sorted(faces, key=lambda f: float(f.bbox[0]))

    def _detect(self, image, det_size=320, max_num=0):
        if self._det_size != det_size:
            self.app.prepare(
                ctx_id=CUDA_DEVICE if self.gpu else -1,
                det_size=(det_size, det_size),
            )
            self._det_size = det_size
        try:
            faces = self.app.get(image, max_num=max_num)
        except TypeError:
            faces = self.app.get(image)
        return self._sort_faces(faces)

    def prepare_source(self, source, source_idx=1, det_size=320):
        faces = self._detect(source, det_size, max_num=0)
        if len(faces) < source_idx:
            raise ValueError(
                f"Source image contains {len(faces)} faces; requested face {source_idx}."
            )
        self._source_face = faces[source_idx - 1]
        self._source_idx = source_idx
        self._source_key = self._img_key(source)
        return self._source_face

    def swap_prepared_source(self, target, target_idx=1, det_size=320, max_faces=0):
        if self._source_face is None:
            raise RuntimeError("Source face has not been prepared.")
        faces = self._detect(target, det_size, max_num=max_faces)
        if len(faces) < target_idx:
            raise ValueError(
                f"Target image contains {len(faces)} faces; requested face {target_idx}."
            )
        return self.swapper.get(
            target,
            faces[target_idx - 1],
            self._source_face,
            paste_back=True,
        )

    def swap_image(self, source, source_idx, target, target_idx, det_size=320):
        self.prepare_source(source, source_idx, det_size)
        max_faces = target_idx if target_idx > 1 else 1
        return self.swap_prepared_source(target, target_idx, det_size, max_faces)

    def _detect_video_face(self, frame, target_idx, det_size):
        max_faces = target_idx if target_idx > 1 else 1
        faces = self._detect(frame, det_size, max_num=max_faces)
        if len(faces) < target_idx:
            raise ValueError("Target face not found.")
        return faces[target_idx - 1]

    def _probe_video(self, video_path):
        cmd = [
            FFPROBE_BIN,
            "-v", "error",
            "-select_streams", "v:0",
            "-show_entries", "stream=avg_frame_rate,nb_frames,width,height",
            "-of", "default=noprint_wrappers=1:nokey=0",
            video_path,
        ]
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            data = {}
            for line in result.stdout.splitlines():
                if "=" in line:
                    k, v = line.strip().split("=", 1)
                    data[k] = v
            fps = 30.0
            afr = data.get("avg_frame_rate", "30/1")
            if "/" in afr:
                a, b = afr.split("/", 1)
                if float(b) != 0:
                    fps = float(a) / float(b)
            else:
                fps = float(afr)
            return {
                "fps": fps,
                "nb_frames": int(data.get("nb_frames", "0") or 0),
                "width": int(data.get("width", "0") or 0),
                "height": int(data.get("height", "0") or 0),
            }
        except Exception:
            return {"fps": 30.0, "nb_frames": 0, "width": 0, "height": 0}

    def _run_ffmpeg(self, cmd):
        subprocess.run(cmd, check=True)

    def _decode_video_to_frames(self, video_path, frames_dir):
        pattern = str(Path(frames_dir) / "frame_%08d.png")

        gpu_cmd = [
            FFMPEG_BIN, "-y", "-loglevel", "error",
            "-hwaccel", "cuda",
            "-hwaccel_output_format", "cuda",
            "-i", video_path,
            "-vf", "hwdownload,format=nv12,format=bgr24",
            pattern,
        ]

        cpu_cmd = [
            FFMPEG_BIN, "-y", "-loglevel", "error",
            "-i", video_path,
            "-vf", "format=bgr24",
            pattern,
        ]

        if self.gpu:
            try:
                self._run_ffmpeg(gpu_cmd)
                print("[FaceSwap] ffmpeg decode path: GPU")
                return "gpu"
            except subprocess.CalledProcessError as exc:
                print(f"[FaceSwap] GPU decode failed, falling back to CPU decode: {exc}")
                for p in Path(frames_dir).glob("frame_*.png"):
                    p.unlink(missing_ok=True)
                self._run_ffmpeg(cpu_cmd)
                print("[FaceSwap] ffmpeg decode path: CPU fallback")
                return "cpu-fallback"

        self._run_ffmpeg(cpu_cmd)
        print("[FaceSwap] ffmpeg decode path: CPU")
        return "cpu"

    def _ffmpeg_encode(self, frames_dir, output_path, fps, audio_source=None):
        pattern = str(Path(frames_dir) / "frame_%08d.jpg")
        cmd = [
            FFMPEG_BIN, "-y",
            "-loglevel", "error",
            "-framerate", f"{fps:.6f}",
            "-i", pattern,
        ]
        if audio_source:
            cmd += ["-i", audio_source, "-map", "0:v:0", "-map", "1:a?"]
        if self.gpu:
            cmd += [
                "-c:v", "h264_nvenc",
                "-preset", "p4",
                "-cq", "18",
                "-pix_fmt", "yuv420p",
            ]
        else:
            cmd += [
                "-c:v", "libx264",
                "-preset", "veryfast",
                "-crf", "18",
                "-pix_fmt", "yuv420p",
            ]
        if audio_source:
            cmd += ["-c:a", "aac", "-b:a", "160k", "-shortest"]
        cmd.append(output_path)
        self._run_ffmpeg(cmd)

    def swap_video(
        self,
        source,
        source_idx,
        video_path,
        target_idx,
        det_size=320,
        detection_interval=1,
        jpeg_quality=92,
        preserve_audio=True,
    ):
        self.prepare_source(source, source_idx, det_size)

        work = tempfile.mkdtemp(prefix="faceswap_")
        decoded_dir = os.path.join(work, "decoded")
        swapped_dir = os.path.join(work, "swapped")
        os.makedirs(decoded_dir, exist_ok=True)
        os.makedirs(swapped_dir, exist_ok=True)
        output = os.path.join(work, "swapped.mp4")

        meta = self._probe_video(video_path)
        fps = meta["fps"] or 30.0

        t0 = time.perf_counter()
        failed = 0
        idx = 0
        last_face = None

        try:
            decode_mode = self._decode_video_to_frames(video_path, decoded_dir)
            frame_paths = sorted(Path(decoded_dir).glob("frame_*.png"))
            if not frame_paths:
                raise ValueError("No frames decoded from video.")

            for frame_path in frame_paths:
                frame = cv2.imread(str(frame_path))
                if frame is None:
                    failed += 1
                    idx += 1
                    continue

                try:
                    if last_face is None or idx % max(1, int(detection_interval)) == 0:
                        last_face = self._detect_video_face(frame, target_idx, det_size)
                    swapped = self.swapper.get(
                        frame,
                        last_face,
                        self._source_face,
                        paste_back=True,
                    )
                except Exception:
                    try:
                        last_face = self._detect_video_face(frame, target_idx, det_size)
                        swapped = self.swapper.get(
                            frame,
                            last_face,
                            self._source_face,
                            paste_back=True,
                        )
                    except Exception:
                        swapped = frame
                        failed += 1

                out_path = os.path.join(swapped_dir, f"frame_{idx:08d}.jpg")
                cv2.imwrite(
                    out_path,
                    swapped,
                    [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_quality)],
                )
                idx += 1

            self._ffmpeg_encode(
                swapped_dir,
                output,
                fps,
                audio_source=video_path if preserve_audio else None,
            )

            stable = tempfile.NamedTemporaryFile(
                prefix="faceswap_result_", suffix=".mp4", delete=False
            ).name
            with open(output, "rb") as srcf, open(stable, "wb") as dstf:
                while True:
                    chunk = srcf.read(8 * 1024 * 1024)
                    if not chunk:
                        break
                    dstf.write(chunk)

            elapsed = time.perf_counter() - t0
            print(
                f"[FaceSwap] processed {idx} frames in {elapsed:.2f}s "
                f"({idx / max(elapsed, 1e-6):.2f} FPS), failures={failed}, "
                f"gpu={self.gpu}, gpu_only={GPU_ONLY}, decode_mode={decode_mode}"
            )
            return stable
        finally:
            shutil.rmtree(work, ignore_errors=True)