import gradio as gr
import insightface
import onnxruntime
from insightface.app import FaceAnalysis

def face_swap(user_image, celebrity_image):

  # Загрузка моделей
  detector = insightface.model_zoo.get_model('retinaface_r50_v1')
  arcface = insightface.model_zoo.get_model('arcface_r100_v1')
  
  # Определение пути к модели замены лица
  onnx_model_path = 'face_swap_model.onnx'

  # Обнаружение лиц
  user_bbox, user_landmarks = detector.detect(user_image)
  celebrity_bbox, celebrity_landmarks = detector.detect(celebrity_image)

  # Выравнивание лиц
  user_aligned = FaceAnalysis.align(user_image, user_landmarks)
  celebrity_aligned = FaceAnalysis.align(celebrity_image, celebrity_landmarks)

  # Получение векторных представлений лиц
  user_embedding = arcface.get_embedding(user_aligned)
  celebrity_embedding = arcface.get_embedding(celebrity_aligned)

  # Загрузка модели ONNX
  session = onnxruntime.InferenceSession(onnx_model_path)

  # Подготовка входных данных
  input_name = session.get_inputs()[0].name
  user_data = user_aligned.transpose(2, 0, 1)
  celebrity_data = celebrity_aligned.transpose(2, 0, 1)
  input_feed = {input_name: np.concatenate([user_data, celebrity_data], axis=0)}

  # Замена лица
  results = session.run(None, input_feed)
  swapped_face = results[0].transpose(1, 2, 0)

  return swapped_face


iface = gr.Interface(fn=face_swap, inputs=["image", "image"], outputs="image")

iface.launch()