import gradio as gr
from gradio_client import Client
import os
import random
from io import BytesIO
from PIL import Image
import base64

KEYS = os.getenv("KEYS").split(",")

def get_random_api_key():
    return random.choice(KEYS)

def swap_face_api(source_img, target_img, doFaceEnhancer):
    try:
        api_key = get_random_api_key()
        client = Client("tuan2308/face-swap")

        # Конвертируем PIL изображения в байты, затем в base64
        source_bytes = BytesIO()
        source_img.save(source_bytes, format="JPEG")
        source_b64 = base64.b64encode(source_bytes.getvalue()).decode("utf-8")

        target_bytes = BytesIO()
        target_img.save(target_bytes, format="JPEG")
        target_b64 = base64.b64encode(target_bytes.getvalue()).decode("utf-8")


        result = client.predict(
            source_file=source_b64,  # Передаем base64 строку
            target_file=target_b64,  # Передаем base64 строку
            doFaceEnhancer=doFaceEnhancer,
            api_name="/predict",
            api_key=api_key # Передаем api_key в predict
        )

        # Конвертируем результат из байтов в PIL Image
        output_image = Image.open(BytesIO(result))
        return output_image

    except Exception as e:
        print(f"Ошибка при вызове API: {e}")
        return None


iface = gr.Interface(
    fn=swap_face_api,
    inputs=[
        gr.Image(type="pil", label="Source Image"),
        gr.Image(type="pil", label="Target Image"),
        gr.Checkbox(label="Face Enhancer?")
    ],
    outputs=gr.Image(type="pil", label="Output Image"),
    title="Face Swap via API"
)

iface.launch()