import gradio as gr
from gradio_client import Client

# Clients for image generation and swapping
image_generation_client = Client("AP123/SDXL-Lightning")
image_swapping_client = Client("sheikhed/face")

# Function to generate and swap images
def generate_and_swap_image(prompt, source_image):
    generated_image = image_generation_client.predict(prompt, "4-Step", api_name="/generate_image")
    swapped_image = image_swapping_client.predict(source_image, generated_image, True, api_name="/predict")
    return swapped_image

# Gradio interface
iface = gr.Blocks()

with iface:
    gr.Markdown("""
    # AI Image with Any Face
    Generate free AI images with your face or any face.<br>
    Contact me to bypass NSFW for more creativity and freedom😉<br>
    [Click here to Donate](https://nowpayments.io/donation/aheed)<br>
    Contact me for bulk processing and more AI software at +92-332-4399819<br>
    <b>Duplicate this space if you get error.</b>
    """)
    
    gr.Interface(
        fn=generate_and_swap_image,
        inputs=[
            gr.Textbox(label="Enter your prompt (English)"),
            gr.Image(label="Source Image", type="filepath"),
        ],
        outputs=gr.Image(label="Swapped Image")
    )
    
iface.launch()