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)