import gradio as gr
from gradio_client import Client
import re

# Instantiate a Client object from gradio_client pointing to the 'selfit-camera/Omni-Image-Editor' Space.
client = Client("selfit-camera/Omni-Image-Editor")

# Define a Python function, generate_image, that takes a single string argument prompt.
def generate_image(prompt):
    # Call the client.predict() method with the user's prompt, aspect_ratio='16:9', and api_name='/text_to_image_interface'.
    result = client.predict(
        prompt=prompt,
        aspect_ratio="16:9",
        api_name="/text_to_image_interface"
    )

    # The predict method returns a tuple. The first element of this tuple is an HTML string containing the image.
    # Extract the image URL from this HTML string.
    html_string = result[0]
    match = re.search(r"src='([^']+)'", html_string)
    if match:
        image_url = match.group(1)
        return image_url
    else:
        # Handle cases where the URL might not be found, e.g., return a default image or raise an error.
        return "https://via.placeholder.com/400x200?text=Error:Image+Not+Found"

# Create a Gradio application using gr.Blocks for more granular control.
with gr.Blocks(title='Omni Image Editor with Gradio') as demo:
    gr.Markdown("## Omni Image Editor with Gradio")
    with gr.Row():
        prompt_input = gr.Textbox(label='Enter your image prompt', placeholder='e.g., A futuristic city at sunset')
        generate_btn = gr.Button("Generate Image")
    output_image = gr.Image(label='Generated Image')

    # Bind the generate_image function to the button click event.
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input],
        outputs=[output_image]
    )

# Launch the Gradio application.
demo.launch()