import gradio as gr
from gradio_client import Client

def generate_omni_image(prompt, aspect_ratio):
    client = Client("selfit-camera/Omni-Image-Editor")
    try:
        # Ensure aspect_ratio is a string for the API
        result = client.predict(
           prompt=prompt,
           aspect_ratio=aspect_ratio,
           api_name="/text_to_image_interface"
        )
        # The API returns a tuple, with the image URL usually as the first element
        if isinstance(result, (list, tuple)) and result:
            # The first element is often the image in HTML format or a direct URL
            # Let's try to extract a direct image URL if possible from the HTML string
            image_output = result[0]
            if isinstance(image_output, str) and "<img src='" in image_output:
                start = image_output.find("src='") + len("src='")
                end = image_output.find("'", start)
                image_url = image_output[start:end]
                return image_url
            else:
                return image_output # Return directly if it's a file path or URL
        return None
    except Exception as e:
        return f"Error generating image: {e}"

with gr.Blocks() as demo:
    gr.Markdown("# Omni Image Editor with Gradio")

    with gr.Row():
        prompt_input = gr.Textbox(
            label="Prompt",
            placeholder="Describe the image you want to generate",
            lines=2
        )
        aspect_ratio_dropdown = gr.Dropdown(
            label="Aspect Ratio",
            choices=["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "5:4", "4:5", "21:9", "9:21"],
            value="16:9"
        )
    
    generate_btn = gr.Button("Generate Image", variant="primary")

    output_image = gr.Image(label="Generated Image", type="pil") # Changed type to 'pil' or 'filepath' for better handling
    output_text = gr.Textbox(label="Status/Error", interactive=False)

    generate_btn.click(
        fn=generate_omni_image,
        inputs=[prompt_input, aspect_ratio_dropdown],
        outputs=[output_image, output_text]
    )

demo.launch()