import gradio as gr
from control import main
from diffusers.utils import load_image


def process_images(image, mask, prompt):
    try:
        # Pass images directly to main function
        result = main(image, mask, prompt)
        return result
    except Exception as e:
        return str(e)


# Create Gradio interface
demo = gr.Interface(
    fn=process_images,
    inputs=[
        gr.Image(label="Input Image", type="pil"),
        gr.Image(label="Mask Image", type="pil"),
        gr.Textbox(label="Prompt"),
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Image Inpainting with FLUX ControlNet",
    description="Upload an image and its mask, then provide a prompt to generate the inpainted result.",
)

if __name__ == "__main__":
    demo.launch()