File size: 2,116 Bytes
752b61f 38bae98 752b61f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import os, cv2, tempfile
import gradio as gr
from moviepy.editor import ImageSequenceClip
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model
os.environ["XDG_RUNTIME_DIR"] = "/tmp/runtime-dir"
# Load InsightFace models once at startup
app = FaceAnalysis(name="buffalo_l", providers=["CUDAExecutionProvider","CPUExecutionProvider"])
app.prepare(ctx_id=0, det_size=(1024,1024))
onnx_model_path = "inswapper_128.onnx" # <-- put model in Space files or download
swap_model = get_model(onnx_model_path,
providers=["CUDAExecutionProvider","CPUExecutionProvider"])
def face_swap_api(source_image, target_video):
src_img = cv2.imread(source_image)
src_faces = app.get(src_img)
if len(src_faces) == 0:
return None
src_face = src_faces[0]
frames_dir = tempfile.mkdtemp()
cap = cv2.VideoCapture(target_video)
fps = cap.get(cv2.CAP_PROP_FPS)
frames = []
idx = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret: break
frame_path = os.path.join(frames_dir, f"frame_{idx:04d}.png")
cv2.imwrite(frame_path, frame)
frames.append(frame_path)
idx += 1
cap.release()
swapped_frames = []
for fpath in frames:
frame = cv2.imread(fpath)
faces = app.get(frame)
for face in faces:
frame = swap_model.get(frame, face, src_face, paste_back=True)
swapped_frames.append(frame)
out_path = os.path.join(tempfile.gettempdir(), "swapped_output.mp4")
clip = ImageSequenceClip(
[cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in swapped_frames], fps=fps
)
clip.write_videofile(out_path, codec="libx264", audio=False, verbose=False)
return out_path
demo = gr.Interface(
fn=face_swap_api,
inputs=[gr.Image(type="filepath", label="Source Face"),
gr.Video(label="Target Video")],
outputs=gr.Video(label="Swapped Output Video"),
title="Face Swap API",
description="Upload a face photo and a video. Returns a face-swapped MP4."
)
if __name__ == "__main__":
demo.launch()
|