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()