import os
import logging
import gdown
import insightface
import gradio as gr



from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
from PIL import Image



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

app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))

# Download 'inswapper_128.onnx' file using gdown
model_url = 'https://drive.google.com/uc?id=1HvZ4MAtzlY74Dk4ASGIS9L6Rg5oZdqvu'
model_output_path = 'inswapper/inswapper_128.onnx'
if not os.path.exists(model_output_path):
    gdown.download(model_url, model_output_path, quiet=False)

swapper = insightface.model_zoo.get_model('inswapper/inswapper_128.onnx', download=False, download_zip=False)

def swap_faces(user_image, celebrity_image):
    try:
        # Swap faces
        input_name = swapper.get_inputs()[0].name
        input_data = np.concatenate([user_image, celebrity_image], axis=0)
        results = swapper.run(None, {input_name: input_data})
        swapped_face = Image.fromarray(results[0])

    except Exception as e:
        logging.error(e)
        gr.Interface.update_output("Couldn't swap faces: " + str(e))
        swapped_face = user_image

    return swapped_face

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