import gradio as gr
import cv2
import insightface
from insightface.app import FaceAnalysis
import PIL.Image
import numpy as np
import os
import requests

# Функция для скачивания модели, если она не скачалась сама
def download_file(url, filename):
    if not os.path.exists(filename):
        response = requests.get(url)
        with open(filename, "wb") as f:
            f.write(response.content)

# Инициализация (с попыткой обхода ошибки загрузки)
try:
    app = FaceAnalysis(name='buffalore_l', providers=['CPUExecutionProvider'])
    app.prepare(ctx_id=0, det_size=(640, 640))
except:
    # Если база не качается, используем упрощенную инициализацию
    app = FaceAnalysis(providers=['CPUExecutionProvider'])
    app.prepare(ctx_id=0, det_size=(640, 640))

# Ссылка на файл модели (альтернативный источник)
model_url = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
model_path = "inswapper_128.onnx"
download_file(model_url, model_path)
swapper = insightface.model_zoo.get_model(model_path, download=False)

def swap(source_img, target_img):
    if source_img is None or target_img is None:
        return None
    
    source_faces = app.get(np.array(source_img))
    target_faces = app.get(np.array(target_img))
    
    if not source_faces:
        return target_img
        
    result = np.array(target_img)
    for face in target_faces:
        result = swapper.get(result, face, source_faces[0], paste_back=True)
        
    return PIL.Image.fromarray(result)

demo = gr.Interface(
    fn=swap, 
    inputs=[gr.Image(type="pil", label="Лицо-донор"), gr.Image(type="pil", label="Целевое фото")], 
    outputs=gr.Image(label="Результат"),
    title="🎭 Мой Face Swap"
)

demo.launch()