import gradio as gr
import os
import random
from PIL import Image
from io import BytesIO
import base64
import requests
import numpy as np

# ... (остальной код без изменений)

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/"

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


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

        payload = {
            "data": [
                source_b64,
                target_b64,
                doFaceEnhancer
            ],
            "api_key": api_key
        }

        response = requests.post(api_url, json=payload)
        response.raise_for_status()

        print(response)
        
        result = response.json()
        result_data = result["data"][0] # Проверьте индекс
        result_decoded = base64.b64decode(result_data)
        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="numpy", label="Source Image"),  # Изменено на numpy
        gr.Image(type="numpy", label="Target Image"),  # Изменено на numpy
        gr.Checkbox(label="Face Enhancer?")
    ],
    outputs=gr.Image(type="pil", label="Output Image"), # Остается pil
    title="Face Swap via API"
)

iface.launch(server_name="0.0.0.0", server_port=7860)