import os
import logging
import gdown
import insightface
import gradio as gr
from insightface.app import FaceAnalysis
from PIL import Image
import numpy as np

logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.INFO)

def prepare_face_analysis():
    """Подготовка приложения FaceAnalysis."""
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    return app

def download_model(model_url, model_path):
    """Загрузка модели, если она еще не существует."""
    if not os.path.exists(model_path):
        gdown.download(model_url, model_path, quiet=False)

def load_model(model_path):
    """Загрузка модели."""
    return insightface.model_zoo.get_model(model_path, download=False)

def swap_faces(user_image, celebrity_image):
    """Меняем лица на заданных изображениях."""
    try:
        # Меняем лица
        input_name = onnx_model.get_inputs()[0].name
        input_data = np.concatenate([user_image, celebrity_image], axis=0)
        results = onnx_model.run(None, {input_name: input_data})
        swapped_face = Image.fromarray(results[0])
    except Exception as e:
        logging.error("Ошибка при смене лиц: %s", e)
        swapped_face = user_image
    return swapped_face

# Подготовка приложения FaceAnalysis
app = prepare_face_analysis()

# Загрузка и загрузка модели
model_url = 'https://drive.google.com/uc?id=1HvZ4MAtzlY74Dk4ASGIS9L6Rg5oZdqvu'
model_path = 'inswapper_128.onnx'
download_model(model_url, model_path)
onnx_model = load_model(model_path)

# Запуск интерфейса Gradio
iface = gr.Interface(fn=swap_faces, inputs=["image", "image"], outputs="image")
iface.launch(share=True)