import gradio as gr
from gradio_client import Client
import os
import shutil

# Load the APIs at the beginning
client_gen_image = Client("AP123/SDXL-Lightning")
client_face_swap = Client("craftgamesnetwork/face-swap")

def generate_and_swap(text_input, source_image_path):
    # First API: Generate image from text input
    gen_result = client_gen_image.predict(
        text_input,  # Text input
        "4-Step",    # Inference steps
        api_name="/generate_image"
    )
    print("Image generated successfully.")

    # Move the generated image to a known location
    generated_image_path = "generated_image.png"
    shutil.move(gen_result, generated_image_path)

    # Second API: Face swap
    swap_result_path = client_face_swap.predict(
        generated_image_path,  # Provide generated image path directly
        source_image_path,    # Provide source image path directly
        api_name="/predict"
    )
    print("Faces swapped successfully.")

    # Read the swapped image file
    with open(swap_result_path, "rb") as f:
        swap_result_content = f.read()

    # Save swapped image to a file
    swapped_image_path = "final_image.png"
    with open(swapped_image_path, "wb") as f:
        f.write(swap_result_content)

    print("Swapped image saved as:", swapped_image_path)
    return swapped_image_path

iface = gr.Interface(
    generate_and_swap,
    [
        gr.Textbox(label="Enter your prompt (English):"),
        gr.Image(type="filepath", label="Upload your source image:")
    ],
    "image",
    title="Face to AI Image"
)
iface.launch()