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


client_gen_image = Client("AP123/SDXL-Lightning")
client_face_swap = Client("craftgamesnetwork/face-swap")

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

    
    generated_image_path = "generated_image.png"
    shutil.move(gen_result, generated_image_path)

    
    swap_result_path = client_face_swap.predict(
        generated_image_path,
        source_image_path,
        api_name="/predict"
    )
    print("Faces swapped successfully.")

    
    with open(swap_result_path, "rb") as f:
        swap_result_content = f.read()

    
    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",
    description="Generate free AI image with your or any face. Support me in making better AI codes as I am a solo developer [Click here to Donate](https://nowpayments.io/donation/aheed) Contact me for bulk processing and better AI software +92-332-4399819",
    title="Face to AI Image"
)
iface.launch()