import gradio as gr
import os
import torch
import cv2
import numpy as np
from roop.utilities.device import get_device
from roop.processors.frame_processor import process_frame

# 选择 CPU 或 GPU 设备
device = get_device()
print(f"Using device: {device}")

def swap_face(source_img, target_img):
    try:
        # 读取输入图像
        source = cv2.imread(source_img)
        target = cv2.imread(target_img)

        # 进行人脸替换
        result = process_frame(source, target, device=device)

        # 保存结果
        output_path = "output.jpg"
        cv2.imwrite(output_path, result)

        return output_path
    except Exception as e:
        return str(e)

with gr.Blocks() as demo:
    gr.Markdown("# Roop Face Swap (Hugging Face)")
    with gr.Row():
        source_img = gr.Image(type="filepath", label="Source Face")
        target_img = gr.Image(type="filepath", label="Target Image")
    btn = gr.Button("Swap Face")
    output_img = gr.Image(label="Output Image")

    btn.click(swap_face, inputs=[source_img, target_img], outputs=output_img)

demo.launch()