import gradio as gr
import os
import random
from PIL import Image
import base64
import requests

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()
        api_url = "https://tuan2308-face-swap-predict.hf.space/api/predict/"  # URL вашего API

        source_b64 = gr.encode_pil_to_base64(source_img)
        target_b64 = gr.encode_pil_to_base64(target_img)

        payload = {
            "data": [
                source_b64,
                target_b64,
                doFaceEnhancer
            ],
            "api_key": api_key  # Передаем api_key в payload
        }

        response = requests.post(api_url, json=payload)
        response.raise_for_status()  # Вызывает исключение, если есть ошибка

        # Получаем результат из ответа
        result = response.json()
        result_data = result["data"]  # Извлекаем данные из ответа
        result_decoded = base64.b64decode(result_data[0]) # Декодируем base64


        output_image = Image.open(BytesIO(result_decoded))
        return output_image

    except requests.exceptions.RequestException as e:
        print(f"Ошибка API: {e}")
        return None

    except Exception as e:
        print(f"Ошибка: {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()